Implement quantize()

src/bpftrace.cpp

@@ -56,6 +56,7 @@ int BPFtrace::print_maps()
     Map &map = *mapmap.second.get();
 
     int key_elems = map.args_.size();
+    if (map.type_ == Type::quantize) key_elems += 1;
     if (key_elems == 0) key_elems = 1;
     auto old_key = std::vector<uint64_t>(key_elems);
     auto key = std::vector<uint64_t>(key_elems);

src/codegen_llvm.cpp

@@ -137,13 +137,15 @@ void CodegenLLVM::visit(AssignMapCallStatement &assignment)
   }
   else if (call.func == "quantize")
   {
-    // TODO
     call.vargs->front()->accept(*this);
     Function *log2_func = module_->getFunction("log2");
-    Value *log2 = b_.CreateAllocaBPF();
-    b_.CreateStore(b_.CreateCall(log2_func, expr_), log2);
-    AllocaInst *key = getMapKey(map);
-    b_.CreateMapUpdateElem(map, key, log2);
+    Value *log2 = b_.CreateCall(log2_func, expr_);
+    AllocaInst *key = getQuantizeMapKey(map, log2);
+
+    Value *oldval = b_.CreateMapLookupElem(map, key);
+    AllocaInst *newval = b_.CreateAllocaBPF();
+    b_.CreateStore(b_.CreateAdd(oldval, b_.getInt64(1)), newval);
+    b_.CreateMapUpdateElem(map, key, newval);
   }
   else
   {
@@ -222,6 +224,28 @@ AllocaInst *CodegenLLVM::getMapKey(Map &map)
   return key;
 }
 
+AllocaInst *CodegenLLVM::getQuantizeMapKey(Map &map, Value *log2)
+{
+  AllocaInst *key;
+  if (map.vargs) {
+    key = b_.CreateAllocaBPF(map.vargs->size() + 1);
+    int i = 0;
+    for (Expression *expr : *map.vargs) {
+      expr->accept(*this);
+      Value *offset = b_.CreateGEP(key, b_.getInt64(i++));
+      b_.CreateStore(expr_, offset);
+    }
+    Value *offset = b_.CreateGEP(key, b_.getInt64(i));
+    b_.CreateStore(log2, offset);
+  }
+  else
+  {
+    key = b_.CreateAllocaBPF();
+    b_.CreateStore(log2, key);
+  }
+  return key;
+}
+
 class BPFtraceMemoryManager : public SectionMemoryManager
 {
 public:

src/codegen_llvm.h

@@ -37,6 +37,7 @@ public:
   void visit(Probe &probe) override;
   void visit(Program &program) override;
   AllocaInst *getMapKey(Map &map);
+  AllocaInst *getQuantizeMapKey(Map &map, Value *log2);
 
   void createLog2Function();
   int compile(bool debug=false);

src/map.cpp

@@ -24,15 +24,15 @@ Map::Map(std::string &name, Type type, std::vector<Type> &args)
           abort();
       }
     }
+    if (type == Type::quantize)
+    {
+      key_size += 8;
+    }
   }
   else
   {
     key_size = 8;
   }
-  if (type == Type::quantize)
-  {
-    key_size += 8;
-  }
 
   int value_size = 8;
   int max_entries = 128;