Calculate key size in Map's constructor

src/CMakeLists.txt

@@ -24,6 +24,7 @@ add_executable(bpftrace
   parser.tab.cc
   printer.cpp
   semantic_analyser.cpp
+  types.cpp
 )
 
 llvm_map_components_to_libnames(llvm_libs bpfcodegen ipo irreader mcjit)

src/bpftrace.cpp

@@ -6,38 +6,6 @@
 namespace ebpf {
 namespace bpftrace {
 
-std::string typestr(Type t)
-{
-  switch (t)
-  {
-    case Type::none:     return "none";     break;
-    case Type::integer:  return "integer";  break;
-    case Type::quantize: return "quantize"; break;
-    case Type::count:    return "count";    break;
-    default: abort();
-  }
-}
-
-bpf_probe_attach_type attachtype(ProbeType t)
-{
-  switch (t)
-  {
-    case ProbeType::kprobe:    return BPF_PROBE_ENTRY;  break;
-    case ProbeType::kretprobe: return BPF_PROBE_RETURN; break;
-    default: abort();
-  }
-}
-
-bpf_prog_type progtype(ProbeType t)
-{
-  switch (t)
-  {
-    case ProbeType::kprobe:    return BPF_PROG_TYPE_KPROBE; break;
-    case ProbeType::kretprobe: return BPF_PROG_TYPE_KPROBE; break;
-    default: abort();
-  }
-}
-
 int BPFtrace::add_probe(ast::Probe &p)
 {
   Probe probe;

src/bpftrace.h

@@ -4,39 +4,13 @@
 #include <memory>
 #include <vector>
 
-#include "libbpf.h"
 #include "ast.h"
 #include "map.h"
+#include "types.h"
 
 namespace ebpf {
 namespace bpftrace {
 
-enum class Type
-{
-  none,
-  integer,
-  quantize,
-  count,
-};
-
-enum class ProbeType
-{
-  kprobe,
-  kretprobe,
-};
-
-std::string typestr(Type t);
-bpf_probe_attach_type attachtype(ProbeType t);
-bpf_prog_type progtype(ProbeType t);
-
-class Probe
-{
-public:
-  ProbeType type;
-  std::string attach_point;
-  std::string name;
-};
-
 class BPFtrace
 {
 public:

src/irbuilderbpf.cpp

@@ -43,29 +43,17 @@ CallInst *IRBuilderBPF::CreateBpfPseudoCall(Map &map)
   int mapfd;
   if (bpftrace_.maps_.find(map.ident) == bpftrace_.maps_.end()) {
     // 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);
+    auto search_val = bpftrace_.map_val_.find(map.ident);
+    if (search_val == bpftrace_.map_val_.end())
+      abort();
+    Type type = search_val->second;
+
+    auto search_args = bpftrace_.map_args_.find(map.ident);
+    if (search_args == bpftrace_.map_args_.end())
+      abort();
+    auto &args = search_args->second;
+
+    bpftrace_.maps_[map.ident] = std::make_unique<ebpf::bpftrace::Map>(map.ident, type, args);
   }
 
   mapfd = bpftrace_.maps_[map.ident]->mapfd_;

src/map.cpp

@@ -7,7 +7,33 @@
 namespace ebpf {
 namespace bpftrace {
 
-Map::Map(std::string &name, int key_size) : name_(name) {
+Map::Map(std::string &name, Type type, std::vector<Type> &args)
+  : name_(name)
+{
+  int key_size = 0;
+  if (args.size() > 0)
+  {
+    for (auto type : args)
+    {
+      switch (type)
+      {
+        case Type::integer:
+          key_size += 8;
+          break;
+        default:
+          abort();
+      }
+    }
+  }
+  else
+  {
+    key_size = 8;
+  }
+  if (type == Type::quantize)
+  {
+    key_size += 8;
+  }
+
   int value_size = 8;
   int max_entries = 128;
   int flags = 0;
@@ -18,7 +44,8 @@ Map::Map(std::string &name, int key_size) : name_(name) {
   }
 }
 
-Map::~Map() {
+Map::~Map()
+{
   close(mapfd_);
 }
 

src/map.h

@@ -2,13 +2,16 @@
 
 #include <map>
 #include <string>
+#include <vector>
+
+#include "types.h"
 
 namespace ebpf {
 namespace bpftrace {
 
 class Map {
 public:
-  Map(std::string &, int key_size);
+  Map(std::string &, Type type, std::vector<Type> &args);
   ~Map();
   Map(const Map &) = delete;
   Map& operator=(const Map &) = delete;

src/semantic_analyser.h

@@ -5,6 +5,7 @@
 #include "ast.h"
 #include "bpftrace.h"
 #include "map.h"
+#include "types.h"
 
 namespace ebpf {
 namespace bpftrace {

src/types.cpp

+#include "types.h"
+
+namespace ebpf {
+namespace bpftrace {
+
+std::string typestr(Type t)
+{
+  switch (t)
+  {
+    case Type::none:     return "none";     break;
+    case Type::integer:  return "integer";  break;
+    case Type::quantize: return "quantize"; break;
+    case Type::count:    return "count";    break;
+    default: abort();
+  }
+}
+
+bpf_probe_attach_type attachtype(ProbeType t)
+{
+  switch (t)
+  {
+    case ProbeType::kprobe:    return BPF_PROBE_ENTRY;  break;
+    case ProbeType::kretprobe: return BPF_PROBE_RETURN; break;
+    default: abort();
+  }
+}
+
+bpf_prog_type progtype(ProbeType t)
+{
+  switch (t)
+  {
+    case ProbeType::kprobe:    return BPF_PROG_TYPE_KPROBE; break;
+    case ProbeType::kretprobe: return BPF_PROG_TYPE_KPROBE; break;
+    default: abort();
+  }
+}
+
+} // namespace bpftrace
+} // namespace ebpf

src/types.h

+#pragma once
+
+#include <unistd.h>
+#include <string>
+
+#include "libbpf.h"
+
+namespace ebpf {
+namespace bpftrace {
+
+enum class Type
+{
+  none,
+  integer,
+  quantize,
+  count,
+};
+
+enum class ProbeType
+{
+  kprobe,
+  kretprobe,
+};
+
+std::string typestr(Type t);
+bpf_probe_attach_type attachtype(ProbeType t);
+bpf_prog_type progtype(ProbeType t);
+
+class Probe
+{
+public:
+  ProbeType type;
+  std::string attach_point;
+  std::string name;
+};
+
+} // namespace bpftrace
+} // namespace ebpf

tests/CMakeLists.txt

@@ -14,6 +14,7 @@ add_executable(bpftrace_test
   ${CMAKE_SOURCE_DIR}/src/bpftrace.cpp
   ${CMAKE_SOURCE_DIR}/src/map.cpp
   ${CMAKE_SOURCE_DIR}/src/semantic_analyser.cpp
+  ${CMAKE_SOURCE_DIR}/src/types.cpp
 )
 
 add_dependencies(bpftrace_test bcc-build)