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 | /*! |
21 | * \file src/relay/doc.cc |
22 | * \brief Doc ADT used for pretty printing. |
23 | * |
24 | * Reference: Philip Wadler. A Prettier Printer. Journal of Functional Programming'98 |
25 | */ |
26 | #include "doc.h" |
27 | |
28 | #include <tvm/runtime/packed_func.h> |
29 | |
30 | #include <sstream> |
31 | #include <vector> |
32 | |
33 | #include "../../support/str_escape.h" |
34 | |
35 | namespace tvm { |
36 | namespace relay { |
37 | |
38 | /*! |
39 | * \brief Represent a piece of text in the doc. |
40 | */ |
41 | class DocTextNode : public DocAtomNode { |
42 | public: |
43 | /*! \brief The str content in the text. */ |
44 | std::string str; |
45 | |
46 | explicit DocTextNode(std::string str_val) : str(str_val) {} |
47 | |
48 | static constexpr const char* _type_key = "printer.DocText" ; |
49 | TVM_DECLARE_FINAL_OBJECT_INFO(DocTextNode, DocAtomNode); |
50 | }; |
51 | |
52 | TVM_REGISTER_OBJECT_TYPE(DocTextNode); |
53 | |
54 | class DocText : public DocAtom { |
55 | public: |
56 | explicit DocText(std::string str) { data_ = runtime::make_object<DocTextNode>(str); } |
57 | |
58 | TVM_DEFINE_OBJECT_REF_METHODS(DocText, DocAtom, DocTextNode); |
59 | }; |
60 | |
61 | /*! |
62 | * \brief Represent a line breaker in the doc. |
63 | */ |
64 | class DocLineNode : public DocAtomNode { |
65 | public: |
66 | /*! \brief The amount of indent in newline. */ |
67 | int indent; |
68 | |
69 | explicit DocLineNode(int indent) : indent(indent) {} |
70 | |
71 | static constexpr const char* _type_key = "printer.DocLine" ; |
72 | TVM_DECLARE_FINAL_OBJECT_INFO(DocLineNode, DocAtomNode); |
73 | }; |
74 | |
75 | TVM_REGISTER_OBJECT_TYPE(DocLineNode); |
76 | |
77 | class DocLine : public DocAtom { |
78 | public: |
79 | explicit DocLine(int indent) { data_ = runtime::make_object<DocLineNode>(indent); } |
80 | |
81 | TVM_DEFINE_OBJECT_REF_METHODS(DocLine, DocAtom, DocLineNode); |
82 | }; |
83 | |
84 | // DSL function implementations |
85 | Doc& Doc::operator<<(const Doc& right) { |
86 | ICHECK(this != &right); |
87 | this->stream_.insert(this->stream_.end(), right.stream_.begin(), right.stream_.end()); |
88 | return *this; |
89 | } |
90 | |
91 | Doc& Doc::operator<<(std::string right) { return *this << DocText(right); } |
92 | |
93 | Doc& Doc::operator<<(const DocAtom& right) { |
94 | this->stream_.push_back(right); |
95 | return *this; |
96 | } |
97 | |
98 | std::string Doc::str() { |
99 | std::ostringstream os; |
100 | for (auto atom : this->stream_) { |
101 | if (auto* text = atom.as<DocTextNode>()) { |
102 | os << text->str; |
103 | } else if (auto* line = atom.as<DocLineNode>()) { |
104 | os << "\n" << std::string(line->indent, ' '); |
105 | } else { |
106 | LOG(FATAL) << "do not expect type " << atom->GetTypeKey(); |
107 | } |
108 | } |
109 | return os.str(); |
110 | } |
111 | |
112 | Doc Doc::NewLine(int indent) { return Doc() << DocLine(indent); } |
113 | |
114 | Doc Doc::Text(std::string text) { return Doc() << DocText(text); } |
115 | |
116 | Doc Doc::RawText(std::string text) { |
117 | return Doc() << DocAtom(runtime::make_object<DocTextNode>(text)); |
118 | } |
119 | |
120 | Doc Doc::Indent(int indent, Doc doc) { |
121 | for (size_t i = 0; i < doc.stream_.size(); ++i) { |
122 | if (auto* line = doc.stream_[i].as<DocLineNode>()) { |
123 | doc.stream_[i] = DocLine(indent + line->indent); |
124 | } |
125 | } |
126 | return doc; |
127 | } |
128 | |
129 | Doc Doc::StrLiteral(const std::string& value, std::string quote) { |
130 | Doc doc; |
131 | return doc << quote << support::StrEscape(value) << quote; |
132 | } |
133 | |
134 | Doc Doc::PyBoolLiteral(bool value) { |
135 | if (value) { |
136 | return Doc::Text("True" ); |
137 | } else { |
138 | return Doc::Text("False" ); |
139 | } |
140 | } |
141 | |
142 | Doc Doc::Brace(std::string open, const Doc& body, std::string close, int indent) { |
143 | Doc doc; |
144 | doc << open; |
145 | doc << Indent(indent, NewLine() << body) << NewLine(); |
146 | doc << close; |
147 | return doc; |
148 | } |
149 | |
150 | Doc Doc::Concat(const std::vector<Doc>& vec, const Doc& sep) { |
151 | Doc seq; |
152 | if (vec.size() != 0) { |
153 | if (vec.size() == 1) return vec[0]; |
154 | seq << vec[0]; |
155 | for (size_t i = 1; i < vec.size(); ++i) { |
156 | seq << sep << vec[i]; |
157 | } |
158 | } |
159 | return seq; |
160 | } |
161 | } // namespace relay |
162 | } // namespace tvm |
163 | |