Caffe2 - C++ API
A deep learning, cross platform ML framework
code_template.h
1 #pragma once
2 #include <sstream>
3 #include <string>
4 #include <unordered_map>
5 #include <vector>
6 
7 namespace torch {
8 namespace jit {
9 
10 // A template environment is a mapping from template variable names, e.g.,
11 // identifier (corresponding to $identifier) to their expansions.
12 //
13 // This template environment supports storing strings, numbers and lists
14 // of strings, and can be chained together (so that lookup proceeds in
15 // in the top level environment, and then recurses into a parent
16 // environment if the key is not found.)
17 struct TemplateEnv {
18  TemplateEnv() : parent(nullptr) {}
19  TemplateEnv(TemplateEnv& parent) : parent(&parent) {}
20 
21  using string_list = std::vector<std::string>;
22 
23  // Add a string 'v' to the map at key 'k'.
24  void s(const std::string& k, const std::string& v) {
25  strings_[k] = v;
26  lists_.erase(k);
27  }
28 
29  // Add a number 'v' to the map at key 'k'
30  template <typename T>
31  void d(const std::string& k, const T& v) {
32  strings_[k] = std::to_string(v);
33  lists_.erase(k);
34  }
35 
36  // Retrieve the string representation of the value stored at 'k' from the map.
37  // Raises an exception if the key is not found.
38  const std::string& s(const std::string& k) const {
39  if (strings_.count(k) == 0) {
40  if (parent) {
41  return parent->s(k);
42  }
43  notFound(k);
44  }
45  return strings_.at(k);
46  }
47 
48  // Store a list of strings 'v' in the map at 'k'.
49  void v(const std::string& k, const string_list& v) {
50  lists_[k] = v;
51  strings_.erase(k);
52  }
53 
54  // Retrieve a list of strings stored at 'k' from the map.
55  // Raises an exception if the key is not found.
56  const string_list& v(const std::string& k) const {
57  if (lists_.count(k) == 0) {
58  if (parent) {
59  return parent->v(k);
60  }
61  notFound(k);
62  }
63  return lists_.at(k);
64  }
65 
66  // Test if a string 'k' is a string (as opposed to a list.)
67  bool keyIsString(const std::string& k) const {
68  if (strings_.count(k) > 0)
69  return true;
70  if (lists_.count(k) > 0)
71  return false;
72  if (parent)
73  return parent->keyIsString(k);
74  notFound(k);
75  }
76 
77  private:
78  [[noreturn]] void notFound(const std::string& k) const {
79  std::stringstream ss;
80  ss << "key not found: " << k;
81  throw std::logic_error(ss.str());
82  }
83  std::unordered_map<std::string, std::string> strings_;
84  std::unordered_map<std::string, string_list> lists_;
85  TemplateEnv* parent;
86 };
87 
88 /*
89 # Match $identifier or ${identifier} and replace with the value in env.
90 # If this identifier is at the beginning of whitespace on a line
91 # and its value is a list then it is treated as
92 # block substitution by indenting all lines of all elements.
93 # If the identifier is on a line starting with non-whitespace and a list
94 # then it is comma separated. ${,foo} will insert a comma before the list
95 # if this list is not empty and ${foo,} will insert one after.
96 */
97 struct CodeTemplate {
98  /* implicit */ CodeTemplate(std::string t) : template_text(std::move(t)) {}
99 
100  std::string format(const TemplateEnv& env) const {
101  std::stringstream out;
102  size_t pos = 0;
103  size_t indent = 0;
104  bool all_whitespace = true;
105  while (pos < template_text.size()) {
106  char c = template_text[pos];
107  if (c == '$') {
108  std::stringstream kss;
109  bool comma_before;
110  bool comma_after;
111  size_t new_pos = parseKey(pos, kss, comma_before, comma_after);
112  std::string k = kss.str();
113  bool is_string = env.keyIsString(k);
114  if (all_whitespace) {
115  if (is_string)
116  emitStringWithIndents(out, indent, env.s(k));
117  else
118  emitLinesIndented(out, indent, env.v(k));
119  } else {
120  if (is_string)
121  out << env.s(k);
122  else
123  emitCommaSeparatedList(out, env.v(k), comma_before, comma_after);
124  }
125  all_whitespace = false;
126  pos = new_pos;
127  } else {
128  out << c;
129  if (!isspace(c))
130  all_whitespace = false;
131  indent++;
132  if (c == '\n') {
133  indent = 0;
134  all_whitespace = true;
135  }
136  pos++;
137  }
138  }
139  return out.str();
140  }
141 
142  private:
143  using string_list = std::vector<std::string>;
144  char charAt(size_t p) const {
145  if (p >= template_text.size())
146  throw std::logic_error("EOS found in key");
147  return template_text[p];
148  }
149  size_t parseKey(
150  size_t pos,
151  std::ostream& k,
152  bool& comma_before,
153  bool& comma_after) const {
154  comma_before = false;
155  comma_after = false;
156  pos++;
157  if (charAt(pos) == '{') {
158  pos++;
159  if (charAt(pos) == ',') {
160  comma_before = true;
161  pos++;
162  }
163  pos = parseIdent(pos, k);
164  if (charAt(pos) == ',') {
165  comma_after = true;
166  pos++;
167  }
168  if (charAt(pos) != '}')
169  throw std::logic_error("missing terminating '}'");
170  pos++;
171  return pos;
172  } else {
173  return parseIdent(pos, k);
174  }
175  }
176  size_t parseIdent(size_t pos, std::ostream& k) const {
177  while (pos < template_text.size() &&
178  (isalnum(template_text[pos]) || template_text[pos] == '_')) {
179  k << template_text[pos];
180  pos++;
181  }
182  return pos;
183  }
184  void emitCommaSeparatedList(
185  std::ostream& out,
186  const string_list& strings,
187  bool comma_before,
188  bool comma_after) const {
189  if (comma_before && strings.size() > 0)
190  out << ", ";
191  for (size_t i = 0; i < strings.size(); ++i) {
192  if (i > 0)
193  out << ", ";
194  out << strings[i];
195  }
196  if (comma_after && strings.size() > 0)
197  out << ", ";
198  }
199  // These indentation functions follow the convention that they never emit
200  // leading or trailing newlines when the input string does not have leading
201  // or trailing newlines. It's the responsibility of the calling function
202  // to indent correctly in the context.
203  void emitIndent(std::ostream& out, size_t indent) const {
204  for (size_t i = 0; i < indent; ++i) {
205  out << " ";
206  }
207  }
208  void emitStringWithIndents(
209  std::ostream& out,
210  size_t indent,
211  const std::string& str) const {
212  for (auto c : str) {
213  out << c;
214  if (c == '\n') {
215  emitIndent(out, indent);
216  }
217  }
218  }
219  void emitLinesIndented(
220  std::stringstream& out,
221  size_t indent,
222  const string_list& strings) const {
223  for (size_t i = 0; i < strings.size(); ++i) {
224  if (i > 0)
225  emitIndent(out, indent);
226  emitStringWithIndents(out, indent, strings[i]);
227  if (i + 1 != strings.size())
228  out << "\n";
229  }
230  }
231  std::string template_text;
232 };
233 static inline std::string format(const std::string& fmt, TemplateEnv& env) {
234  return CodeTemplate(fmt).format(env);
235 }
236 
237 } // namespace jit
238 } // namespace torch
Definition: jit_type.h:17