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/tvm/relay/doc.cc
22  * \brief Doc ADT used for pretty printing.
23  * Based on Section 1 of https://homepages.inf.ed.ac.uk/wadler/papers/prettier/prettier.pdf.
24  */
25 #include <memory>
26 #include <vector>
27 #include "doc.h"
28 
29 namespace tvm {
30 namespace relay {
31 
32 // Text constructor
Text(const std::string & str)33 DocAtom Text(const std::string& str) {
34   return std::make_shared<TextNode>(str);
35 }
36 
37 // Line constructor
Line(int indent=0)38 DocAtom Line(int indent = 0) {
39   return std::make_shared<LineNode>(indent);
40 }
41 
Doc(const std::string & str)42 Doc::Doc(const std::string& str) {
43   if (str == "\n") {
44     this->stream_ = {Line()};
45   } else {
46     this->stream_ = {Text(str)};
47   }
48 }
49 
50 // DSL function implementations
51 
operator <<(const Doc & right)52 Doc& Doc::operator<<(const Doc& right) {
53   CHECK(this != &right);
54   this->stream_.insert(this->stream_.end(), right.stream_.begin(), right.stream_.end());
55   return *this;
56 }
57 
operator <<(const std::string & right)58 Doc& Doc::operator<<(const std::string& right) {
59   return *this << Doc(right);
60 }
61 
operator <<(const DocAtom & right)62 Doc& Doc::operator<<(const DocAtom& right) {
63   this->stream_.push_back(right);
64   return *this;
65 }
66 
Indent(int indent,const Doc & doc)67 Doc Indent(int indent, const Doc& doc) {
68   Doc ret;
69   for (auto atom : doc.stream_) {
70     if (auto text = std::dynamic_pointer_cast<TextNode>(atom)) {
71       ret.stream_.push_back(text);
72     } else if (auto line = std::dynamic_pointer_cast<LineNode>(atom)) {
73       ret.stream_.push_back(Line(indent + line->indent));
74     } else {CHECK(false);}
75   }
76   return ret;
77 }
78 
str()79 std::string Doc::str() {
80   std::ostringstream os;
81   for (auto atom : this->stream_) {
82     if (auto text = std::dynamic_pointer_cast<TextNode>(atom)) {
83       os << text->str;
84     } else if (auto line = std::dynamic_pointer_cast<LineNode>(atom)) {
85       os << "\n" << std::string(line->indent, ' ');
86     } else {CHECK(false);}
87   }
88   return os.str();
89 }
90 
PrintSep(const std::vector<Doc> & vec,const Doc & sep)91 Doc PrintSep(const std::vector<Doc>& vec, const Doc& sep) {
92   Doc seq;
93   if (vec.size() != 0) {
94     seq = vec[0];
95     for (size_t i = 1; i < vec.size(); i++) {
96       seq << sep << vec[i];
97     }
98   }
99   return seq;
100 }
101 
PrintBool(bool value)102 Doc PrintBool(bool value) {
103   if (value) {
104     return Doc("True");
105   } else {
106     return Doc("False");
107   }
108 }
109 
PrintDType(DataType dtype)110 Doc PrintDType(DataType dtype) {
111   return Doc(runtime::TVMType2String(Type2TVMType(dtype)));
112 }
113 
PrintString(const std::string & value)114 Doc PrintString(const std::string& value) {
115   // TODO(M.K.): add escape.
116   Doc doc;
117   return doc << "\"" << value << "\"";
118 }
119 
PrintNewLine(int ident)120 Doc PrintNewLine(int ident) {
121   Doc doc;
122   return doc << Line(ident);
123 }
124 
125 }  // namespace relay
126 }  // namespace tvm
127