polish + colors

AnyDSL · Mar 13, 2024 · a5c71d2 · a5c71d2
1 parent 035df75
commit a5c71d2
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 46 deletions.
diff --git a/include/thorin/def.h b/include/thorin/def.h
@@ -35,9 +35,15 @@
 namespace thorin {
 
 namespace Node {
+
 #define CODE(node, name) node,
 enum : node_t { THORIN_NODE(CODE) };
 #undef CODE
+
+#define CODE(node, name) +size_t(1)
+constexpr auto Num_Nodes = size_t(0) THORIN_NODE(CODE);
+#undef CODE
+
 } // namespace Node
 
 class App;

diff --git a/src/thorin/dot.cpp b/src/thorin/dot.cpp
@@ -32,7 +32,7 @@ class Dot {
         tab_.println(os_, "digraph {{");
         ++tab_;
         tab_.println(os_, "ordering=out;");
-        tab_.println(os_, "node [shape=box];");
+        tab_.println(os_, "node [shape=box,style=filled];");
     }
 
     void epilogue() {
@@ -47,69 +47,60 @@ class Dot {
     }
 
     void recurse(const Def* def, uint32_t max) {
-        if (max <= 0 || !done_.emplace(def).second) return;
-        auto l     = label(def);
-        auto style = def->isa_mut() ? ",style=bold" : "";
-        tab_.print(os_, "_{}[label=<{}<br/>{}>{}", def->gid(), def->unique_name(), l, style);
-        tooltip(def);
-        println(os_, "];");
-
-        if (def->is_set()) {
-            for (size_t i = 0, e = def->num_ops(); i != e; ++i) {
-                auto op = def->op(i);
-                recurse(op, max - 1);
-                if (op->isa<Lit>() || op->isa<Axiom>() || def->isa<Var>() || def->isa<Nat>() || def->isa<Idx>())
-                    tab_.println(os_, "_{} -> _{}[color=\"#00000000\",constraint=false];", def->gid(), op->gid());
-                else
-                    tab_.println(os_, "_{} -> _{}[label=\"{}\"];", def->gid(), op->gid(), i);
-            }
-            if (auto type = def->type(); type && types_) {
-                recurse(type, max - 1);
-                tab_.println(os_, "_{} -> _{}[color=\"#00000000\",constraint=false,style=dashed];", def->gid(),
-                             type->gid());
-            }
+        if (max == 0 || !done_.emplace(def).second) return;
+        tab_.print(os_, "_{}[", def->gid());
+        if (def->isa_mut()) os_ << "style=\"filled,diagonals\"";
+        label(def) << ',';
+        color(def) << ',';
+        tooltip(def) << "];\n";
+
+        if (!def->is_set()) return;
+
+        for (size_t i = 0, e = def->num_ops(); i != e; ++i) {
+            auto op = def->op(i);
+            recurse(op, max - 1);
+            if (op->isa<Lit>() || op->isa<Axiom>() || def->isa<Var>() || def->isa<Nat>() || def->isa<Idx>())
+                tab_.println(os_, "_{} -> _{}[color=\"#00000000\",constraint=false];", def->gid(), op->gid());
+            else
+                tab_.println(os_, "_{} -> _{}[label=\"{}\"];", def->gid(), op->gid(), i);
         }
-    }
 
-    std::string label(const Def* def) {
-        if (auto lit = def->isa<Lit>()) {
-            std::ostringstream oss;
-            lit->stream(oss, 0);
-            return oss.str();
+        if (auto type = def->type(); type && types_) {
+            recurse(type, max - 1);
+            tab_.println(os_, "_{} -> _{}[color=\"#00000000\",constraint=false,style=dashed];", def->gid(),
+                         type->gid());
         }
-        return std::string(def->node_name());
     }
 
-    void tooltip(const Def* def) {
+    std::ostream& label(const Def* def) {
+        print(os_, "label=<{}<br/>", def->unique_name());
+        if (auto lit = def->isa<Lit>())
+            lit->stream(os_, 0);
+        else
+            os_ << def->node_name();
+        return os_ << '>';
+    }
+
+    std::ostream& color(const Def* def) {
+        return print(os_, "fillcolor=\"{} 0.5 0.75\"", def->node() / (float)Node::Num_Nodes);
+    }
+
+    std::ostream& tooltip(const Def* def) {
         static constexpr auto NL = "&#13;&#10;";
         auto loc                 = escape(def->loc());
         auto type                = escape(def->type());
         std::ostringstream oss;
         oss << std::hex << def->flags();
         auto flags = oss.str();
         escape(loc);
-        print(os_, ",tooltip=\"");
+        print(os_, "tooltip=\"");
         print(os_, "<b>expr:</b> {}{}", def, NL);
         print(os_, "<b>type:</b> {}{}", type, NL);
         print(os_, "<b>name:</b> {}{}", def->sym(), NL);
         print(os_, "<b>gid:</b> {}{}", def->gid(), NL);
         print(os_, "<b>flags:</b> {}{}", flags, NL);
         print(os_, "<b>loc:</b> {}", loc);
-        print(os_, "\"");
-    }
-
-    std::string label(const Def* def) const {
-        if (auto lit = Lit::isa(def)) return std::to_string(*lit);
-        switch (def->node()) {
-            case Node::Axiom: return def->sym().str();
-            case Node::Pi: return "Π";
-            case Node::Lam: return "λ";
-            case Node::App: return "";
-            case Node::Sigma: return "Σ";
-            case Node::Tuple: return "()";
-            case Node::Extract: return "#";
-            default: return std::string(def->node_name());
-        }
+        return print(os_, "\"");
     }
 
 private: