Make function calls usable as expressions

src/ast.cpp

@@ -40,10 +40,6 @@ void AssignMapStatement::accept(Visitor &v) {
   v.visit(*this);
 }
 
-void AssignMapCallStatement::accept(Visitor &v) {
-  v.visit(*this);
-}
-
 void Predicate::accept(Visitor &v) {
   v.visit(*this);
 }

src/ast.h

@@ -18,7 +18,10 @@ public:
   SizedType type;
 };
 
+class Map;
 class Expression : public Node {
+public:
+  Map *map = nullptr; // Only set when this expression is assigned to a map
 };
 using ExpressionList = std::vector<Expression *>;
 
@@ -98,22 +101,15 @@ public:
 
 class AssignMapStatement : public Statement {
 public:
-  AssignMapStatement(Map *map, Expression *expr) : map(map), expr(expr) { }
+  AssignMapStatement(Map *map, Expression *expr) : map(map), expr(expr) {
+    expr->map = map;
+  }
   Map *map;
   Expression *expr;
 
   void accept(Visitor &v) override;
 };
 
-class AssignMapCallStatement : public Statement {
-public:
-  AssignMapCallStatement(Map *map, Call *call) : map(map), call(call) { }
-  Map *map;
-  Call *call;
-
-  void accept(Visitor &v) override;
-};
-
 class Predicate : public Node {
 public:
   explicit Predicate(Expression *expr) : expr(expr) { }
@@ -160,7 +156,6 @@ public:
   virtual void visit(Unop &unop) = 0;
   virtual void visit(ExprStatement &expr) = 0;
   virtual void visit(AssignMapStatement &assignment) = 0;
-  virtual void visit(AssignMapCallStatement &assignment) = 0;
   virtual void visit(Predicate &pred) = 0;
   virtual void visit(Probe &probe) = 0;
   virtual void visit(Program &program) = 0;

src/codegen_llvm.cpp

@@ -93,7 +93,41 @@ void CodegenLLVM::visit(Builtin &builtin)
 
 void CodegenLLVM::visit(Call &call)
 {
+  if (call.func == "count")
+  {
+    Map &map = *call.map;
+    AllocaInst *key = getMapKey(map);
+    Value *oldval = b_.CreateMapLookupElem(map, key);
+    AllocaInst *newval = b_.CreateAllocaBPF(map.type, map.ident + "_val");
+    b_.CreateStore(b_.CreateAdd(oldval, b_.getInt64(1)), newval);
+    b_.CreateMapUpdateElem(map, key, newval);
+    expr_ = nullptr;
+  }
+  else if (call.func == "quantize")
+  {
+    Map &map = *call.map;
+    call.vargs->front()->accept(*this);
+    Function *log2_func = module_->getFunction("log2");
+    Value *log2 = b_.CreateCall(log2_func, expr_, "log2");
+    AllocaInst *key = getQuantizeMapKey(map, log2);
+
+    Value *oldval = b_.CreateMapLookupElem(map, key);
+    AllocaInst *newval = b_.CreateAllocaBPF(map.type, map.ident + "_val");
+    b_.CreateStore(b_.CreateAdd(oldval, b_.getInt64(1)), newval);
+    b_.CreateMapUpdateElem(map, key, newval);
+    expr_ = nullptr;
+  }
+  else if (call.func == "delete")
+  {
+    Map &map = *call.map;
+    AllocaInst *key = getMapKey(map);
+    b_.CreateMapDeleteElem(map, key);
+    expr_ = nullptr;
+  }
+  else
+  {
     abort();
+  }
 }
 
 void CodegenLLVM::visit(Map &map)
@@ -179,6 +213,9 @@ void CodegenLLVM::visit(AssignMapStatement &assignment)
 
   assignment.expr->accept(*this);
 
+  if (!expr_) // Some functions do the assignments themselves
+    return;
+
   if (assignment.expr->type.type == Type::string)
   {
     Value *val = expr_;
@@ -194,42 +231,6 @@ void CodegenLLVM::visit(AssignMapStatement &assignment)
   }
 }
 
-void CodegenLLVM::visit(AssignMapCallStatement &assignment)
-{
-  Map &map = *assignment.map;
-  Call &call = *assignment.call;
-
-  if (call.func == "count")
-  {
-    AllocaInst *key = getMapKey(map);
-    Value *oldval = b_.CreateMapLookupElem(map, key);
-    AllocaInst *newval = b_.CreateAllocaBPF(map.type, map.ident + "_val");
-    b_.CreateStore(b_.CreateAdd(oldval, b_.getInt64(1)), newval);
-    b_.CreateMapUpdateElem(map, key, newval);
-  }
-  else if (call.func == "quantize")
-  {
-    call.vargs->front()->accept(*this);
-    Function *log2_func = module_->getFunction("log2");
-    Value *log2 = b_.CreateCall(log2_func, expr_, "log2");
-    AllocaInst *key = getQuantizeMapKey(map, log2);
-
-    Value *oldval = b_.CreateMapLookupElem(map, key);
-    AllocaInst *newval = b_.CreateAllocaBPF(map.type, map.ident + "_val");
-    b_.CreateStore(b_.CreateAdd(oldval, b_.getInt64(1)), newval);
-    b_.CreateMapUpdateElem(map, key, newval);
-  }
-  else if (call.func == "delete")
-  {
-    AllocaInst *key = getMapKey(map);
-    b_.CreateMapDeleteElem(map, key);
-  }
-  else
-  {
-    abort();
-  }
-}
-
 void CodegenLLVM::visit(Predicate &pred)
 {
   Function *parent = b_.GetInsertBlock()->getParent();

src/codegen_llvm.h

@@ -32,7 +32,6 @@ public:
   void visit(Unop &unop) override;
   void visit(ExprStatement &expr) override;
   void visit(AssignMapStatement &assignment) override;
-  void visit(AssignMapCallStatement &assignment) override;
   void visit(Predicate &pred) override;
   void visit(Probe &probe) override;
   void visit(Program &program) override;

src/parser.yy

@@ -124,13 +124,13 @@ stmts : stmts ";" stmt { $$ = $1; $1->push_back($3); }
 
 stmt : expr         { $$ = new ast::ExprStatement($1); }
      | map "=" expr { $$ = new ast::AssignMapStatement($1, $3); }
-     | map "=" call { $$ = new ast::AssignMapCallStatement($1, $3); }
      ;
 
 expr : INT             { $$ = new ast::Integer($1); }
      | STRING          { $$ = new ast::String($1.substr(1, $1.size()-2)); }
      | IDENT           { $$ = new ast::Builtin($1); }
      | map             { $$ = $1; }
+     | call            { $$ = $1; }
      | "(" expr ")"    { $$ = $2; }
      | expr EQ expr    { $$ = new ast::Binop($1, token::EQ, $3); }
      | expr NE expr    { $$ = new ast::Binop($1, token::NE, $3); }

src/printer.cpp

@@ -87,17 +87,6 @@ void Printer::visit(AssignMapStatement &assignment)
   --depth_;
 }
 
-void Printer::visit(AssignMapCallStatement &assignment)
-{
-  std::string indent(depth_, ' ');
-  out_ << indent << "=" << std::endl;
-
-  ++depth_;
-  assignment.map->accept(*this);
-  assignment.call->accept(*this);
-  --depth_;
-}
-
 void Printer::visit(Predicate &pred)
 {
   std::string indent(depth_, ' ');

src/printer.h

@@ -19,7 +19,6 @@ public:
   void visit(Unop &unop) override;
   void visit(ExprStatement &expr) override;
   void visit(AssignMapStatement &assignment) override;
-  void visit(AssignMapCallStatement &assignment) override;
   void visit(Predicate &pred) override;
   void visit(Probe &probe) override;
   void visit(Program &program) override;

src/semantic_analyser.cpp

@@ -66,6 +66,9 @@ void SemanticAnalyser::visit(Call &call)
   }
 
   if (call.func == "quantize") {
+    if (!call.map) {
+      err_ << "quantize() should be assigned to a map" << std::endl;
+    }
     if (nargs != 1) {
       err_ << "quantize() should take 1 argument (";
       err_ << nargs << " provided)" << std::endl;
@@ -77,6 +80,9 @@ void SemanticAnalyser::visit(Call &call)
     call.type = SizedType(Type::quantize, 8);
   }
   else if (call.func == "count") {
+    if (!call.map) {
+      err_ << "count() should be assigned to a map" << std::endl;
+    }
     if (nargs != 0) {
       err_ << "count() should take 0 arguments (";
       err_ << nargs << " provided)" << std::endl;
@@ -84,15 +90,18 @@ void SemanticAnalyser::visit(Call &call)
     call.type = SizedType(Type::count, 8);
   }
   else if (call.func == "delete") {
+    if (!call.map) {
+      err_ << "delete() should be assigned to a map" << std::endl;
+    }
     if (nargs != 0) {
       err_ << "delete() should take 0 arguments (";
       err_ << nargs << " provided)" << std::endl;
     }
-    // Don't assign a type
+    call.type = SizedType(Type::del, 0);
   }
   else {
-    call.type = SizedType(Type::none, 0);
     err_ << "Unknown function: '" << call.func << "'" << std::endl;
+    call.type = SizedType(Type::none, 0);
   }
 }
 
@@ -188,7 +197,8 @@ void SemanticAnalyser::visit(AssignMapStatement &assignment)
         search->second = assignment.expr->type;
       }
     }
-    else if (search->second.type != assignment.expr->type.type) {
+    else if (search->second.type != assignment.expr->type.type &&
+             assignment.expr->type.type != Type::del) {
       err_ << "Type mismatch for " << map_ident << ": ";
       err_ << "trying to assign value of type '" << assignment.expr->type;
       err_ << "'\n\twhen map already contains a value of type '";
@@ -201,36 +211,6 @@ void SemanticAnalyser::visit(AssignMapStatement &assignment)
   }
 }
 
-void SemanticAnalyser::visit(AssignMapCallStatement &assignment)
-{
-  assignment.map->accept(*this);
-  assignment.call->accept(*this);
-
-  std::string map_ident = assignment.map->ident;
-  auto search = map_val_.find(map_ident);
-  if (search != map_val_.end()) {
-    if (search->second.type == Type::none) {
-      if (is_final_pass()) {
-        err_ << "Undefined map: " << map_ident << std::endl;
-      }
-      else {
-        search->second = assignment.call->type;
-      }
-    }
-    else if (search->second.type != assignment.call->type.type &&
-             assignment.call->func != "delete") {
-      err_ << "Type mismatch for " << map_ident << ": ";
-      err_ << "trying to assign result of '" << assignment.call->func;
-      err_ << "()'\n\twhen map already contains a value of type '";
-      err_ << search->second << "'\n" << std::endl;
-    }
-  }
-  else {
-    // This map hasn't been seen before
-    map_val_.insert({map_ident, assignment.call->type});
-  }
-}
-
 void SemanticAnalyser::visit(Predicate &pred)
 {
   pred.expr->accept(*this);

src/semantic_analyser.h

@@ -26,7 +26,6 @@ public:
   void visit(Unop &unop) override;
   void visit(ExprStatement &expr) override;
   void visit(AssignMapStatement &assignment) override;
-  void visit(AssignMapCallStatement &assignment) override;
   void visit(Predicate &pred) override;
   void visit(Probe &probe) override;
   void visit(Program &program) override;

src/types.h

@@ -14,6 +14,7 @@ const int STRING_SIZE = 32;
 enum class Type
 {
   none,
+  del,
   integer,
   quantize,
   count,

tests/parser.cpp

@@ -144,7 +144,7 @@ TEST(Parser, expressions)
 
 TEST(Parser, call)
 {
-  test("kprobe:sys_open { @x = foo(); @y = bar(1,2,3); }",
+  test("kprobe:sys_open { @x = foo(); @y = bar(1,2,3); myfunc(@x); }",
       "Program\n"
       " kprobe:sys_open\n"
       "  =\n"
@@ -155,7 +155,9 @@ TEST(Parser, call)
       "   call: bar\n"
       "    int: 1\n"
       "    int: 2\n"
-      "    int: 3\n");
+      "    int: 3\n"
+      "  call: myfunc\n"
+      "   map: @x\n");
 }
 
 TEST(Parser, multiple_probes)

tests/semantic_analyser.cpp

@@ -62,4 +62,25 @@ TEST(semantic_analyser, mismatched_call_types)
   test("kprobe:f { @x = 1; @x = delete(); }", 0);
 }
 
+TEST(semantic_analyser, call_quantize)
+{
+  test("kprobe:f { @x = quantize(1); }", 0);
+  test("kprobe:f { @x = quantize(); }", 1);
+  test("kprobe:f { quantize(); }", 1);
+}
+
+TEST(semantic_analyser, call_count)
+{
+  test("kprobe:f { @x = count(); }", 0);
+  test("kprobe:f { @x = count(1); }", 1);
+  test("kprobe:f { count(); }", 1);
+}
+
+TEST(semantic_analyser, call_delete)
+{
+  test("kprobe:f { @x = delete(); }", 0);
+  test("kprobe:f { @x = delete(1); }", 1);
+  test("kprobe:f { delete(); }", 1);
+}
+
 } // namespace bpftrace