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
29namespace tvm {
30
31TVM_REGISTER_PASS_CONFIG_OPTION("relay.frontend.fill_span", Bool);
32
33ObjectPtr<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
49ObjectPtr<Object> GetSourceNameNodeByStr(const std::string& name) {
50 return GetSourceNameNode(name);
51}
52
53SourceName SourceName::Get(const String& name) { return SourceName(GetSourceNameNode(name)); }
54
55TVM_REGISTER_GLOBAL("ir.SourceName").set_body_typed(SourceName::Get);
56
57TVM_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
63TVM_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
69Span::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
79Span 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
89TVM_REGISTER_NODE_TYPE(SpanNode);
90
91TVM_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
96TVM_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
103TVM_REGISTER_NODE_TYPE(SourceNode);
104
105/*! \brief Construct a source from a string. */
106Source::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
136tvm::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
153TVM_REGISTER_NODE_TYPE(SourceMapNode);
154
155SourceMap::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
161void SourceMap::Add(const Source& source) { (*this)->source_map.Set(source->source_name, source); }
162
163TVM_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