Implement while loop support

With the loop supported added in the 5.0 kernel we're now able to
implement C-style while and for loops. Allowing us to write code like:

```
i:s:1 {
	$i = 0;
	while ($i < 100) {
		if ( ($i/10) * 10 == $i ) {
			printf("\n");
		}
		if ( ($i/6)*6 == $i ){
			printf("XX ");
			$i+=1;
			continue;
		}
		printf("%2d ", $i);
		$i+=1;
		if ($i >= 50) {
			break;
		}
	}
}
```

To generate:

```
Attaching 1 probe...

XX  1  2  3  4  5 XX  7  8  9
10 11 XX 13 14 15 16 17 XX 19
20 21 22 23 XX 25 26 27 28 29
XX 31 32 33 34 35 XX 37 38 39
40 41 XX 43 44 45 46 47 XX 49
```

Using only a few instructions:

```
203: perf_event  name 1  tag 4881cacfe5abad41  gpl
	loaded_at 2020-03-26T18:17:17+0000  uid 0
	xlated 448B  jited 340B  memlock 4096B  map_ids 85
```

Note that LLVM tries to unroll quite aggressively causing a loop the
following be fully unrolled:

```
i:s:1 { $i = 0; while ($i < 30) { @=$i; $i++ } }
```

docs/reference_guide.md

@@ -35,6 +35,8 @@ discussion to other files in /docs, the /tools/\*\_examples.txt files, or blog p
     - [9. `++ and --`: increment operators](#9--and----increment-operators)
     - [10. `[]`: Array access](#10--array-access)
     - [11. Integer casts](#11-integer-casts)
+    - [12. Looping constructs](#12-looping-constructs)
+    - [13. `return`: Terminate Early](#13-return-terminate-early)
 - [Probes](#probes)
     - [1. `kprobe`/`kretprobe`: Dynamic Tracing, Kernel-Level](#1-kprobekretprobe-dynamic-tracing-kernel-level)
     - [2. `kprobe`/`kretprobe`: Dynamic Tracing, Kernel-Level Arguments](#2-kprobekretprobe-dynamic-tracing-kernel-level-arguments)
@@ -748,6 +750,25 @@ Attaching 1 probe...
 ^C
 ```
 
+## 12. Looping Constructs
+
+**Experimental**
+
+Kernel: 5.3
+
+bpftrace supports C style while loops:
+
+```
+# bpftrace -e 'i:ms:100 { $i = 0; while ($i <= 100) { printf("%d ", $i); $i++} exit(); }'
+```
+
+Loops can be short circuited by using the `continue` and `break` keywords.
+
+## 13. `return`: Terminate Early
+
+The `return` keyword is used to exit the current probe. This differs from
+`exit()` in that it doesn't exit bpftrace.
+
 # Probes
 
 - `kprobe` - kernel function start
@@ -2526,7 +2547,7 @@ they don't corrupt the terminal display. The resulting string can be provided as
 printf() using the `%r` format specifier:
 
 ```
-# bpftrace -e 'tracepoint:syscalls:sys_enter_sendto 
+# bpftrace -e 'tracepoint:syscalls:sys_enter_sendto
     { printf("Datagram bytes: %r\n", buf(args->buff, args->len)); }' -c 'ping 8.8.8.8 -c1'
 Attaching 1 probe...
 PING 8.8.8.8 (8.8.8.8) 56(84) bytes of data.

src/ast/codegen_llvm.cpp

@@ -1516,41 +1516,56 @@ void CodegenLLVM::visit(AssignVarStatement &assignment)
 void CodegenLLVM::visit(If &if_block)
 {
   Function *parent = b_.GetInsertBlock()->getParent();
-  BasicBlock *if_true = BasicBlock::Create(module_->getContext(), "if_stmt", parent);
-  BasicBlock *if_false = BasicBlock::Create(module_->getContext(), "else_stmt", parent);
+  BasicBlock *if_true = BasicBlock::Create(module_->getContext(),
+                                           "if_body",
+                                           parent);
+  BasicBlock *if_end = BasicBlock::Create(module_->getContext(),
+                                          "if_end",
+                                          parent);
+  BasicBlock *if_else = nullptr;
 
   if_block.cond->accept(*this);
-  Value *cond = expr_;
+  Value *zero_value = Constant::getNullValue(expr_->getType());
+  Value *cond = b_.CreateICmpNE(expr_, zero_value, "true_cond");
 
-  Value *zero_value = Constant::getNullValue(cond->getType());
-  b_.CreateCondBr(b_.CreateICmpNE(cond, zero_value, "true_cond"),
-                  if_true,
-                  if_false);
+  // 3 possible flows:
+  //
+  // if condition is true
+  //   parent -> if_body -> if_end
+  //
+  // if condition is false, no else
+  //   parent -> if_end
+  //
+  // if condition is false, with else
+  //   parent -> if_else -> if_end
+  //
+  if (if_block.else_stmts)
+  {
+    // LLVM doesn't accept empty basic block, only create when needed
+    if_else = BasicBlock::Create(module_->getContext(), "else_body", parent);
+    b_.CreateCondBr(cond, if_true, if_else);
+  }
+  else
+  {
+    b_.CreateCondBr(cond, if_true, if_end);
+  }
 
   b_.SetInsertPoint(if_true);
   for (Statement *stmt : *if_block.stmts)
-  {
     stmt->accept(*this);
-  }
+
+  b_.CreateBr(if_end);
+
+  b_.SetInsertPoint(if_end);
 
   if (if_block.else_stmts)
   {
-    BasicBlock *done = BasicBlock::Create(module_->getContext(), "done", parent);
-    b_.CreateBr(done);
-
-    b_.SetInsertPoint(if_false);
+    b_.SetInsertPoint(if_else);
     for (Statement *stmt : *if_block.else_stmts)
-    {
       stmt->accept(*this);
-    }
-    b_.CreateBr(done);
 
-    b_.SetInsertPoint(done);
-  }
-  else
-  {
-      b_.CreateBr(if_false);
-      b_.SetInsertPoint(if_false);
+    b_.CreateBr(if_end);
+    b_.SetInsertPoint(if_end);
   }
 }
 
@@ -1566,12 +1581,79 @@ void CodegenLLVM::visit(Unroll &unroll)
 
 void CodegenLLVM::visit(Jump &jump)
 {
-  return;
+  switch (jump.ident)
+  {
+    case bpftrace::Parser::token::RETURN:
+      // return can be used outside of loops
+      b_.CreateRet(ConstantInt::get(module_->getContext(), APInt(64, 0)));
+      break;
+    case bpftrace::Parser::token::BREAK:
+      b_.CreateBr(std::get<1>(loops_.back()));
+      break;
+    case bpftrace::Parser::token::CONTINUE:
+      b_.CreateBr(std::get<0>(loops_.back()));
+      break;
+    default:
+      throw std::runtime_error("Unknown jump: " + opstr(jump));
+  }
+
+  // LLVM doesn't like having instructions after an unconditional branch (segv)
+  // This can be avoided by putting all instructions in a unreachable basicblock
+  // which will be optimize out.
+  //
+  // e.g. in the case of `while (..) { $i++; break; $i++ }` the ir will be:
+  //
+  // while_body:
+  //   ...
+  //   br label %while_end
+  //
+  // while_end:
+  //   ...
+  //
+  // unreach:
+  //   $i++
+  //   br label %while_cond
+  //
+
+  Function *parent = b_.GetInsertBlock()->getParent();
+  BasicBlock *unreach = BasicBlock::Create(module_->getContext(),
+                                           "unreach",
+                                           parent);
+  b_.SetInsertPoint(unreach);
 }
 
 void CodegenLLVM::visit(While &while_block)
 {
-  return;
+  Function *parent = b_.GetInsertBlock()->getParent();
+  BasicBlock *while_cond = BasicBlock::Create(module_->getContext(),
+                                              "while_cond",
+                                              parent);
+  BasicBlock *while_body = BasicBlock::Create(module_->getContext(),
+                                              "while_body",
+                                              parent);
+  BasicBlock *while_end = BasicBlock::Create(module_->getContext(),
+                                             "while_end",
+                                             parent);
+
+  loops_.push_back(std::make_tuple(while_cond, while_end));
+
+  b_.CreateBr(while_cond);
+
+  b_.SetInsertPoint(while_cond);
+  while_block.cond->accept(*this);
+  Value *zero_value = Constant::getNullValue(expr_->getType());
+  auto *cond = b_.CreateICmpNE(expr_, zero_value, "true_cond");
+  b_.CreateCondBr(cond, while_body, while_end);
+
+  b_.SetInsertPoint(while_body);
+  for (Statement *stmt : *while_block.stmts)
+  {
+    stmt->accept(*this);
+  }
+  b_.CreateBr(while_cond);
+
+  b_.SetInsertPoint(while_end);
+  loops_.pop_back();
 }
 
 void CodegenLLVM::visit(Predicate &pred)

src/ast/codegen_llvm.h

@@ -98,6 +98,8 @@ class CodegenLLVM : public Visitor {
   {
     return layout_.getTypeAllocSize(s);
   }
+
+  std::vector<std::tuple<BasicBlock *, BasicBlock *>> loops_;
 };
 
 } // namespace ast

src/ast/semantic_analyser.cpp

@@ -818,6 +818,13 @@ void SemanticAnalyser::visit(Call &call)
   }
   else if (call.func == "print") {
     check_assignment(call, false, false, false);
+    if (in_loop() && is_final_pass())
+    {
+      warning("Due to it's asynchronous nature using 'print()' in a loop can "
+              "lead to unexpected behavior. The map will likely be updated "
+              "before the runtime can 'print' it.",
+              call.loc);
+    }
     if (check_varargs(call, 1, 3)) {
       auto &arg = *call.vargs->at(0);
       if (!arg.is_map)
@@ -1393,44 +1400,54 @@ void SemanticAnalyser::visit(If &if_block)
       ERR("Invalid condition in if(): " << cond, if_block.loc);
   }
 
-  for (Statement *stmt : *if_block.stmts) {
-    stmt->accept(*this);
-  }
+  accept_statements(if_block.stmts);
 
-  if (if_block.else_stmts) {
-    for (Statement *stmt : *if_block.else_stmts) {
-      stmt->accept(*this);
-    }
-  }
+  if (if_block.else_stmts)
+    accept_statements(if_block.else_stmts);
 }
 
 void SemanticAnalyser::visit(Unroll &unroll)
 {
   if (unroll.var > 20)
-  {
     error("unroll maximum value is 20", location(0));
-  }
   else if (unroll.var < 1)
-  {
     error("unroll minimum value is 1", location(0));
-  }
 
-  for (int i=0; i < unroll.var; i++) {
-    for (Statement *stmt : *unroll.stmts)
-    {
-      stmt->accept(*this);
-    }
-  }
+  for (int i = 0; i < unroll.var; i++)
+    accept_statements(unroll.stmts);
 }
 
 void SemanticAnalyser::visit(Jump &jump)
 {
-  error(opstr(jump) + " has not yet been implemented", jump.loc);
+  switch (jump.ident)
+  {
+    case bpftrace::Parser::token::RETURN:
+      // return can be used outside of loops
+      break;
+    case bpftrace::Parser::token::BREAK:
+    case bpftrace::Parser::token::CONTINUE:
+      if (!in_loop())
+        error(opstr(jump) + " used outside of a loop", jump.loc);
+      break;
+    default:
+      error("Unknown jump: '" + opstr(jump) + "'", jump.loc);
+  }
 }
 
 void SemanticAnalyser::visit(While &while_block)
 {
-  error("While has not yet been implemented", while_block.loc);
+  if (is_final_pass() && !feature_.has_loop())
+  {
+    warning("Kernel does not support bounded loops. Depending"
+            " on LLVMs loop unroll to generate loadable code.",
+            while_block.loc);
+  }
+
+  while_block.cond->accept(*this);
+
+  loop_depth_++;
+  accept_statements(while_block.stmts);
+  loop_depth_--;
 }
 
 void SemanticAnalyser::visit(FieldAccess &acc)
@@ -2384,5 +2401,24 @@ void SemanticAnalyser::assign_map_type(const Map &map, const SizedType &type)
   }
 }
 
+void SemanticAnalyser::accept_statements(StatementList *stmts)
+{
+  for (size_t i = 0; i < stmts->size(); i++)
+  {
+    auto stmt = stmts->at(i);
+    stmt->accept(*this);
+
+    if (is_final_pass())
+    {
+      auto *jump = dynamic_cast<Jump *>(stmt);
+      if (jump && i < (stmts->size() - 1))
+      {
+        warning("All code after a '" + opstr(*jump) + "' is unreachable.",
+                jump->loc);
+      }
+    }
+  }
+}
+
 } // namespace ast
 } // namespace bpftrace

src/ast/semantic_analyser.h

@@ -92,6 +92,12 @@ class SemanticAnalyser : public Visitor {
   void builtin_args_tracepoint(AttachPoint *attach_point, Builtin &builtin);
   ProbeType single_provider_type(void);
 
+  bool in_loop(void)
+  {
+    return loop_depth_ > 0;
+  };
+  void accept_statements(StatementList *stmts);
+
   Probe *probe_;
   std::string func_;
   std::map<std::string, SizedType> variable_val_;
@@ -100,6 +106,7 @@ class SemanticAnalyser : public Visitor {
   std::map<std::string, ExpressionList> map_args_;
   std::map<std::string, SizedType> ap_args_;
   std::unordered_set<StackType> needs_stackid_maps_;
+  uint32_t loop_depth_ = 0;
   bool needs_join_map_ = false;
   bool needs_elapsed_map_ = false;
   bool has_begin_probe_ = false;

tests/codegen/basic_while_loop.cpp

@@ -0,0 +1,14 @@
+#include "common.h"
+
+namespace bpftrace {
+namespace test {
+namespace codegen {
+
+TEST(codegen, basic_while_loop)
+{
+  test("i:s:1 { $a = 1; while ($a <= 150) { @=$a++; }}", NAME);
+}
+
+} // namespace codegen
+} // namespace test
+} // namespace bpftrace