[IR] Refactor IR Printer: HLOExpr to Doc (#198)

* [IR] Refactor IR Printer: HLOExpr to Doc * typo * missing files
bytedance · Mar 17, 2023 · 8a816ba · 8a816ba
1 parent 23d0e01
commit 8a816ba
Show file tree

Hide file tree

Showing 11 changed files with 276 additions and 0 deletions.
diff --git a/include/matxscript/ir/printer/doc.h b/include/matxscript/ir/printer/doc.h
@@ -310,6 +310,19 @@ class LiteralDoc : public ExprDoc {
     return LiteralDoc::Str(dtype, p);
   }
 
+  /*!
+   * \brief Create a LiteralDoc to represent string.
+   * \param v The string value.
+   * \param p The object path
+   */
+  static LiteralDoc HLOType(const Type& v, const Optional<ObjectPath>& p) {
+    if (auto const* pt = v.as<PrimTypeNode>()) {
+      return LiteralDoc::DataType(pt->dtype, p);
+    }
+    StringRef dtype = v->GetPythonTypeName().encode();
+    return LiteralDoc::Str(dtype, p);
+  }
+
   MATXSCRIPT_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LiteralDoc, ExprDoc, LiteralDocNode);
 };
 

diff --git a/src/ir/_base/cow_array_ref.cc b/src/ir/_base/cow_array_ref.cc
@@ -29,6 +29,8 @@
 #include <matxscript/ir/_base/repr_printer.h>
 #include <matxscript/ir/_base/structural_equal.h>
 #include <matxscript/ir/_base/structural_hash.h>
+#include <matxscript/ir/printer/doc.h>
+#include <matxscript/ir/printer/ir_docsifier.h>
 #include <matxscript/runtime/functor.h>
 #include <matxscript/runtime/registry.h>
 
@@ -197,5 +199,19 @@ MATXSCRIPT_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
       p->stream << ']';
     });
 
+using namespace ::matxscript::ir::printer;
+MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
+    .set_dispatch<Array<ObjectRef>>(  //
+        "",
+        [](Array<ObjectRef> array, ObjectPath p, IRDocsifier d) -> Doc {
+          int n = array.size();
+          Array<ExprDoc> results;
+          results.reserve(n);
+          for (int i = 0; i < n; ++i) {
+            results.push_back(d->AsDoc<ExprDoc>(array[i], p->ArrayIndex(i)));
+          }
+          return ListDoc(results);
+        });
+
 }  // namespace ir
 }  // namespace matxscript
diff --git a/src/ir/_base/cow_map_ref.cc b/src/ir/_base/cow_map_ref.cc
@@ -29,6 +29,8 @@
 #include <matxscript/ir/_base/repr_printer.h>
 #include <matxscript/ir/_base/structural_equal.h>
 #include <matxscript/ir/_base/structural_hash.h>
+#include <matxscript/ir/printer/doc.h>
+#include <matxscript/ir/printer/ir_docsifier.h>
 #include <matxscript/runtime/functor.h>
 #include <matxscript/runtime/registry.h>
 
@@ -364,5 +366,41 @@ MATXSCRIPT_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 
 MATX_DLL constexpr uint64_t DenseMapNode::kNextProbeLocation[];
 
+using namespace ::matxscript::ir::printer;
+MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
+    .set_dispatch<Map<ObjectRef, ObjectRef>>(  //
+        "",
+        [](Map<ObjectRef, ObjectRef> dict, ObjectPath p, IRDocsifier d) -> Doc {
+          using POO = std::pair<ObjectRef, ObjectRef>;
+          std::vector<POO> items{dict.begin(), dict.end()};
+          bool is_str_map = true;
+          for (const auto& kv : items) {
+            if (!kv.first.as<StringNode>()) {
+              is_str_map = false;
+              break;
+            }
+          }
+          if (is_str_map) {
+            std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) {
+              return runtime::Downcast<StringRef>(lhs.first) <
+                     runtime::Downcast<StringRef>(rhs.first);
+            });
+          } else {
+            std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) {
+              return lhs.first.get() < rhs.first.get();
+            });
+          }
+          int n = dict.size();
+          Array<ExprDoc> ks;
+          Array<ExprDoc> vs;
+          ks.reserve(n);
+          vs.reserve(n);
+          for (int i = 0; i < n; ++i) {
+            ks.push_back(d->AsDoc<ExprDoc>(items[i].first, p->MissingMapEntry()));
+            vs.push_back(d->AsDoc<ExprDoc>(items[i].second, p->MapValue(items[i].first)));
+          }
+          return DictDoc(ks, vs);
+        });
+
 }  // namespace ir
 }  // namespace matxscript
diff --git a/src/ir/_base/string_ref.cc b/src/ir/_base/string_ref.cc
@@ -23,6 +23,8 @@
 #include <initializer_list>
 
 #include <matxscript/ir/_base/reflection.h>
+#include <matxscript/ir/printer/doc.h>
+#include <matxscript/ir/printer/ir_docsifier.h>
 #include <matxscript/runtime/container/string_helper.h>
 #include <matxscript/runtime/object.h>
 #include <matxscript/runtime/registry.h>
@@ -262,5 +264,12 @@ typename StringRef::const_reverse_iterator StringRef::rend() const {
   return const_reverse_iterator(begin());
 }
 
+using namespace ::matxscript::ir::printer;
+MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
+    .set_dispatch<StringRef>("", [](StringRef s, ObjectPath p, IRDocsifier d) -> Doc {
+      // TODO: optimize MultipleLines
+      return LiteralDoc::Str(s, p);
+    });
+
 }  // namespace ir
 }  // namespace matxscript
diff --git a/src/ir/adt.cc b/src/ir/adt.cc
@@ -30,13 +30,16 @@
  */
 #include <matxscript/ir/adt.h>
 
+#include <matxscript/ir/printer/doc.h>
+#include <matxscript/ir/printer/ir_docsifier.h>
 #include <matxscript/ir/type.h>
 #include <matxscript/runtime/registry.h>
 
 namespace matxscript {
 namespace ir {
 
 using namespace runtime;
+using namespace ::matxscript::ir::printer;
 
 Constructor::Constructor(Type ret_type,
                          StringRef name_hint,
@@ -66,6 +69,11 @@ MATXSCRIPT_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
                 << node->belong_to << ")";
     });
 
+MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable)  //
+    .set_dispatch<Constructor>("", [](Constructor node, ObjectPath p, IRDocsifier d) -> Doc {
+      return IdDoc(node->name_hint);
+    });
+
 // ClassType
 ClassType::ClassType(uint64_t py_type_id,
                      GlobalTypeVar header,