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