Unroll (#175)

* fixes: https://github.com/iovisor/bpftrace/issues/9

* test fix

* tests

* lifetime

* crash fix

src/ast/ast.cpp

@@ -72,6 +72,10 @@ void If::accept(Visitor &v) {
   v.visit(*this);
 }
 
+void Unroll::accept(Visitor &v) {
+  v.visit(*this);
+}
+
 void Probe::accept(Visitor &v) {
   v.visit(*this);
 }

src/ast/ast.h

@@ -167,6 +167,16 @@ public:
   void accept(Visitor &v) override;
 };
 
+class Unroll : public Statement {
+public:
+  Unroll(long int var, StatementList *stmts) : var(var), stmts(stmts) {}
+
+  StatementList *stmts;
+  long int var = 0;
+
+  void accept(Visitor &v) override;
+};
+
 class Predicate : public Node {
 public:
   explicit Predicate(Expression *expr) : expr(expr) { }
@@ -262,6 +272,7 @@ public:
   virtual void visit(AssignMapStatement &assignment) = 0;
   virtual void visit(AssignVarStatement &assignment) = 0;
   virtual void visit(If &if_block) = 0;
+  virtual void visit(Unroll &unroll) = 0;
   virtual void visit(Predicate &pred) = 0;
   virtual void visit(AttachPoint &ap) = 0;
   virtual void visit(Probe &probe) = 0;

src/ast/codegen_llvm.cpp

@@ -966,6 +966,34 @@ void CodegenLLVM::visit(If &if_block)
   }
 }
 
+void CodegenLLVM::visit(Unroll &unroll)
+{
+  Function *parent = b_.GetInsertBlock()->getParent();
+  BasicBlock *loop = BasicBlock::Create(module_->getContext(), "loop", parent);
+  BasicBlock *done = BasicBlock::Create(module_->getContext(), "done", parent);
+
+  AllocaInst *val = b_.CreateAllocaBPF(SizedType(Type::integer, 8), "loop_count");
+  b_.CreateStore(b_.getInt64(unroll.var), val);
+
+  b_.CreateCondBr(b_.CreateICmpNE(b_.getInt64(unroll.var), b_.getInt64(0), "true_cond"), loop, done);
+
+  b_.SetInsertPoint(loop);
+
+  for (Statement *stmt : *unroll.stmts)
+  {
+    stmt->accept(*this);
+  }
+
+  Value *var = b_.CreateLoad(val);
+  Value *newValue = b_.CreateSub(var, b_.getInt64(1), "subtmp");
+  b_.CreateStore(newValue, val);
+
+  b_.CreateCondBr(b_.CreateICmpNE(newValue, b_.getInt64(0), "loop_cond"), loop, done);
+
+  b_.SetInsertPoint(done);
+  b_.CreateLifetimeEnd(val);
+}
+
 void CodegenLLVM::visit(Predicate &pred)
 {
   Function *parent = b_.GetInsertBlock()->getParent();

src/ast/codegen_llvm.h

@@ -42,6 +42,7 @@ public:
   void visit(AssignMapStatement &assignment) override;
   void visit(AssignVarStatement &assignment) override;
   void visit(If &if_block) override;
+  void visit(Unroll &unroll) override;
   void visit(Predicate &pred) override;
   void visit(AttachPoint &ap) override;
   void visit(Probe &probe) override;

src/ast/printer.cpp

@@ -175,6 +175,18 @@ void Printer::visit(If &if_block)
   depth_ -= 2;
 }
 
+void Printer::visit(Unroll &unroll)
+{
+  std::string indent(depth_, ' ');
+  out_ << indent << "unroll " << unroll.var << std::endl;
+  ++depth_;
+
+  for (Statement *stmt : *unroll.stmts) {
+    stmt->accept(*this);
+  }
+  --depth_;
+}
+
 void Printer::visit(Predicate &pred)
 {
   std::string indent(depth_, ' ');

src/ast/printer.h

@@ -25,6 +25,7 @@ public:
   void visit(AssignMapStatement &assignment) override;
   void visit(AssignVarStatement &assignment) override;
   void visit(If &if_block) override;
+  void visit(Unroll &unroll) override;
   void visit(Predicate &pred) override;
   void visit(AttachPoint &ap) override;
   void visit(Probe &probe) override;

src/ast/semantic_analyser.cpp

@@ -496,6 +496,23 @@ void SemanticAnalyser::visit(If &if_block)
   }
 }
 
+void SemanticAnalyser::visit(Unroll &unroll)
+{
+  if (unroll.var > 20)
+  {
+    err_ << "unroll maximum value is 20.\n" << std::endl;
+  }
+  else if (unroll.var == 0)
+  {
+    err_ << "unroll minimum value is 1.\n" << std::endl;
+  }
+
+  for (Statement *stmt : *unroll.stmts)
+  {
+    stmt->accept(*this);
+  }
+}
+
 void SemanticAnalyser::visit(FieldAccess &acc)
 {
   acc.expr->accept(*this);

src/ast/semantic_analyser.h

@@ -33,6 +33,7 @@ public:
   void visit(AssignMapStatement &assignment) override;
   void visit(AssignVarStatement &assignment) override;
   void visit(If &if_block) override;
+  void visit(Unroll &unroll) override;
   void visit(Predicate &pred) override;
   void visit(AttachPoint &ap) override;
   void visit(Probe &probe) override;

src/lexer.l

@@ -91,6 +91,7 @@ pid|tid|cgroup|uid|gid|nsecs|cpu|comm|stack|ustack|arg[0-9]|retval|func|name|cur
 "if"                    { return Parser::make_IF(yytext, loc); }
 "else"                  { return Parser::make_ELSE(yytext, loc); }
 "?"                     { return Parser::make_QUES(loc); }
+"unroll"                { return Parser::make_UNROLL(yytext, loc); }
 
 \"                      { BEGIN(STR); buffer.clear(); }
 <STR>\"                 { BEGIN(INITIAL); return Parser::make_STRING(buffer, loc); }

src/parser.yy

@@ -80,9 +80,10 @@ void yyerror(bpftrace::Driver &driver, const char *s);
 %token <std::string> STRING "string"
 %token <std::string> MAP "map"
 %token <std::string> VAR "variable"
+%token <long> INT "integer"
 %nonassoc <std::string> IF "if"
 %nonassoc <std::string> ELSE "else"
-%token <long> INT "integer"
+%nonassoc <std::string> UNROLL "unroll"
 
 %type <std::string> c_definitions
 %type <ast::ProbeList *> probes
@@ -172,6 +173,7 @@ stmt : expr         { $$ = new ast::ExprStatement($1); }
      | var "=" expr { $$ = new ast::AssignVarStatement($1, $3); }
      | IF "(" expr ")" block  { $$ = new ast::If($3, $5); }
      | IF "(" expr ")" block ELSE block { $$ = new ast::If($3, $5, $7); }
+     | UNROLL "(" INT ")" block { $$ = new ast::Unroll($3, $5); }
      ;
 
 expr : INT             { $$ = new ast::Integer($1); }

tests/codegen.cpp

@@ -3144,6 +3144,76 @@ attributes #1 = { argmemonly nounwind }
 )EXPECTED");
 }
 
+TEST(codegen, unroll)
+{
+  test("kprobe:f { $i = 0; unroll(5) { printf(\"i: %d\\n\", $i); $i = $i + 1; } }",
+
+R"EXPECTED(%printf_t = type { i64, i64 }
+
+; Function Attrs: nounwind
+declare i64 @llvm.bpf.pseudo(i64, i64) #0
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+
+define i64 @"kprobe:f"(i8*) local_unnamed_addr section "s_kprobe:f_1" {
+entry:
+  %printf_args = alloca %printf_t, align 8
+  %1 = bitcast %printf_t* %printf_args to i8*
+  %2 = getelementptr inbounds %printf_t, %printf_t* %printf_args, i64 0, i32 1
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
+  %3 = bitcast %printf_t* %printf_args to i8*
+  call void @llvm.memset.p0i8.i64(i8* nonnull %3, i8 0, i64 16, i32 8, i1 false)
+  %pseudo = call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %get_cpu_id = call i64 inttoptr (i64 8 to i64 ()*)()
+  %perf_event_output = call i64 inttoptr (i64 25 to i64 (i8*, i8*, i64, i8*, i64)*)(i8* %0, i64 %pseudo, i64 %get_cpu_id, %printf_t* nonnull %printf_args, i64 16)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
+  %4 = getelementptr inbounds %printf_t, %printf_t* %printf_args, i64 0, i32 0
+  store i64 0, i64* %4, align 8
+  store i64 1, i64* %2, align 8
+  %pseudo.1 = call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %get_cpu_id.1 = call i64 inttoptr (i64 8 to i64 ()*)()
+  %perf_event_output.1 = call i64 inttoptr (i64 25 to i64 (i8*, i8*, i64, i8*, i64)*)(i8* %0, i64 %pseudo.1, i64 %get_cpu_id.1, %printf_t* nonnull %printf_args, i64 16)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
+  %5 = getelementptr inbounds %printf_t, %printf_t* %printf_args, i64 0, i32 0
+  store i64 0, i64* %5, align 8
+  store i64 2, i64* %2, align 8
+  %pseudo.2 = call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %get_cpu_id.2 = call i64 inttoptr (i64 8 to i64 ()*)()
+  %perf_event_output.2 = call i64 inttoptr (i64 25 to i64 (i8*, i8*, i64, i8*, i64)*)(i8* %0, i64 %pseudo.2, i64 %get_cpu_id.2, %printf_t* nonnull %printf_args, i64 16)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
+  %6 = getelementptr inbounds %printf_t, %printf_t* %printf_args, i64 0, i32 0
+  store i64 0, i64* %6, align 8
+  store i64 3, i64* %2, align 8
+  %pseudo.3 = call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %get_cpu_id.3 = call i64 inttoptr (i64 8 to i64 ()*)()
+  %perf_event_output.3 = call i64 inttoptr (i64 25 to i64 (i8*, i8*, i64, i8*, i64)*)(i8* %0, i64 %pseudo.3, i64 %get_cpu_id.3, %printf_t* nonnull %printf_args, i64 16)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
+  %7 = getelementptr inbounds %printf_t, %printf_t* %printf_args, i64 0, i32 0
+  store i64 0, i64* %7, align 8
+  store i64 4, i64* %2, align 8
+  %pseudo.4 = call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %get_cpu_id.4 = call i64 inttoptr (i64 8 to i64 ()*)()
+  %perf_event_output.4 = call i64 inttoptr (i64 25 to i64 (i8*, i8*, i64, i8*, i64)*)(i8* %0, i64 %pseudo.4, i64 %get_cpu_id.4, %printf_t* nonnull %printf_args, i64 16)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
+  ret i64 0
+}
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.memset.p0i8.i64(i8* nocapture writeonly, i8, i64, i32, i1) #1
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { argmemonly nounwind }
+)EXPECTED");
+}
+
 TEST(codegen, struct_long)
 {
   auto expected = R"EXPECTED(; Function Attrs: nounwind

tests/parser.cpp

@@ -350,6 +350,25 @@ TEST(Parser, if_else)
        "   variable: $s\n");
 }
 
+TEST(Parser, unroll)
+{
+  test("kprobe:sys_open { $i = 0; unroll(5) { printf(\"i: %d\\n\", $i); $i = $i + 1; } }",
+       "Program\n"
+       " kprobe:sys_open\n"
+       "  =\n"
+       "   variable: $i\n"
+       "   int: 0\n"
+       "  unroll 5\n"
+       "   call: printf\n"
+       "    string: i: %d\\n\n"
+       "    variable: $i\n"
+       "   =\n"
+       "    variable: $i\n"
+       "    +\n"
+       "     variable: $i\n"
+       "     int: 1\n");
+}
+
 TEST(Parser, ternary_str)
 {
   test("kprobe:sys_open { @x = pid < 10000 ? \"lo\" : \"high\" }",

tests/semantic_analyser.cpp

@@ -322,6 +322,13 @@ TEST(semantic_analyser, variable_type)
   EXPECT_EQ(st, assignment->var->type);
 }
 
+TEST(semantic_analyser, unroll)
+{
+  test("kprobe:f { $i = 0; unroll(5) { printf(\"i: %d\\n\", $i); $i = $i + 1; } }", 0);
+  test("kprobe:f { $i = 0; unroll(21) { printf(\"i: %d\\n\", $i); $i = $i + 1; } }", 1);
+  test("kprobe:f { $i = 0; unroll(0) { printf(\"i: %d\\n\", $i); $i = $i + 1; } }", 1);
+}
+
 TEST(semantic_analyser, map_integer_sizes)
 {
   Driver driver;