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