Add support for dumping DOT graphs from KQIR (#2205)

src/common/type_util.h

@@ -39,3 +39,11 @@ using RemoveCVRef = typename std::remove_cv_t<typename std::remove_reference_t<T
 // dependent false for static_assert with constexpr if, see CWG2518/P2593R1
 template <typename T>
 constexpr bool AlwaysFalse = false;
+
+template <typename>
+struct GetClassFromMember;
+
+template <typename C, typename T>
+struct GetClassFromMember<T C::*> {
+  using type = C;  // NOLINT
+};

src/search/ir.h

@@ -22,33 +22,49 @@
 
 #include <fmt/format.h>
 
+#include <initializer_list>
 #include <limits>
 #include <memory>
 #include <optional>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <variant>
 #include <vector>
 
 #include "fmt/core.h"
+#include "ir_iterator.h"
 #include "string_util.h"
+#include "type_util.h"
 
 // kqir stands for Kvorcks Query Intermediate Representation
 namespace kqir {
 
 struct Node {
   virtual std::string Dump() const = 0;
+  virtual std::string_view Name() const = 0;
+  virtual std::string Content() const { return {}; }
+
+  virtual NodeIterator ChildBegin() { return {}; };
+  virtual NodeIterator ChildEnd() { return {}; };
+
   virtual ~Node() = default;
 
   template <typename T, typename U = Node, typename... Args>
   static std::unique_ptr<U> Create(Args &&...args) {
     return std::unique_ptr<U>(new T(std::forward<Args>(args)...));
   }
 
+  template <typename T>
+  static std::unique_ptr<T> MustAs(std::unique_ptr<Node> &&original) {
+    auto casted = As<T>(std::move(original));
+    CHECK(casted != nullptr);
+    return casted;
+  }
+
   template <typename T>
   static std::unique_ptr<T> As(std::unique_ptr<Node> &&original) {
     auto casted = dynamic_cast<T *>(original.release());
-    CHECK(casted);
     return std::unique_ptr<T>(casted);
   }
 };
@@ -58,15 +74,19 @@ struct FieldRef : Node {
 
   explicit FieldRef(std::string name) : name(std::move(name)) {}
 
+  std::string_view Name() const override { return "FieldRef"; }
   std::string Dump() const override { return name; }
+  std::string Content() const override { return Dump(); }
 };
 
 struct StringLiteral : Node {
   std::string val;
 
   explicit StringLiteral(std::string val) : val(std::move(val)) {}
 
+  std::string_view Name() const override { return "StringLiteral"; }
   std::string Dump() const override { return fmt::format("\"{}\"", util::EscapeString(val)); }
+  std::string Content() const override { return Dump(); }
 };
 
 struct QueryExpr : Node {};
@@ -80,15 +100,21 @@ struct TagContainExpr : BoolAtomExpr {
   TagContainExpr(std::unique_ptr<FieldRef> &&field, std::unique_ptr<StringLiteral> &&tag)
       : field(std::move(field)), tag(std::move(tag)) {}
 
+  std::string_view Name() const override { return "TagContainExpr"; }
   std::string Dump() const override { return fmt::format("{} hastag {}", field->Dump(), tag->Dump()); }
+
+  NodeIterator ChildBegin() override { return {field.get(), tag.get()}; };
+  NodeIterator ChildEnd() override { return {}; };
 };
 
 struct NumericLiteral : Node {
   double val;
 
   explicit NumericLiteral(double val) : val(val) {}
 
+  std::string_view Name() const override { return "NumericLiteral"; }
   std::string Dump() const override { return fmt::format("{}", val); }
+  std::string Content() const override { return Dump(); }
 };
 
 // NOLINTNEXTLINE
@@ -156,15 +182,22 @@ struct NumericCompareExpr : BoolAtomExpr {
     __builtin_unreachable();
   }
 
+  std::string_view Name() const override { return "NumericCompareExpr"; }
   std::string Dump() const override { return fmt::format("{} {} {}", field->Dump(), ToOperator(op), num->Dump()); };
+  std::string Content() const override { return ToOperator(op); }
+
+  NodeIterator ChildBegin() override { return {field.get(), num.get()}; };
+  NodeIterator ChildEnd() override { return {}; };
 };
 
 struct BoolLiteral : BoolAtomExpr {
   bool val;
 
   explicit BoolLiteral(bool val) : val(val) {}
 
+  std::string_view Name() const override { return "BoolLiteral"; }
   std::string Dump() const override { return val ? "true" : "false"; }
+  std::string Content() const override { return Dump(); }
 };
 
 struct QueryExpr;
@@ -174,27 +207,39 @@ struct NotExpr : QueryExpr {
 
   explicit NotExpr(std::unique_ptr<QueryExpr> &&inner) : inner(std::move(inner)) {}
 
+  std::string_view Name() const override { return "NotExpr"; }
   std::string Dump() const override { return fmt::format("not {}", inner->Dump()); }
+
+  NodeIterator ChildBegin() override { return NodeIterator{inner.get()}; };
+  NodeIterator ChildEnd() override { return {}; };
 };
 
 struct AndExpr : QueryExpr {
   std::vector<std::unique_ptr<QueryExpr>> inners;
 
   explicit AndExpr(std::vector<std::unique_ptr<QueryExpr>> &&inners) : inners(std::move(inners)) {}
 
+  std::string_view Name() const override { return "AndExpr"; }
   std::string Dump() const override {
     return fmt::format("(and {})", util::StringJoin(inners, [](const auto &v) { return v->Dump(); }));
   }
+
+  NodeIterator ChildBegin() override { return NodeIterator(inners.begin()); };
+  NodeIterator ChildEnd() override { return NodeIterator(inners.end()); };
 };
 
 struct OrExpr : QueryExpr {
   std::vector<std::unique_ptr<QueryExpr>> inners;
 
   explicit OrExpr(std::vector<std::unique_ptr<QueryExpr>> &&inners) : inners(std::move(inners)) {}
 
+  std::string_view Name() const override { return "OrExpr"; }
   std::string Dump() const override {
     return fmt::format("(or {})", util::StringJoin(inners, [](const auto &v) { return v->Dump(); }));
   }
+
+  NodeIterator ChildBegin() override { return NodeIterator(inners.begin()); };
+  NodeIterator ChildEnd() override { return NodeIterator(inners.end()); };
 };
 
 struct Limit : Node {
@@ -203,7 +248,9 @@ struct Limit : Node {
 
   Limit(size_t offset, size_t count) : offset(offset), count(count) {}
 
+  std::string_view Name() const override { return "Limit"; }
   std::string Dump() const override { return fmt::format("limit {}, {}", offset, count); }
+  std::string Content() const override { return fmt::format("{}, {}", offset, count); }
 };
 
 struct SortBy : Node {
@@ -213,50 +260,72 @@ struct SortBy : Node {
   SortBy(Order order, std::unique_ptr<FieldRef> &&field) : order(order), field(std::move(field)) {}
 
   static constexpr const char *OrderToString(Order order) { return order == ASC ? "asc" : "desc"; }
+
+  std::string_view Name() const override { return "SortBy"; }
   std::string Dump() const override { return fmt::format("sortby {}, {}", field->Dump(), OrderToString(order)); }
+  std::string Content() const override { return OrderToString(order); }
+
+  NodeIterator ChildBegin() override { return NodeIterator(field.get()); };
+  NodeIterator ChildEnd() override { return {}; };
 };
 
 struct SelectExpr : Node {
   std::vector<std::unique_ptr<FieldRef>> fields;
 
   explicit SelectExpr(std::vector<std::unique_ptr<FieldRef>> &&fields) : fields(std::move(fields)) {}
 
+  std::string_view Name() const override { return "SelectExpr"; }
   std::string Dump() const override {
     if (fields.empty()) return "select *";
     return fmt::format("select {}", util::StringJoin(fields, [](const auto &v) { return v->Dump(); }));
   }
+
+  NodeIterator ChildBegin() override { return NodeIterator(fields.begin()); };
+  NodeIterator ChildEnd() override { return NodeIterator(fields.end()); };
 };
 
 struct IndexRef : Node {
   std::string name;
 
   explicit IndexRef(std::string name) : name(std::move(name)) {}
 
+  std::string_view Name() const override { return "IndexRef"; }
   std::string Dump() const override { return name; }
+  std::string Content() const override { return Dump(); }
 };
 
 struct SearchStmt : Node {
+  std::unique_ptr<SelectExpr> select_expr;
   std::unique_ptr<IndexRef> index;
   std::unique_ptr<QueryExpr> query_expr;  // optional
   std::unique_ptr<Limit> limit;           // optional
   std::unique_ptr<SortBy> sort_by;        // optional
-  std::unique_ptr<SelectExpr> select_expr;
 
   SearchStmt(std::unique_ptr<IndexRef> &&index, std::unique_ptr<QueryExpr> &&query_expr, std::unique_ptr<Limit> &&limit,
              std::unique_ptr<SortBy> &&sort_by, std::unique_ptr<SelectExpr> &&select_expr)
-      : index(std::move(index)),
+      : select_expr(std::move(select_expr)),
+        index(std::move(index)),
         query_expr(std::move(query_expr)),
         limit(std::move(limit)),
-        sort_by(std::move(sort_by)),
-        select_expr(std::move(select_expr)) {}
+        sort_by(std::move(sort_by)) {}
 
+  std::string_view Name() const override { return "SearchStmt"; }
   std::string Dump() const override {
     std::string opt;
     if (query_expr) opt += " where " + query_expr->Dump();
     if (sort_by) opt += " " + sort_by->Dump();
     if (limit) opt += " " + limit->Dump();
     return fmt::format("{} from {}{}", select_expr->Dump(), index->Dump(), opt);
   }
+
+  static inline const std::vector<std::function<Node *(Node *)>> ChildMap = {
+      NodeIterator::MemFn<&SearchStmt::select_expr>, NodeIterator::MemFn<&SearchStmt::index>,
+      NodeIterator::MemFn<&SearchStmt::query_expr>,  NodeIterator::MemFn<&SearchStmt::limit>,
+      NodeIterator::MemFn<&SearchStmt::sort_by>,
+  };
+
+  NodeIterator ChildBegin() override { return NodeIterator(this, ChildMap.begin()); };
+  NodeIterator ChildEnd() override { return NodeIterator(this, ChildMap.end()); };
 };
 
 }  // namespace kqir

src/search/ir_dot_dumper.h

@@ -0,0 +1,54 @@
+/*
+ * 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.
+ *
+ */
+
+#pragma once
+
+#include "ir.h"
+#include "string_util.h"
+
+namespace kqir {
+
+struct DotDumper {
+  std::ostream &os;
+
+  void Dump(Node *node) {
+    os << "digraph {\n";
+    dump(node);
+    os << "}\n";
+  }
+
+ private:
+  static std::string nodeId(Node *node) { return fmt::format("x{:x}", (uint64_t)node); }
+
+  void dump(Node *node) {
+    os << "  " << nodeId(node) << " [ label = \"" << node->Name();
+    if (auto content = node->Content(); !content.empty()) {
+      os << " (" << util::EscapeString(content) << ")\" ];\n";
+    } else {
+      os << "\" ];\n";
+    }
+    for (auto i = node->ChildBegin(); i != node->ChildEnd(); ++i) {
+      os << "  " << nodeId(node) << " -> " << nodeId(*i) << ";\n";
+      dump(*i);
+    }
+  }
+};
+
+}  // namespace kqir