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
35namespace tvm {
36namespace relay {
37
38/*!
39 * \brief Represent a piece of text in the doc.
40 */
41class 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
52TVM_REGISTER_OBJECT_TYPE(DocTextNode);
53
54class 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 */
64class 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
75TVM_REGISTER_OBJECT_TYPE(DocLineNode);
76
77class 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
85Doc& 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
91Doc& Doc::operator<<(std::string right) { return *this << DocText(right); }
92
93Doc& Doc::operator<<(const DocAtom& right) {
94 this->stream_.push_back(right);
95 return *this;
96}
97
98std::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
112Doc Doc::NewLine(int indent) { return Doc() << DocLine(indent); }
113
114Doc Doc::Text(std::string text) { return Doc() << DocText(text); }
115
116Doc Doc::RawText(std::string text) {
117 return Doc() << DocAtom(runtime::make_object<DocTextNode>(text));
118}
119
120Doc 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
129Doc Doc::StrLiteral(const std::string& value, std::string quote) {
130 Doc doc;
131 return doc << quote << support::StrEscape(value) << quote;
132}
133
134Doc Doc::PyBoolLiteral(bool value) {
135 if (value) {
136 return Doc::Text("True");
137 } else {
138 return Doc::Text("False");
139 }
140}
141
142Doc 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
150Doc 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