Implement key access for maps

src/bpftrace.cpp

@@ -86,14 +86,33 @@ int BPFtrace::print_maps()
   for(auto &mapmap : maps_)
   {
     Map &map = *mapmap.second.get();
-    uint64_t key = 0;
-    uint64_t next_key = 99;
+    auto map_args = map_args_.find(map.name_);
+    if (map_args == map_args_.end())
+      abort();
+
+    int key_elems = map_args->second.size(); // TODO handle more than just integer keys
+    if (key_elems == 0) key_elems = 1;
+    auto key = std::vector<uint64_t>(key_elems);
+    auto next_key = std::vector<uint64_t>(key_elems);
     uint64_t value;
-    int ret;
-    ret = bpf_get_next_key(map.mapfd_, &key, &next_key);
-    std::cout << ret << map.name_ << ":" << key << ":" << next_key << std::endl;
-    ret = bpf_lookup_elem(map.mapfd_, &key, &value);
-    std::cout << ret << map.name_ << ":" << key << ":" << value << std::endl;
+    int err;
+
+    err = bpf_get_next_key(map.mapfd_, key.data(), next_key.data());
+    if (err)
+      next_key = key;
+
+    std::cout << map.name_ << "[ ";
+    for (int i=0; i<key_elems; i++)
+      std::cout << next_key.at(i) << " ";
+    std::cout << "]: ";
+
+    err = bpf_lookup_elem(map.mapfd_, next_key.data(), &value);
+    std::cout << value << std::endl;
+    if (err)
+    {
+      std::cerr << "Error looking up elem: " << err << std::endl;
+      return -1;
+    }
   }
 }
 

src/codegen_llvm.cpp

@@ -59,8 +59,21 @@ void CodegenLLVM::visit(Call &call)
 
 void CodegenLLVM::visit(Map &map)
 {
-  AllocaInst *key = b_.CreateAllocaBPF(b_.getInt64Ty());
-  b_.CreateStore(b_.getInt64(0), key); // TODO variable key
+  AllocaInst *key;
+  if (map.vargs) {
+    key = b_.CreateAllocaBPF(map.vargs->size());
+    int i = 0;
+    for (Expression *expr : *map.vargs) {
+      expr->accept(*this);
+      Value *offset = b_.CreateGEP(key, b_.getInt64(i++));
+      b_.CreateStore(expr_, offset);
+    }
+  }
+  else
+  {
+    key = b_.CreateAllocaBPF();
+    b_.CreateStore(b_.getInt64(0), key);
+  }
   expr_ = b_.CreateMapLookupElem(map, key);
 }
 
@@ -79,8 +92,8 @@ void CodegenLLVM::visit(Binop &binop)
     case ebpf::bpftrace::Parser::token::GE:    expr_ = b_.CreateICmpSGE(lhs, rhs); break;
     case ebpf::bpftrace::Parser::token::LT:    expr_ = b_.CreateICmpSLT(lhs, rhs); break;
     case ebpf::bpftrace::Parser::token::GT:    expr_ = b_.CreateICmpSGT(lhs, rhs); break;
-    case ebpf::bpftrace::Parser::token::LAND:  break;//expr_ = b_.CreateAnd(lhs, rhs); break; TODO
-    case ebpf::bpftrace::Parser::token::LOR:   break;//expr_ = b_.CreateOR(lhs, rhs); break; TODO
+    case ebpf::bpftrace::Parser::token::LAND:  abort();// TODO
+    case ebpf::bpftrace::Parser::token::LOR:   abort();// TODO
     case ebpf::bpftrace::Parser::token::PLUS:  expr_ = b_.CreateAdd    (lhs, rhs); break;
     case ebpf::bpftrace::Parser::token::MINUS: expr_ = b_.CreateSub    (lhs, rhs); break;
     case ebpf::bpftrace::Parser::token::MUL:   expr_ = b_.CreateMul    (lhs, rhs); break;
@@ -112,10 +125,24 @@ void CodegenLLVM::visit(ExprStatement &expr)
 void CodegenLLVM::visit(AssignMapStatement &assignment)
 {
   Map &map = *assignment.map;
-  AllocaInst *key = b_.CreateAllocaBPF(b_.getInt64Ty());
-  AllocaInst *val = b_.CreateAllocaBPF(b_.getInt64Ty());
+  AllocaInst *val = b_.CreateAllocaBPF();
+
+  AllocaInst *key;
+  if (map.vargs) {
+    key = b_.CreateAllocaBPF(map.vargs->size());
+    int i = 0;
+    for (Expression *expr : *map.vargs) {
+      expr->accept(*this);
+      Value *offset = b_.CreateGEP(key, b_.getInt64(i++));
+      b_.CreateStore(expr_, offset);
+    }
+  }
+  else
+  {
+    key = b_.CreateAllocaBPF();
+    b_.CreateStore(b_.getInt64(0), key);
+  }
 
-  b_.CreateStore(b_.getInt64(0), key); // TODO variable key
   assignment.expr->accept(*this);
   b_.CreateStore(expr_, val);
 
@@ -129,10 +156,10 @@ void CodegenLLVM::visit(AssignMapCallStatement &assignment)
 
   if (call.func == "count")
   {
-    AllocaInst *key = b_.CreateAllocaBPF(b_.getInt64Ty());
+    AllocaInst *key = b_.CreateAllocaBPF();
     b_.CreateStore(b_.getInt64(0), key); // TODO variable key
     Value *oldval = b_.CreateMapLookupElem(map, key);
-    AllocaInst *newval = b_.CreateAllocaBPF(b_.getInt64Ty());
+    AllocaInst *newval = b_.CreateAllocaBPF();
     b_.CreateStore(b_.CreateAdd(oldval, b_.getInt64(1)), newval);
     b_.CreateMapUpdateElem(map, key, newval);
   }

src/irbuilderbpf.cpp

@@ -26,22 +26,48 @@ IRBuilderBPF::IRBuilderBPF(LLVMContext &context,
       &module_);
 }
 
-AllocaInst *IRBuilderBPF::CreateAllocaBPF(llvm::Type *ty, const std::string &name) const
+AllocaInst *IRBuilderBPF::CreateAllocaBPF(int num_items, const std::string &name)
 {
+  llvm::Type *ty = getInt64Ty();
+  Value *array_size = getInt64(num_items);
   Function *parent = GetInsertBlock()->getParent();
   BasicBlock &entry_block = parent->getEntryBlock();
   if (entry_block.empty())
-    return new AllocaInst(ty, "", &entry_block);
+    return new AllocaInst(ty, array_size, "", &entry_block);
   else
-    return new AllocaInst(ty, "", &entry_block.front());
+    return new AllocaInst(ty, array_size, "", &entry_block.front());
 }
 
 CallInst *IRBuilderBPF::CreateBpfPseudoCall(Map &map)
 {
   int mapfd;
   if (bpftrace_.maps_.find(map.ident) == bpftrace_.maps_.end()) {
-    bpftrace_.maps_[map.ident] = std::make_unique<ebpf::bpftrace::Map>(map.ident);
+    // Create map since it doesn't already exist
+    int key_size = 0;
+    if (map.vargs)
+    {
+      auto map_args = bpftrace_.map_args_.find(map.ident);
+      if (map_args == bpftrace_.map_args_.end())
+        abort();
+      for (auto type : map_args->second)
+      {
+        switch (type)
+        {
+          case Type::integer:
+            key_size += 8;
+            break;
+          default:
+            abort();
+        }
+      }
+    }
+    else
+    {
+      key_size = 8;
+    }
+    bpftrace_.maps_[map.ident] = std::make_unique<ebpf::bpftrace::Map>(map.ident, key_size);
   }
+
   mapfd = bpftrace_.maps_[map.ident]->mapfd_;
   Function *pseudo_func = module_.getFunction("llvm.bpf.pseudo");
   return CreateCall(pseudo_func, {getInt64(BPF_PSEUDO_MAP_FD), getInt64(mapfd)});
@@ -70,7 +96,7 @@ LoadInst *IRBuilderBPF::CreateMapLookupElem(Map &map, AllocaInst *key)
   BasicBlock *lookup_failure_block = BasicBlock::Create(module_.getContext(), "lookup_failure", parent);
   BasicBlock *lookup_merge_block = BasicBlock::Create(module_.getContext(), "lookup_merge", parent);
 
-  Value *value = CreateAllocaBPF(getInt64Ty());
+  Value *value = CreateAllocaBPF();
   Value *condition = CreateICmpNE(
       CreateIntCast(call, getInt8PtrTy(), true),
       ConstantExpr::getCast(Instruction::IntToPtr, getInt64(0), getInt8PtrTy()),

src/irbuilderbpf.h

@@ -18,7 +18,7 @@ public:
                Module &module,
                BPFtrace &bpftrace);
 
-  AllocaInst *CreateAllocaBPF(llvm::Type *ty, const std::string &name="") const;
+  AllocaInst *CreateAllocaBPF(int array_size=1, const std::string &name="");
   CallInst   *CreateBpfPseudoCall(Map &map);
   LoadInst   *CreateMapLookupElem(Map &map, AllocaInst *key);
   void        CreateMapUpdateElem(Map &map, AllocaInst *key, AllocaInst *val);

src/map.cpp

@@ -7,8 +7,11 @@
 namespace ebpf {
 namespace bpftrace {
 
-Map::Map(std::string &name) : name_(name) {
-  mapfd_ = bpf_create_map(BPF_MAP_TYPE_HASH, 8, 8, 128, 0);
+Map::Map(std::string &name, int key_size) : name_(name) {
+  int value_size = 8;
+  int max_entries = 128;
+  int flags = 0;
+  mapfd_ = bpf_create_map(BPF_MAP_TYPE_HASH, key_size, value_size, max_entries, flags);
   if (mapfd_ < 0)
   {
     std::cerr << "Error creating map: '" << name_ << "'" << std::endl;

src/map.h

@@ -8,7 +8,7 @@ namespace bpftrace {
 
 class Map {
 public:
-  explicit Map(std::string &);
+  Map(std::string &, int key_size);
   ~Map();
   Map(const Map &) = delete;
   Map& operator=(const Map &) = delete;

src/semantic_analyser.cpp

@@ -76,7 +76,7 @@ void SemanticAnalyser::visit(Map &map)
       err_ << "Argument mismatch for " << map.ident << ": ";
       err_ << "trying to access with arguments: [ ";
       for (Type t : args) { err_ << typestr(t) << " "; }
-      err_ << "]\n\twhen map already uses the arguments: [ ";
+      err_ << "]\n\twhen map expects arguments: [ ";
       for (Type t : search->second) { err_ << typestr(t) << " "; }
       err_ << "]\n" << std::endl;
     }