Pretty print tree using Visitors

visitor-tree-pretty-print.cpp

#include <cassert>
#include <cstdint>
#include <iostream>
#include <memory>
#include <sstream>
#include <utility>
#include <vector>

// Visited hierarchy
struct AstBase {
  const enum class Kind { BinaryExpr, LiteralExpr } kind;

  AstBase(Kind kind) : kind(kind) {}
};

struct AstExpr : public AstBase {
  AstExpr(Kind kind) : AstBase(kind) {}
};

struct AstBinaryExpr : public AstExpr {
  std::unique_ptr<AstExpr> left;
  char op;
  std::unique_ptr<AstExpr> right;

  AstBinaryExpr(std::unique_ptr<AstExpr> left, char op,
                std::unique_ptr<AstExpr> right)
      : AstExpr(Kind::BinaryExpr), left(std::move(left)), op(op),
        right(std::move(right)) {}
};

struct AstLiteralExpr : public AstExpr {
  int literal;

  AstLiteralExpr(int literal) : AstExpr(Kind::LiteralExpr), literal(literal) {}
};

// Visitor base (using CRTP)
template <typename Derived, typename RetTy = void, typename... ArgTys>
struct AstVisitor {
private:
  Derived &derived() { return *static_cast<Derived *>(this); }

  RetTy visitBinaryExpr(AstBinaryExpr &node, ArgTys... args) {
      visit(*node.left);
      visit(*node.right);

      return RetTy();
  }

  RetTy visitLiteralExpr(AstLiteralExpr &node, ArgTys... args) {
      return RetTy();
  }

public:
  RetTy visit(AstBase &node, ArgTys... args) {
    if (node.kind == AstBase::Kind::BinaryExpr)
      return derived().visitBinaryExpr(static_cast<AstBinaryExpr &>(node),
                                       args...);
    else if (node.kind == AstBase::Kind::LiteralExpr)
      return derived().visitLiteralExpr(static_cast<AstLiteralExpr &>(node),
                                        args...);
    else
      assert(false && "Unhandled AST type!");
  }
};

// Concrete visitors
struct PrettyPrinter
    : public AstVisitor<PrettyPrinter, void, const std::string &, bool> {

  std::ostream &os;
  bool ascii;

  PrettyPrinter(std::ostream &os, bool ascii = false) : os(os), ascii(ascii) {}

  class FieldsStream {
  private:
    std::ostream &os;
    bool first = true;

  public:
    FieldsStream(std::ostream &os) : os(os) {}

    template <typename T> FieldsStream &operator<<(const T &t) {
      os << t;
      return *this;
    }

    template <typename U, typename T>
    FieldsStream &operator<<(const std::pair<U, T> &pair) {
      os << (first ? ": " : ", ") << pair.first << " = '" << pair.second << "'";
      first = false;

      return *this;
    }
  };

  FieldsStream printHelper(std::string &indent, bool last, const std::string& name) {
    os << indent;

    if (last) {
      os << (ascii ? "`-- " : "└── ");
      indent += "    ";
    } else {
      os << (ascii ? "|-- " : "├── ");
      indent += (ascii ? "|   " : "│   ");
    }

    os << name;

    return FieldsStream(os);
  }


  void visitBinaryExpr(AstBinaryExpr &e, std::string indent, bool last) {
    printHelper(indent, last, "BinaryExpr") << std::make_pair("op", e.op) << "\n";

    visit(*e.left, indent, false);
    visit(*e.right, indent, true);
  }

  void visitLiteralExpr(AstLiteralExpr &e, std::string indent, bool last) {
    printHelper(indent, last, "LiteralExpr") << std::make_pair("value", e.literal) << "\n";
  }
};

struct DotPrinter : public AstVisitor<DotPrinter> {

  std::ostream &os;

  DotPrinter(std::ostream &os) : os(os) {}

  void print(AstBase &node) {
    os << "digraph {\n";
    os << "node [shape=record];\n";
    visit(node);
    os << "}\n";
  }

  void node(const AstBase &node, const std::string &name,
            const std::vector<std::pair<std::string, std::string>> &records) {
    os << reinterpret_cast<std::uintptr_t>(&node);
    os << " [label=\"";

    os << "{" << name << "|{";

    bool first = true;
    for (const auto &record : records) {
      if (!first)
        os << "|";
      os << "<" << record.first << ">" << record.second;
      first = false;
    }

    os << "}}";

    os << "\"];\n";
  }

  void edge(const AstBase &from, const std::string part, const AstBase &to) {
    os << reinterpret_cast<std::uintptr_t>(&from) << ":" << part << " -> "
       << reinterpret_cast<std::uintptr_t>(&to) << ";\n";
  }

  void visitBinaryExpr(AstBinaryExpr &e) {
    node(e, "BinaryExpr",
         {{"lhs", "LHS"}, {"op", std::string(1, e.op)}, {"rhs", "RHS"}});

    edge(e, "lhs", *e.left);
    edge(e, "rhs", *e.right);

    visit(*e.left);
    visit(*e.right);
  }

  void visitLiteralExpr(AstLiteralExpr &e) {
    node(e, "LiteralExpr", {{"literal", std::to_string(e.literal)}});
  }
};

// Test code
int main() {
  std::unique_ptr<AstBase> toplevel = std::make_unique<AstBinaryExpr>(
      std::make_unique<AstLiteralExpr>(1), '+',
      std::make_unique<AstBinaryExpr>(std::make_unique<AstLiteralExpr>(2), '*',
                                      std::make_unique<AstLiteralExpr>(3)));

  PrettyPrinter printerUnicode(std::cout, false);
  printerUnicode.visit(*toplevel, "", true);

  PrettyPrinter printerAscii(std::cout, true);
  printerAscii.visit(*toplevel, "", true);

  DotPrinter dotPrinter(std::cout);
  dotPrinter.print(*toplevel);

  return 0;
}