1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | /*! |
20 | * \file source_map.cc |
21 | * \brief The implementation of the source map data structure. |
22 | */ |
23 | #include <tvm/ir/source_map.h> |
24 | #include <tvm/ir/transform.h> |
25 | #include <tvm/runtime/registry.h> |
26 | |
27 | #include <algorithm> |
28 | |
29 | namespace tvm { |
30 | |
31 | TVM_REGISTER_PASS_CONFIG_OPTION("relay.frontend.fill_span" , Bool); |
32 | |
33 | ObjectPtr<Object> GetSourceNameNode(const String& name) { |
34 | // always return pointer as the reference can change as map re-allocate. |
35 | // or use another level of indirection by creating a unique_ptr |
36 | static std::unordered_map<String, ObjectPtr<SourceNameNode>> source_map; |
37 | |
38 | auto sn = source_map.find(name); |
39 | if (sn == source_map.end()) { |
40 | ObjectPtr<SourceNameNode> n = make_object<SourceNameNode>(); |
41 | source_map[name] = n; |
42 | n->name = std::move(name); |
43 | return n; |
44 | } else { |
45 | return sn->second; |
46 | } |
47 | } |
48 | |
49 | ObjectPtr<Object> GetSourceNameNodeByStr(const std::string& name) { |
50 | return GetSourceNameNode(name); |
51 | } |
52 | |
53 | SourceName SourceName::Get(const String& name) { return SourceName(GetSourceNameNode(name)); } |
54 | |
55 | TVM_REGISTER_GLOBAL("ir.SourceName" ).set_body_typed(SourceName::Get); |
56 | |
57 | TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) |
58 | .set_dispatch<SourceNameNode>([](const ObjectRef& ref, ReprPrinter* p) { |
59 | auto* node = static_cast<const SourceNameNode*>(ref.get()); |
60 | p->stream << "SourceName(" << node->name << ", " << node << ")" ; |
61 | }); |
62 | |
63 | TVM_REGISTER_NODE_TYPE(SourceNameNode) |
64 | .set_creator(GetSourceNameNodeByStr) |
65 | .set_repr_bytes([](const Object* n) -> std::string { |
66 | return static_cast<const SourceNameNode*>(n)->name; |
67 | }); |
68 | |
69 | Span::Span(SourceName source_name, int line, int end_line, int column, int end_column) { |
70 | auto n = make_object<SpanNode>(); |
71 | n->source_name = std::move(source_name); |
72 | n->line = line; |
73 | n->end_line = end_line; |
74 | n->column = column; |
75 | n->end_column = end_column; |
76 | data_ = std::move(n); |
77 | } |
78 | |
79 | Span Span::Merge(const Span& other) const { |
80 | ICHECK(this->defined() && other.defined()) << "Span::Merge: both spans must be defined" ; |
81 | |
82 | ICHECK((*this)->source_name == other->source_name); |
83 | return Span((*this)->source_name, std::min((*this)->line, other->line), |
84 | std::max((*this)->end_line, other->end_line), |
85 | std::min((*this)->column, other->column), |
86 | std::max((*this)->end_column, other->end_column)); |
87 | } |
88 | |
89 | TVM_REGISTER_NODE_TYPE(SpanNode); |
90 | |
91 | TVM_REGISTER_GLOBAL("ir.Span" ).set_body_typed([](SourceName source_name, int line, int end_line, |
92 | int column, int end_column) { |
93 | return Span(source_name, line, end_line, column, end_column); |
94 | }); |
95 | |
96 | TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) |
97 | .set_dispatch<SpanNode>([](const ObjectRef& ref, ReprPrinter* p) { |
98 | auto* node = static_cast<const SpanNode*>(ref.get()); |
99 | p->stream << "Span(" << node->source_name << ", " << node->line << ", " << node->end_line |
100 | << ", " << node->column << ", " << node->end_column << ")" ; |
101 | }); |
102 | |
103 | TVM_REGISTER_NODE_TYPE(SourceNode); |
104 | |
105 | /*! \brief Construct a source from a string. */ |
106 | Source::Source(SourceName src_name, std::string source) { |
107 | auto n = make_object<SourceNode>(); |
108 | n->source_name = std::move(src_name); |
109 | n->source = std::move(source); |
110 | |
111 | int index = 0; |
112 | int length = 0; |
113 | n->line_map.push_back({index, length}); |
114 | // NB(@jroesch): |
115 | std::string source_str = n->source; |
116 | for (auto c : source_str) { |
117 | if (c == '\n') { |
118 | // Record the length of the line. |
119 | n->line_map.back().second = length; |
120 | // Bump past the newline. |
121 | index += 1; |
122 | // Record the start of the next line, and put placeholder for length. |
123 | n->line_map.push_back({index, 0}); |
124 | // Reset length to zero. |
125 | length = 0; |
126 | } else { |
127 | length += 1; |
128 | index += 1; |
129 | } |
130 | } |
131 | n->line_map.back().second = length; |
132 | |
133 | data_ = n; |
134 | } |
135 | |
136 | tvm::String Source::GetLine(int line) { |
137 | VLOG(1) << "Source::GetLine: line=" << line; |
138 | ICHECK(line - 1 < static_cast<int64_t>((*this)->line_map.size())) |
139 | << "requested line: " << line << "at index: " << (line - 1) |
140 | << "line_map size: " << (*this)->line_map.size() << "source: " << (*this)->source; |
141 | |
142 | // Adjust for zero indexing, now have (line_start, line_length); |
143 | auto range = (*this)->line_map.at(line - 1); |
144 | int line_start = range.first; |
145 | int line_length = range.second; |
146 | VLOG(1) << "Source::GetLine: line_start=" << line_start << " line_length=" << line_length; |
147 | // TODO(@jroesch): expose substring on tvm::String. |
148 | auto line_text = std::string((*this)->source).substr(line_start, line_length); |
149 | VLOG(1) << "Source::GetLine: line_text=" << line_text; |
150 | return line_text; |
151 | } |
152 | |
153 | TVM_REGISTER_NODE_TYPE(SourceMapNode); |
154 | |
155 | SourceMap::SourceMap(Map<SourceName, Source> source_map) { |
156 | auto n = make_object<SourceMapNode>(); |
157 | n->source_map = std::move(source_map); |
158 | data_ = std::move(n); |
159 | } |
160 | |
161 | void SourceMap::Add(const Source& source) { (*this)->source_map.Set(source->source_name, source); } |
162 | |
163 | TVM_REGISTER_GLOBAL("SourceMapAdd" ).set_body_typed([](SourceMap map, String name, String content) { |
164 | auto src_name = SourceName::Get(name); |
165 | Source source(src_name, content); |
166 | map.Add(source); |
167 | return src_name; |
168 | }); |
169 | |
170 | } // namespace tvm |
171 | |