Parser: Accept "#include <file.h>" at start of script

src/ast/ast.cpp

@@ -60,6 +60,10 @@ void Probe::accept(Visitor &v) {
   v.visit(*this);
 }
 
+void Include::accept(Visitor &v) {
+  v.visit(*this);
+}
+
 void Program::accept(Visitor &v) {
   v.visit(*this);
 }

src/ast/ast.h

@@ -176,9 +176,20 @@ public:
 };
 using ProbeList = std::vector<Probe *>;
 
+class Include : public Node {
+public:
+  explicit Include(const std::string &file) : file(file) { }
+  std::string file;
+
+  void accept(Visitor &v) override;
+};
+using IncludeList = std::vector<Include *>;
+
 class Program : public Node {
 public:
-  explicit Program(ProbeList *probes) : probes(probes) { }
+  Program(IncludeList *includes, ProbeList *probes)
+    : includes(includes), probes(probes) { }
+  IncludeList *includes;
   ProbeList *probes;
 
   void accept(Visitor &v) override;
@@ -201,6 +212,7 @@ public:
   virtual void visit(Predicate &pred) = 0;
   virtual void visit(AttachPoint &ap) = 0;
   virtual void visit(Probe &probe) = 0;
+  virtual void visit(Include &include) = 0;
   virtual void visit(Program &program) = 0;
 };
 

src/ast/codegen_llvm.cpp

@@ -410,11 +410,16 @@ void CodegenLLVM::visit(Probe &probe)
   b_.CreateRet(ConstantInt::get(module_->getContext(), APInt(64, 0)));
 }
 
+void CodegenLLVM::visit(Include &include)
+{
+}
+
 void CodegenLLVM::visit(Program &program)
 {
-  for (Probe *probe : *program.probes) {
+  for (Include *include : *program.includes)
+    include->accept(*this);
+  for (Probe *probe : *program.probes)
     probe->accept(*this);
-  }
 }
 
 AllocaInst *CodegenLLVM::getMapKey(Map &map)

src/ast/codegen_llvm.h

@@ -41,6 +41,7 @@ public:
   void visit(Predicate &pred) override;
   void visit(AttachPoint &ap) override;
   void visit(Probe &probe) override;
+  void visit(Include &include) override;
   void visit(Program &program) override;
   AllocaInst *getMapKey(Map &map);
   AllocaInst *getQuantizeMapKey(Map &map, Value *log2);

src/ast/printer.cpp

@@ -131,8 +131,6 @@ void Printer::visit(AttachPoint &ap)
 
 void Printer::visit(Probe &probe)
 {
-  std::string indent(depth_, ' ');
-
   for (AttachPoint *ap : *probe.attach_points) {
     ap->accept(*this);
   }
@@ -147,15 +145,22 @@ void Printer::visit(Probe &probe)
   --depth_;
 }
 
+void Printer::visit(Include &include)
+{
+  std::string indent(depth_, ' ');
+  out_ << indent << "#include " << include.file << std::endl;
+}
+
 void Printer::visit(Program &program)
 {
   std::string indent(depth_, ' ');
   out_ << indent << "Program" << std::endl;
 
   ++depth_;
-  for (Probe *probe : *program.probes) {
+  for (Include *include : *program.includes)
+    include->accept(*this);
+  for (Probe *probe : *program.probes)
     probe->accept(*this);
-  }
   --depth_;
 }
 

src/ast/printer.h

@@ -24,6 +24,7 @@ public:
   void visit(Predicate &pred) override;
   void visit(AttachPoint &ap) override;
   void visit(Probe &probe) override;
+  void visit(Include &include) override;
   void visit(Program &program) override;
 
   int depth_ = 0;

src/ast/semantic_analyser.cpp

@@ -425,11 +425,16 @@ void SemanticAnalyser::visit(Probe &probe)
   }
 }
 
+void SemanticAnalyser::visit(Include &include)
+{
+}
+
 void SemanticAnalyser::visit(Program &program)
 {
-  for (Probe *probe : *program.probes) {
+  for (Include *include : *program.includes)
+    include->accept(*this);
+  for (Probe *probe : *program.probes)
     probe->accept(*this);
-  }
 }
 
 int SemanticAnalyser::analyse()

src/ast/semantic_analyser.h

@@ -31,6 +31,7 @@ public:
   void visit(Predicate &pred) override;
   void visit(AttachPoint &ap) override;
   void visit(Probe &probe) override;
+  void visit(Include &include) override;
   void visit(Program &program) override;
 
   int analyse();

src/lexer.l

@@ -25,6 +25,7 @@ hspace [ \t]
 vspace [\n\r]
 space  {hspace}|{vspace}
 path   :(\\.|[_\-\./a-zA-Z0-9])*:
+header <(\\.|[_\-\./a-zA-Z0-9])*>
 %x STR
 
 %%
@@ -39,6 +40,7 @@ path   :(\\.|[_\-\./a-zA-Z0-9])*:
 
 {ident}                 { return Parser::make_IDENT(yytext, loc); }
 {path}                  { return Parser::make_PATH(yytext, loc); }
+{header}                { return Parser::make_HEADER(yytext, loc); }
 {map}                   { return Parser::make_MAP(yytext, loc); }
 {var}                   { return Parser::make_VAR(yytext, loc); }
 {int}                   { return Parser::make_INT(strtoul(yytext, NULL, 0), loc); }
@@ -71,6 +73,7 @@ path   :(\\.|[_\-\./a-zA-Z0-9])*:
 "^"                     { return Parser::make_BXOR(loc); }
 "!"                     { return Parser::make_LNOT(loc); }
 "~"                     { return Parser::make_BNOT(loc); }
+"#include"              { return Parser::make_INCLUDE(loc); }
 
 \"                      { BEGIN(STR); string_buffer.clear(); }
 <STR>\"                 { BEGIN(INITIAL); return Parser::make_STRING(string_buffer, loc); }

src/parser.yy

@@ -65,19 +65,23 @@ void yyerror(bpftrace::Driver &driver, const char *s);
   BXOR     "^"
   LNOT     "!"
   BNOT     "~"
+  INCLUDE  "#include"
 ;
 
 %token <std::string> IDENT "identifier"
 %token <std::string> PATH "path"
+%token <std::string> HEADER "header"
 %token <std::string> STRING "string"
 %token <std::string> MAP "map"
 %token <std::string> VAR "variable"
 %token <int> INT "integer"
 
+%type <ast::IncludeList *> includes
+%type <ast::Include *> include
 %type <ast::ProbeList *> probes
-%type <ast::StatementList *> block stmts
 %type <ast::Probe *> probe
 %type <ast::Predicate *> pred
+%type <ast::StatementList *> block stmts
 %type <ast::Statement *> stmt
 %type <ast::Expression *> expr
 %type <ast::Call *> call
@@ -104,7 +108,15 @@ void yyerror(bpftrace::Driver &driver, const char *s);
 
 %%
 
-program : probes { driver.root_ = new ast::Program($1); }
+program : includes probes { driver.root_ = new ast::Program($1, $2); }
+        ;
+
+includes : includes include { $$ = $1; $1->push_back($2); }
+         |                  { $$ = new ast::IncludeList; }
+         ;
+
+include : INCLUDE STRING { $$ = new ast::Include($2); }
+        | INCLUDE HEADER { $$ = new ast::Include($2.substr(1, $2.size()-2)); }
         ;
 
 probes : probes probe { $$ = $1; $1->push_back($2); }

tests/parser.cpp

@@ -300,6 +300,41 @@ TEST(Parser, short_map_name)
       "   int: 1\n");
 }
 
+TEST(Parser, include)
+{
+  test("#include <stdio.h> kprobe:sys_read { @x = 1 }",
+      "Program\n"
+      " #include stdio.h\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @x\n"
+      "   int: 1\n");
+}
+
+TEST(Parser, include_quote)
+{
+  test("#include \"stdio.h\" kprobe:sys_read { @x = 1 }",
+      "Program\n"
+      " #include stdio.h\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @x\n"
+      "   int: 1\n");
+}
+
+TEST(Parser, include_multiple)
+{
+  test("#include <stdio.h> #include \"blah\" #include <foo.h> kprobe:sys_read { @x = 1 }",
+      "Program\n"
+      " #include stdio.h\n"
+      " #include blah\n"
+      " #include foo.h\n"
+      " kprobe:sys_read\n"
+      "  =\n"
+      "   map: @x\n"
+      "   int: 1\n");
+}
+
 } // namespace parser
 } // namespace test
 } // namespace bpftrace