1 #include <torch/csrc/jit/frontend/lexer.h>
2
3 #include <c10/util/Exception.h>
4
5 #include <string>
6 #include <unordered_map>
7
8 namespace torch::jit {
9
10 static const std::unordered_map<int, int> binary_prec = {
11 {TK_IF, 1},
12 {TK_FOR, 1},
13 {TK_AND, 2},
14 {TK_OR, 2},
15 // reserve a level for unary not
16 {TK_IN, 4},
17 {TK_NOTIN, 4},
18 {'<', 4},
19 {'>', 4},
20 {TK_IS, 4},
21 {TK_ISNOT, 4},
22 {TK_EQ, 4},
23 {TK_LE, 4},
24 {TK_GE, 4},
25 {TK_NE, 4},
26 {'|', 5},
27 {'^', 6},
28 {'&', 7},
29 {TK_LSHIFT, 8},
30 {TK_RSHIFT, 8},
31 {'+', 9},
32 {'-', 9},
33 {'*', 10},
34 {'/', 10},
35 {TK_FLOOR_DIV, 10},
36 {'%', 10},
37 {'@', 10},
38 {TK_POW, 11},
39 };
40
41 static const std::unordered_map<int, int> unary_prec = {
42 {TK_NOT, 3},
43 {'~', 3},
44 {'-', 10},
45 {'*', 10},
46 };
47
isUnary(int kind,int * prec)48 bool SharedParserData::isUnary(int kind, int* prec) {
49 auto it = unary_prec.find(kind);
50 if (it != unary_prec.end()) {
51 *prec = it->second;
52 return true;
53 }
54 return false;
55 }
isBinary(int kind,int * prec)56 bool SharedParserData::isBinary(int kind, int* prec) {
57 auto it = binary_prec.find(kind);
58 if (it != binary_prec.end()) {
59 *prec = it->second;
60 return true;
61 }
62 return false;
63 }
64
stringToKind(const std::string & str)65 C10_EXPORT int stringToKind(const std::string& str) {
66 static std::unordered_map<std::string, int> str_to_kind = []() {
67 std::unordered_map<std::string, int> ret_str_to_kind;
68 for (char tok : std::string(valid_single_char_tokens))
69 // NOLINTNEXTLINE(bugprone-signed-char-misuse)
70 ret_str_to_kind[std::string(1, tok)] = tok;
71 #define DEFINE_CASE(tok, _, str) \
72 if (std::string(str) != "") \
73 ret_str_to_kind[str] = tok;
74 TC_FORALL_TOKEN_KINDS(DEFINE_CASE)
75 #undef DEFINE_CASE
76 return ret_str_to_kind;
77 }();
78 try {
79 return str_to_kind.at(str);
80 } catch (std::out_of_range&) {
81 throw std::out_of_range("unknown token in stringToKind");
82 }
83 }
84
kindToString(int kind)85 C10_EXPORT std::string kindToString(int kind) {
86 if (kind < 256)
87 return std::string(1, static_cast<char>(kind));
88 switch (kind) {
89 #define DEFINE_CASE(tok, str, _) \
90 case tok: \
91 return str;
92 TC_FORALL_TOKEN_KINDS(DEFINE_CASE)
93 #undef DEFINE_CASE
94 default:
95 throw std::runtime_error("Unknown kind: " + std::to_string(kind));
96 }
97 }
98
sharedParserData()99 C10_EXPORT SharedParserData& sharedParserData() {
100 static SharedParserData data; // safely handles multi-threaded init
101 return data;
102 }
103
104 } // namespace torch::jit
105