1 | /* |
---|---|
2 | * SPDX-License-Identifier: Apache-2.0 |
3 | */ |
4 | |
5 | // ATTENTION: The code in this file is highly EXPERIMENTAL. |
6 | // Adventurous users should note that the APIs will probably change. |
7 | |
8 | #include <stdint.h> |
9 | #include <mutex> |
10 | #include <string> |
11 | #include <unordered_map> |
12 | #include <vector> |
13 | |
14 | #include "onnx/common/assertions.h" |
15 | #include "onnx/common/interned_strings.h" |
16 | |
17 | namespace ONNX_NAMESPACE { |
18 | |
19 | struct InternedStrings { |
20 | InternedStrings() : next_sym(kLastSymbol) { |
21 | #define REGISTER_SYMBOL(s) \ |
22 | string_to_sym_[#s] = k##s; \ |
23 | sym_to_string_[k##s] = #s; |
24 | FORALL_BUILTIN_SYMBOLS(REGISTER_SYMBOL) |
25 | #undef REGISTER_SYMBOL |
26 | } |
27 | uint32_t symbol(const std::string& s) { |
28 | std::lock_guard<std::mutex> guard(mutex_); |
29 | auto it = string_to_sym_.find(s); |
30 | if (it != string_to_sym_.end()) |
31 | return it->second; |
32 | uint32_t k = next_sym++; |
33 | string_to_sym_[s] = k; |
34 | sym_to_string_[k] = s; |
35 | return k; |
36 | } |
37 | const char* string(Symbol sym) { |
38 | // Builtin Symbols are also in the maps, but |
39 | // we can bypass the need to acquire a lock |
40 | // to read the map for Builtins because we already |
41 | // know their string value |
42 | switch (sym) { |
43 | #define DEFINE_CASE(s) \ |
44 | case k##s: \ |
45 | return #s; |
46 | FORALL_BUILTIN_SYMBOLS(DEFINE_CASE) |
47 | #undef DEFINE_CASE |
48 | default: |
49 | return customString(sym); |
50 | } |
51 | } |
52 | |
53 | private: |
54 | const char* customString(Symbol sym) { |
55 | std::lock_guard<std::mutex> guard(mutex_); |
56 | auto it = sym_to_string_.find(sym); |
57 | ONNX_ASSERT(it != sym_to_string_.end()); |
58 | return it->second.c_str(); |
59 | } |
60 | std::unordered_map<std::string, uint32_t> string_to_sym_; |
61 | std::unordered_map<uint32_t, std::string> sym_to_string_; |
62 | uint32_t next_sym; |
63 | std::mutex mutex_; |
64 | }; |
65 | |
66 | static InternedStrings& globalStrings() { |
67 | static InternedStrings s; |
68 | return s; |
69 | } |
70 | |
71 | const char* Symbol::toString() const { |
72 | return globalStrings().string(*this); |
73 | } |
74 | |
75 | Symbol::Symbol(const std::string& s) : value(globalStrings().symbol(s)) {} |
76 | |
77 | } // namespace ONNX_NAMESPACE |
78 |