Use visitor pattern to walk AST

src/CMakeLists.txt

@@ -10,4 +10,4 @@ flex_target(lexer lexer.l ${CMAKE_CURRENT_BINARY_DIR}/lex.yy.cc)
 include_directories(${CMAKE_CURRENT_BINARY_DIR})
 include_directories(${CMAKE_CURRENT_SOURCE_DIR})
 
-add_executable(bpftrace main.cpp ast.cpp parser.tab.cc lex.yy.cc)
+add_executable(bpftrace main.cpp parser.tab.cc lex.yy.cc ast.cpp printer.cpp)

src/ast.cpp

 #include "ast.h"
-#include "parser.tab.hh"
 
 namespace ebpf {
 namespace bpftrace {
 namespace ast {
 
-void Integer::print_ast(std::ostream &out, unsigned int depth) const
-{
-  std::string indent(depth, ' ');
-  out << indent << "int: " << n << std::endl;
+void Integer::accept(Visitor &v) {
+  v.visit(*this);
 }
 
-void Variable::print_ast(std::ostream &out, unsigned int depth) const
-{
-  std::string indent(depth, ' ');
-  out << indent << "var: " << ident << std::endl;
+void Variable::accept(Visitor &v) {
+  v.visit(*this);
+  ++v.depth_;
   if (vargs) {
     for (Expression *expr : *vargs) {
-      expr->print_ast(out, depth+1);
+      expr->accept(v);
     }
   }
+  --v.depth_;
 }
 
-void Binop::print_ast(std::ostream &out, unsigned int depth) const
-{
-  std::string indent(depth, ' ');
-  std::string opstr;
-  switch (op) {
-    case ebpf::bpftrace::Parser::token::EQ:
-      opstr = "==";
-      break;
-    case ebpf::bpftrace::Parser::token::NE:
-      opstr = "!=";
-      break;
-    case ebpf::bpftrace::Parser::token::LE:
-      opstr = "<=";
-      break;
-    case ebpf::bpftrace::Parser::token::GE:
-      opstr = ">=";
-      break;
-    case ebpf::bpftrace::Parser::token::LT:
-      opstr = "<";
-      break;
-    case ebpf::bpftrace::Parser::token::GT:
-      opstr = ">";
-      break;
-    case ebpf::bpftrace::Parser::token::LAND:
-      opstr = "&&";
-      break;
-    case ebpf::bpftrace::Parser::token::LOR:
-      opstr = "||";
-      break;
-    case ebpf::bpftrace::Parser::token::PLUS:
-      opstr = "+";
-      break;
-    case ebpf::bpftrace::Parser::token::MINUS:
-      opstr = "-";
-      break;
-    case ebpf::bpftrace::Parser::token::MUL:
-      opstr = "*";
-      break;
-    case ebpf::bpftrace::Parser::token::DIV:
-      opstr = "/";
-      break;
-    case ebpf::bpftrace::Parser::token::MOD:
-      opstr = "%";
-      break;
-    case ebpf::bpftrace::Parser::token::BAND:
-      opstr = "&";
-      break;
-    case ebpf::bpftrace::Parser::token::BOR:
-      opstr = "|";
-      break;
-    case ebpf::bpftrace::Parser::token::BXOR:
-      opstr = "^";
-      break;
-    default:
-      opstr = "???";
-      break;
-  }
-
-  out << indent << opstr << std::endl;
-  left->print_ast(out, depth+1);
-  right->print_ast(out, depth+1);
+void Binop::accept(Visitor &v) {
+  v.visit(*this);
+  ++v.depth_;
+  left->accept(v);
+  right->accept(v);
+  --v.depth_;
 }
 
-void Unop::print_ast(std::ostream &out, unsigned int depth) const
-{
-  std::string indent(depth, ' ');
-  std::string opstr;
-  switch (op) {
-    case ebpf::bpftrace::Parser::token::LNOT:
-      opstr = "!";
-      break;
-    case ebpf::bpftrace::Parser::token::BNOT:
-      opstr = "~";
-      break;
-    default:
-      opstr = "???";
-      break;
-  }
-
-  out << indent << opstr << std::endl;
-  expr->print_ast(out, depth+1);
+void Unop::accept(Visitor &v) {
+  v.visit(*this);
+  ++v.depth_;
+  expr->accept(v);
+  --v.depth_;
 }
 
-void ExprStatement::print_ast(std::ostream &out, unsigned int depth) const
-{
-  std::string indent(depth, ' ');
-  expr->print_ast(out, depth);
+void ExprStatement::accept(Visitor &v) {
+  v.visit(*this);
+  ++v.depth_;
+  expr->accept(v);
+  --v.depth_;
 }
 
-void AssignStatement::print_ast(std::ostream &out, unsigned int depth) const
-{
-  std::string indent(depth, ' ');
-  out << indent << "=" << std::endl;
-  var->print_ast(out, depth+1);
-  expr->print_ast(out, depth+1);
+void AssignStatement::accept(Visitor &v) {
+  v.visit(*this);
+  ++v.depth_;
+  var->accept(v);
+  expr->accept(v);
+  --v.depth_;
 }
 
-void Probe::print_ast(std::ostream &out, unsigned int depth) const
-{
-  std::string indent(depth, ' ');
-  out << indent << type << ":" << attach_point << std::endl;
-  out << indent << " pred" << std::endl;
+void Probe::accept(Visitor &v) {
+  v.visit(*this);
+  ++v.depth_;
   if (pred) {
-    pred->print_ast(out, depth+2);
+    pred->accept(v);
   }
   for (Statement *stmt : *stmts) {
-    stmt->print_ast(out, depth+1);
+    stmt->accept(v);
   }
+  --v.depth_;
 }
 
-void Program::print_ast(std::ostream &out, unsigned int depth) const
-{
-  std::string indent(depth, ' ');
-  out << indent << "Program" << std::endl;
+void Program::accept(Visitor &v) {
+  v.visit(*this);
+  ++v.depth_;
   for (Probe *probe : *probes) {
-    probe->print_ast(out, depth+1);
+    probe->accept(v);
   }
+  --v.depth_;
 }
 
 } // namespace ast

src/ast.h

@@ -2,16 +2,17 @@
 
 #include <string>
 #include <vector>
-#include <ostream>
 
 namespace ebpf {
 namespace bpftrace {
 namespace ast {
 
+class Visitor;
+
 class Node {
 public:
   virtual ~Node() { }
-  virtual void print_ast(std::ostream &out, unsigned int depth = 0) const = 0;
+  virtual void accept(Visitor &v) = 0;
 };
 
 class Expression : public Node {
@@ -21,33 +22,37 @@ using ExpressionList = std::vector<Expression *>;
 class Integer : public Expression {
 public:
   explicit Integer(int n) : n(n) { }
-  void print_ast(std::ostream &out, unsigned int depth = 0) const override;
   int n;
+
+  void accept(Visitor &v) override;
 };
 
 class Variable : public Expression {
 public:
   explicit Variable(std::string &ident) : ident(ident), vargs(nullptr) { }
   Variable(std::string &ident, ExpressionList *vargs) : ident(ident), vargs(vargs) { }
-  void print_ast(std::ostream &out, unsigned int depth = 0) const override;
   std::string ident;
   ExpressionList *vargs;
+
+  void accept(Visitor &v) override;
 };
 
 class Binop : public Expression {
 public:
   Binop(Expression *left, int op, Expression *right) : left(left), right(right), op(op) { }
-  void print_ast(std::ostream &out, unsigned int depth = 0) const override;
   Expression *left, *right;
   int op;
+
+  void accept(Visitor &v) override;
 };
 
 class Unop : public Expression {
 public:
   Unop(int op, Expression *expr) : expr(expr), op(op) { }
-  void print_ast(std::ostream &out, unsigned int depth = 0) const override;
   Expression *expr;
   int op;
+
+  void accept(Visitor &v) override;
 };
 
 class Statement : public Node {
@@ -57,16 +62,18 @@ using StatementList = std::vector<Statement *>;
 class ExprStatement : public Statement {
 public:
   explicit ExprStatement(Expression *expr) : expr(expr) { }
-  void print_ast(std::ostream &out, unsigned int depth = 0) const override;
   Expression *expr;
+
+  void accept(Visitor &v) override;
 };
 
 class AssignStatement : public Statement {
 public:
   AssignStatement(Variable *var, Expression *expr) : var(var), expr(expr) { }
-  void print_ast(std::ostream &out, unsigned int depth = 0) const override;
   Variable *var;
   Expression *expr;
+
+  void accept(Visitor &v) override;
 };
 
 class Probe : public Node {
@@ -75,20 +82,37 @@ public:
     : type(type), attach_point(attach_point), pred(nullptr), stmts(stmts) { }
   Probe(std::string &type, std::string &attach_point, Expression *pred, StatementList *stmts)
     : type(type), attach_point(attach_point), pred(pred), stmts(stmts) { }
-  void print_ast(std::ostream &out, unsigned int depth = 0) const override;
 
   std::string type;
   std::string attach_point;
   Expression *pred;
   StatementList *stmts;
+
+  void accept(Visitor &v) override;
 };
 using ProbeList = std::vector<Probe *>;
 
 class Program : public Node {
 public:
   explicit Program(ProbeList *probes) : probes(probes) { }
-  void print_ast(std::ostream &out, unsigned int depth = 0) const override;
   ProbeList *probes;
+
+  void accept(Visitor &v) override;
+};
+
+class Visitor {
+public:
+  virtual ~Visitor() { }
+  virtual void visit(Integer &integer) = 0;
+  virtual void visit(Variable &var) = 0;
+  virtual void visit(Binop &binop) = 0;
+  virtual void visit(Unop &unop) = 0;
+  virtual void visit(ExprStatement &expr) = 0;
+  virtual void visit(AssignStatement &assignment) = 0;
+  virtual void visit(Probe &probe) = 0;
+  virtual void visit(Program &program) = 0;
+
+  int depth_ = 0;
 };
 
 } // namespace ast

src/main.cpp

+#include <iostream>
 #include "driver.h"
 #include "ast.h"
+#include "printer.h"
 
 int main(int argc, char *argv[])
 {
@@ -12,8 +14,10 @@ int main(int argc, char *argv[])
     result = driver.parse(argv[1]);
   }
 
-  if (!result)
-    driver.root_->print_ast(std::cout);
+  if (!result) {
+    ebpf::bpftrace::ast::Printer p = ebpf::bpftrace::ast::Printer(std::cout);
+    driver.root_->accept(p);
+  }
 
   return result;
 }

src/printer.cpp

+#include "printer.h"
+#include "ast.h"
+#include "parser.tab.hh"
+
+namespace ebpf {
+namespace bpftrace {
+namespace ast {
+
+void Printer::visit(Integer &integer)
+{
+  std::string indent(depth_, ' ');
+  out_ << indent << "int: " << integer.n << std::endl;
+}
+
+void Printer::visit(Variable &var)
+{
+  std::string indent(depth_, ' ');
+  out_ << indent << "var: " << var.ident << std::endl;
+}
+
+void Printer::visit(Binop &binop)
+{
+  std::string indent(depth_, ' ');
+  std::string opstr;
+  switch (binop.op) {
+    case ebpf::bpftrace::Parser::token::EQ:
+      opstr = "==";
+      break;
+    case ebpf::bpftrace::Parser::token::NE:
+      opstr = "!=";
+      break;
+    case ebpf::bpftrace::Parser::token::LE:
+      opstr = "<=";
+      break;
+    case ebpf::bpftrace::Parser::token::GE:
+      opstr = ">=";
+      break;
+    case ebpf::bpftrace::Parser::token::LT:
+      opstr = "<";
+      break;
+    case ebpf::bpftrace::Parser::token::GT:
+      opstr = ">";
+      break;
+    case ebpf::bpftrace::Parser::token::LAND:
+      opstr = "&&";
+      break;
+    case ebpf::bpftrace::Parser::token::LOR:
+      opstr = "||";
+      break;
+    case ebpf::bpftrace::Parser::token::PLUS:
+      opstr = "+";
+      break;
+    case ebpf::bpftrace::Parser::token::MINUS:
+      opstr = "-";
+      break;
+    case ebpf::bpftrace::Parser::token::MUL:
+      opstr = "*";
+      break;
+    case ebpf::bpftrace::Parser::token::DIV:
+      opstr = "/";
+      break;
+    case ebpf::bpftrace::Parser::token::MOD:
+      opstr = "%";
+      break;
+    case ebpf::bpftrace::Parser::token::BAND:
+      opstr = "&";
+      break;
+    case ebpf::bpftrace::Parser::token::BOR:
+      opstr = "|";
+      break;
+    case ebpf::bpftrace::Parser::token::BXOR:
+      opstr = "^";
+      break;
+    default:
+      opstr = "???";
+      break;
+  }
+
+  out_ << indent << opstr << std::endl;
+}
+
+void Printer::visit(Unop &unop)
+{
+  std::string indent(depth_, ' ');
+  std::string opstr;
+  switch (unop.op) {
+    case ebpf::bpftrace::Parser::token::LNOT:
+      opstr = "!";
+      break;
+    case ebpf::bpftrace::Parser::token::BNOT:
+      opstr = "~";
+      break;
+    default:
+      opstr = "???";
+      break;
+  }
+
+  out_ << indent << opstr << std::endl;
+}
+
+void Printer::visit(ExprStatement &)
+{
+}
+
+void Printer::visit(AssignStatement &)
+{
+  std::string indent(depth_, ' ');
+  out_ << indent << "=" << std::endl;
+}
+
+void Printer::visit(Probe &probe)
+{
+  std::string indent(depth_, ' ');
+  out_ << indent << probe.type << ":" << probe.attach_point << std::endl;
+  if (probe.pred) {
+    out_ << indent << " pred";
+  }
+}
+
+void Printer::visit(Program &)
+{
+  std::string indent(depth_, ' ');
+  out_ << indent << "Program" << std::endl;
+}
+
+} // namespace ast
+} // namespace bpftrace
+} // namespace ebpf

src/printer.h

+#pragma once
+
+#include <ostream>
+#include "ast.h"
+
+namespace ebpf {
+namespace bpftrace {
+namespace ast {
+
+class Printer : public Visitor {
+public:
+  explicit Printer(std::ostream &out) : out_(out) { }
+
+  void visit(Integer &integer) override;
+  void visit(Variable &var) override;
+  void visit(Binop &binop) override;
+  void visit(Unop &unop) override;
+  void visit(ExprStatement &expr) override;
+  void visit(AssignStatement &assignment) override;
+  void visit(Probe &probe) override;
+  void visit(Program &program) override;
+
+private:
+  std::ostream &out_;
+};
+
+} // namespace ast
+} // namespace bpftrace
+} // namespace ebpf