[REFACTOR] Establish printer in the source folder.

As we move towards the unified IR, we will eventually want to build a unified
printers for both relay and TIR.

This PR isolate the printer component into a separate folder in src as a first step.

- Refactored the Doc DSL using Object, clean up APIs.
- Isolate out the meta data into a header.
- move printer into relay_text_printer, add comments about further TODos.

CMakeLists.txt

@@ -132,6 +132,7 @@ file(GLOB_RECURSE COMPILER_SRCS
     src/autotvm/*.cc
     src/tir/*.cc
     src/driver/*.cc
+    src/printer/*.cc
     src/api/*.cc
     )
 

include/tvm/ir/module.h

@@ -308,5 +308,20 @@ class IRModule : public ObjectRef {
   TVM_DLL static IRModule FromText(const std::string& text, const std::string& source_path);
 };
 
+/*!
+ * \brief Render the node as a string in the text format.
+ *
+ * \param node The node to be rendered.
+ * \param show_meta_data Whether to print meta data section.
+ * \param annotate An optional callback function for attaching
+ *        additional comment block to an expr.
+ *
+ * \note We support a limited set of IR nodes that are part of
+ *       relay IR and
+ * \return The text representation.
+ */
+TVM_DLL std::string AsText(const ObjectRef& node,
+                           bool show_meta_data = true,
+                           runtime::TypedPackedFunc<std::string(ObjectRef)> annotate = nullptr);
 }  // namespace tvm
 #endif  // TVM_IR_MODULE_H_

include/tvm/relay/expr.h

@@ -542,18 +542,6 @@ class TempExpr : public Expr {
 /*! \brief Pretty print a Relay node, producing a fragment of the Relay text format. */
 std::string PrettyPrint(const ObjectRef& node);
 
-/*!
- * \brief Render the node as a string in the Relay text format.
- * \param node The node to be rendered.
- * \param show_meta_data Whether to print meta data section.
- * \param annotate An optional callback function for attaching
- *        additional comment block to an expr.
- * \return The text representation.
- */
-std::string AsText(const ObjectRef& node,
-                   bool show_meta_data = true,
-                   runtime::TypedPackedFunc<std::string(Expr)> annotate = nullptr);
-
 /*! \brief namespace of the attributes that are attached to a function. */
 namespace attr {
 /*! \brief Mark the function as a primitive function. */

src/ir/error.cc

@@ -111,7 +111,7 @@ void ErrorReporter::RenderErrors(const IRModule& module, bool use_color) {
     //
     // The annotation callback will annotate the error messages
     // contained in the map.
-    annotated_prog << relay::AsText(func, false, [&err_map](tvm::relay::Expr expr) {
+    annotated_prog << AsText(func, false, [&err_map](const ObjectRef& expr) {
       auto it = err_map.find(expr);
       if (it != err_map.end()) {
         CHECK_NE(it->second.size(), 0);

src/printer/doc.cc

@@ -0,0 +1,173 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file src/tvm/relay/doc.cc
+ * \brief Doc ADT used for pretty printing.
+ *
+ *  Reference: Philip Wadler. A Prettier Printer. Journal of Functional Programming'98
+ */
+#include <tvm/runtime/packed_func.h>
+#include <vector>
+#include <sstream>
+#include "doc.h"
+
+namespace tvm {
+
+/*!
+ * \brief Represent a piece of text in the doc.
+ */
+class DocTextNode : public DocAtomNode {
+ public:
+  /*! \brief The str content in the text. */
+  std::string str;
+
+  explicit DocTextNode(std::string str_val)
+      : str(str_val) {
+    if (str.find_first_of("\t\n") != str.npos) {
+      LOG(WARNING) << "text node: '" << str << "' should not has tab or newline.";
+    }
+  }
+
+  static constexpr const char* _type_key = "printer.DocText";
+  TVM_DECLARE_FINAL_OBJECT_INFO(DocTextNode, DocAtomNode);
+};
+
+TVM_REGISTER_OBJECT_TYPE(DocTextNode);
+
+class DocText : public DocAtom {
+ public:
+  explicit DocText(std::string str) {
+    data_ = runtime::make_object<DocTextNode>(str);
+  }
+
+  TVM_DEFINE_OBJECT_REF_METHODS(DocText, DocAtom, DocTextNode);
+};
+
+/*!
+ * \brief Represent a line breaker in the doc.
+ */
+class DocLineNode : public DocAtomNode {
+ public:
+  /*! \brief The amount of indent in newline. */
+  int indent;
+
+  explicit DocLineNode(int indent)
+      : indent(indent) {}
+
+  static constexpr const char* _type_key = "printer.DocLine";
+  TVM_DECLARE_FINAL_OBJECT_INFO(DocLineNode, DocAtomNode);
+};
+
+TVM_REGISTER_OBJECT_TYPE(DocLineNode);
+
+class DocLine : public DocAtom {
+ public:
+  explicit DocLine(int indent) {
+    data_ = runtime::make_object<DocLineNode>(indent);
+  }
+
+  TVM_DEFINE_OBJECT_REF_METHODS(DocLine, DocAtom, DocLineNode);
+};
+
+// DSL function implementations
+Doc& Doc::operator<<(const Doc& right) {
+  CHECK(this != &right);
+  this->stream_.insert(
+      this->stream_.end(), right.stream_.begin(), right.stream_.end());
+  return *this;
+}
+
+Doc& Doc::operator<<(std::string right) {
+  return *this << DocText(right);
+}
+
+Doc& Doc::operator<<(const DocAtom& right) {
+  this->stream_.push_back(right);
+  return *this;
+}
+
+std::string Doc::str() {
+  std::ostringstream os;
+  for (auto atom : this->stream_) {
+    if (auto* text = atom.as<DocTextNode>()) {
+      os << text->str;
+    } else if (auto* line = atom.as<DocLineNode>()) {
+      os << "\n" << std::string(line->indent, ' ');
+    } else {
+      LOG(FATAL) << "do not expect type " << atom->GetTypeKey();
+    }
+  }
+  return os.str();
+}
+
+Doc Doc::NewLine(int indent) {
+  return Doc() << DocLine(indent);
+}
+
+Doc Doc::Text(std::string text) {
+  return Doc() << DocText(text);
+}
+
+Doc Doc::Indent(int indent, Doc doc) {
+  for (size_t i = 0; i < doc.stream_.size(); ++i) {
+    if (auto* line = doc.stream_[i].as<DocLineNode>()) {
+      doc.stream_[i] = DocLine(indent + line->indent);
+    }
+  }
+  return doc;
+}
+
+Doc Doc::StrLiteral(const std::string& value, std::string quote) {
+  // TODO(M.K.): add escape.
+  Doc doc;
+  return doc << quote << value << quote;
+}
+
+Doc Doc::PyBoolLiteral(bool value) {
+  if (value) {
+    return Doc::Text("True");
+  } else {
+    return Doc::Text("False");
+  }
+}
+
+Doc Doc::Brace(std::string open,
+               const Doc& body,
+               std::string close,
+               int indent) {
+  Doc doc;
+  doc << open;
+  doc << Indent(indent, NewLine() << body) << NewLine();
+  doc << close;
+  return doc;
+}
+
+Doc Doc::Concat(const std::vector<Doc>& vec, const Doc& sep) {
+  Doc seq;
+  if (vec.size() != 0) {
+    if (vec.size() == 1) return vec[0];
+    seq << vec[0];
+    for (size_t i = 1; i < vec.size(); ++i) {
+      seq << sep << vec[i];
+    }
+  }
+  return seq;
+}
+}  // namespace tvm

src/printer/doc.h

@@ -0,0 +1,165 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/printer/doc.h
+ * \brief Doc ADT used for pretty printing.
+ *
+ *  Reference: Philip Wadler. A Prettier Printer. Journal of Functional Programming'98
+ */
+#ifndef TVM_PRINTER_DOC_H_
+#define TVM_PRINTER_DOC_H_
+
+#include <tvm/runtime/data_type.h>
+#include <tvm/runtime/object.h>
+#include <tvm/node/node.h>
+#include <string>
+#include <vector>
+#include <type_traits>
+
+namespace tvm {
+
+/*!
+ * \brief Doc atom node for the ADT.
+ * \sa DocAtom
+ */
+class DocAtomNode : public Object {
+ public:
+  static constexpr const char* _type_key = "printer.DocAtom";
+  TVM_DECLARE_BASE_OBJECT_INFO(DocAtomNode, Object);
+};
+
+/*!
+ * \brief Managed reference to DocAtomNode.
+ * \sa DocAtomNode.
+*/
+class DocAtom : public ObjectRef {
+ public:
+  TVM_DEFINE_OBJECT_REF_METHODS(DocAtom, ObjectRef, DocAtomNode);
+};
+
+/*!
+ * \brief Stream-like interface for Doc DSL.
+ *
+ * The Doc DSL de-couples the layout decision from the printing decision.
+ *
+ * The layout(code formating) decisions include:
+ * - Change indentation.
+ * - Break single line into multiple ones(subjected to future improvements).
+ */
+class Doc {
+ public:
+  /*! \brief default constructor */
+  Doc() {}
+  /*!
+   * \brief Append right to the end of the current doc stream.
+   * \param right The doc to be appended.
+   * \return reference to self.
+   */
+  Doc& operator<<(const Doc& right);
+  /*!
+   * \brief Append right to the end of the current doc stream.
+   * \param right The doc to be appended.
+   * \return reference to self.
+   * \note pass by value to allow copy elison optimization.
+   */
+  Doc& operator<<(std::string right);
+  /*!
+   * \brief Append right to the end of the current doc stream.
+   * \param right The doc to be appended.
+   * \return reference to self.
+   */
+  Doc& operator<<(const DocAtom& right);
+  /*!
+   * \brief Convert value to string via std::ostreamstream
+   *        the append to the current doc stream.
+   * \param right The doc to be appended.
+   * \tparam T the type of the value.
+   * \return reference to self.
+   */
+  template<typename T,
+           typename = typename std::enable_if<!std::is_class<T>::value>::type>
+  Doc& operator<<(const T& value) {
+    std::ostringstream os;
+    os << value;
+    return *this << os.str();
+  }
+  /*!
+   * \brief Convert the doc stream into string.
+   * \return The string representation.
+   */
+  std::string str();
+  /*!
+   * \brief Create a doc that represents text content.
+   * \return The created doc.
+   */
+  static Doc Text(std::string value);
+  /*!
+   * \brief Create a doc that represents a new line.
+   * \return The created doc.
+   */
+  static Doc NewLine(int indent = 0);
+  /*!
+   * \brief Create a new doc that adds indentation to everyline of the doc.
+   * \param indent The indent to be added.
+   * \param doc The doc to be indented.
+   * \return The created doc.
+   * \note pass by value to allow copy elison optimization.
+   */
+  static Doc Indent(int indent, Doc doc);
+  /*!
+   * \brief Create a Doc that represents a string literal.
+   * \param value The content of the string literal.
+   * \param quote The quote in the literal.
+   * \return The created doc.
+   */
+  static Doc StrLiteral(const std::string& value, std::string quote = "\"");
+  /*!
+   * \brief Create a Doc that represents a boolean literal in python syntax.
+   * \param value The bool value.
+   * \return The created doc.
+   */
+  static Doc PyBoolLiteral(bool value);
+  /*!
+   * \brief Enclose body by brace and add indent.
+   * \param body The body
+   * \param open The open brace.
+   * \param close The close brace.
+   * \param indent amount of indentation.
+   * \return The created doc.
+   */
+  static Doc Brace(std::string open,
+                   const Doc& body,
+                   std::string close,
+                   int indent = 2);
+  /*!
+   * \brief Create a doc by concatenating  together with separator.
+   * \param vec The docs to be concatenated.
+   * \param sep The seperator.
+   * \return The created doc.
+   */
+  static Doc Concat(const std::vector<Doc>& vec, const Doc& sep = Text(", "));
+
+ private:
+  /*! \brief Internal doc stream. */
+  std::vector<DocAtom> stream_;
+};
+
+}  // namespace tvm
+#endif  // TVM_PRINTER_DOC_H_