Allow different sized map keys (not just 64-bit ints)

src/CMakeLists.txt

@@ -21,6 +21,7 @@ add_executable(bpftrace
   lex.yy.cc
   main.cpp
   map.cpp
+  mapkey.cpp
   parser.tab.cc
   printer.cpp
   semantic_analyser.cpp

src/bpftrace.cpp

@@ -70,10 +70,10 @@ int BPFtrace::print_maps()
 
 int BPFtrace::print_map(Map &map)
 {
-  std::vector<uint64_t> old_key;
+  std::vector<uint8_t> old_key;
   try
   {
-    old_key = find_empty_key(map, map.args_.size());
+    old_key = find_empty_key(map, map.key_.size());
   }
   catch (std::runtime_error &e)
   {
@@ -85,7 +85,7 @@ int BPFtrace::print_map(Map &map)
 
   while (bpf_get_next_key(map.mapfd_, old_key.data(), key.data()) == 0)
   {
-    std::cout << map.name_ << argument_list(key, map.args_.size()) << ": ";
+    std::cout << map.name_ << map.key_.argument_value_list(key) << ": ";
 
     uint64_t value;
     int err = bpf_lookup_elem(map.mapfd_, key.data(), &value);
@@ -111,10 +111,10 @@ int BPFtrace::print_map_quantize(Map &map)
   // e.g. A map defined as: @x[1, 2] = @quantize(3);
   // would actually be stored with the key: [1, 2, 3]
 
-  std::vector<uint64_t> old_key;
+  std::vector<uint8_t> old_key;
   try
   {
-    old_key = find_empty_key(map, map.args_.size() + 1);
+    old_key = find_empty_key(map, map.key_.size() + 8);
   }
   catch (std::runtime_error &e)
   {
@@ -124,14 +124,14 @@ int BPFtrace::print_map_quantize(Map &map)
   }
   auto key(old_key);
 
-  std::map<std::vector<uint64_t>, std::vector<uint64_t>> values;
+  std::map<std::vector<uint8_t>, std::vector<uint64_t>> values_by_key;
 
   while (bpf_get_next_key(map.mapfd_, old_key.data(), key.data()) == 0)
   {
-    auto key_prefix = std::vector<uint64_t>(map.args_.size());
-    int bucket = key.at(map.args_.size());
+    auto key_prefix = std::vector<uint8_t>(map.key_.size());
+    int bucket = key.at(map.key_.size());
 
-    for (int i=0; i<map.args_.size(); i++)
+    for (size_t i=0; i<map.key_.size(); i++)
       key_prefix.at(i) = key.at(i);
 
     uint64_t value;
@@ -142,19 +142,19 @@ int BPFtrace::print_map_quantize(Map &map)
       return -1;
     }
 
-    if (values.find(key_prefix) == values.end())
+    if (values_by_key.find(key_prefix) == values_by_key.end())
     {
       // New key - create a list of buckets for it
-      values[key_prefix] = std::vector<uint64_t>(65);
+      values_by_key[key_prefix] = std::vector<uint64_t>(65);
     }
-    values[key_prefix].at(bucket) = value;
+    values_by_key[key_prefix].at(bucket) = value;
 
     old_key = key;
   }
 
-  for (auto &map_elem : values)
+  for (auto &map_elem : values_by_key)
   {
-    std::cout << map.name_ << argument_list(map_elem.first) << ": " << std::endl;
+    std::cout << map.name_ << map.key_.argument_value_list(map_elem.first) << ": " << std::endl;
 
     print_quantize(map_elem.second);
 
@@ -238,26 +238,20 @@ std::string BPFtrace::quantize_index_label(int power)
   return label.str();
 }
 
-std::vector<uint64_t> BPFtrace::find_empty_key(Map &map, int num_elems)
+std::vector<uint8_t> BPFtrace::find_empty_key(Map &map, size_t size)
 {
-  if (num_elems == 0) num_elems = 1;
-  auto key = std::vector<uint64_t>(num_elems);
-  uint64_t value;
+  if (size == 0) size = 8;
+  auto key = std::vector<uint8_t>(size);
+  uint8_t value;
 
   if (bpf_lookup_elem(map.mapfd_, key.data(), &value))
     return key;
 
-  for (auto &elem : key)
-  {
-    elem = 0xffffffffffffffff;
-  }
+  for (auto &elem : key) elem = 0xff;
   if (bpf_lookup_elem(map.mapfd_, key.data(), &value))
     return key;
 
-  for (auto &elem : key)
-  {
-    elem = 0x5555555555555555;
-  }
+  for (auto &elem : key) elem = 0x55;
   if (bpf_lookup_elem(map.mapfd_, key.data(), &value))
     return key;
 

src/bpftrace.h

@@ -31,7 +31,7 @@ private:
   int print_map_quantize(Map &map);
   int print_quantize(std::vector<uint64_t> values);
   std::string quantize_index_label(int power);
-  std::vector<uint64_t> find_empty_key(Map &map, int num_elems);
+  std::vector<uint8_t> find_empty_key(Map &map, size_t size);
 };
 
 } // namespace bpftrace

src/map.cpp

@@ -6,32 +6,14 @@
 
 namespace bpftrace {
 
-Map::Map(std::string &name, Type type, std::vector<Type> &args)
-  : name_(name), type_(type), args_(args)
+Map::Map(std::string &name, Type type, MapKey key)
+  : name_(name), type_(type), key_(key)
 {
-  int key_size = 0;
-  if (args.size() > 0)
-  {
-    for (auto type : args)
-    {
-      switch (type)
-      {
-        case Type::integer:
-          key_size += 8;
-          break;
-        default:
-          abort();
-      }
-    }
-    if (type == Type::quantize)
-    {
-      key_size += 8;
-    }
-  }
-  else
-  {
+  int key_size = key.size();
+  if (type == Type::quantize)
+    key_size += 8;
+  if (key_size == 0)
     key_size = 8;
-  }
 
   int value_size = 8;
   int max_entries = 128;

src/map.h

@@ -4,13 +4,14 @@
 #include <string>
 #include <vector>
 
+#include "mapkey.h"
 #include "types.h"
 
 namespace bpftrace {
 
 class Map {
 public:
-  Map(std::string &, Type type, std::vector<Type> &args);
+  Map(std::string &name, Type type, MapKey key);
   ~Map();
   Map(const Map &) = delete;
   Map& operator=(const Map &) = delete;
@@ -18,7 +19,7 @@ public:
   int mapfd_;
   std::string name_;
   Type type_;
-  std::vector<Type> args_;
+  MapKey key_;
 };
 
 } // namespace bpftrace

src/mapkey.cpp

+#include "mapkey.h"
+
+namespace bpftrace {
+
+bool MapKey::operator!=(const MapKey &k) const
+{
+  return args_ != k.args_;
+}
+
+size_t MapKey::size() const
+{
+  size_t size = 0;
+  for (auto &arg : args_)
+    size += arg.size;
+  return size;
+}
+
+std::string MapKey::argument_type_list() const
+{
+  size_t n = args_.size();
+  if (n == 0)
+    return "[]";
+
+  std::ostringstream list;
+  list << "[";
+  for (size_t i = 0; i < n-1; i++)
+    list << args_.at(i) << ", ";
+  list << args_.at(n-1) << "]";
+  return list.str();
+}
+
+std::string MapKey::argument_value_list(const std::vector<uint8_t> &data) const
+{
+  size_t n = args_.size();
+  if (n == 0)
+    return "";
+
+  std::ostringstream list;
+  list << "[";
+  int offset = 0;
+  for (size_t i = 0; i < n-1; i++)
+  {
+    const MapKeyArgument &arg = args_.at(i);
+    list << argument_value(arg, &data.at(offset)) << ", ";
+    offset += arg.size;
+  }
+  const MapKeyArgument &arg = args_.at(n-1);
+  list << argument_value(arg, &data.at(offset)) << "]";
+  return list.str();
+}
+
+std::string MapKey::argument_value(const MapKeyArgument &arg, const void *data) const
+{
+  switch (arg.type)
+  {
+    case Type::integer:
+      switch (arg.size)
+      {
+        case 8:
+          return std::to_string(*(uint64_t*)data);
+        case 4:
+          return std::to_string(*(uint32_t*)data);
+      }
+  }
+  abort();
+}
+
+} // namespace bpftrace

src/mapkey.h

+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "types.h"
+
+namespace bpftrace {
+
+class MapKey
+{
+public:
+  std::vector<MapKeyArgument> args_;
+
+  bool operator!=(const MapKey &k) const;
+
+  size_t size() const;
+  std::string argument_type_list() const;
+  std::string argument_value_list(const std::vector<uint8_t> &data) const;
+
+private:
+  std::string argument_value(const MapKeyArgument &arg, const void *data) const;
+};
+
+} // namespace bpftrace

src/semantic_analyser.cpp

@@ -66,27 +66,27 @@ void SemanticAnalyser::visit(Call &call)
 
 void SemanticAnalyser::visit(Map &map)
 {
-  std::vector<Type> args;
+  MapKey key;
   if (map.vargs) {
     for (Expression *expr : *map.vargs) {
       expr->accept(*this);
-      args.push_back(type_);
+      key.args_.push_back({type_, 8});
     }
   }
 
-  auto search = map_args_.find(map.ident);
-  if (search != map_args_.end()) {
-    if (search->second != args) {
+  auto search = map_key_.find(map.ident);
+  if (search != map_key_.end()) {
+    if (search->second != key) {
       err_ << "Argument mismatch for " << map.ident << ": ";
       err_ << "trying to access with arguments: ";
-      err_ << argument_list(args, true);
+      err_ << key.argument_type_list();
       err_ << "\n\twhen map expects arguments: ";
-      err_ << argument_list(search->second, true);
+      err_ << search->second.argument_type_list();
       err_ << "\n" << std::endl;
     }
   }
   else {
-    map_args_.insert({map.ident, args});
+    map_key_.insert({map.ident, key});
   }
 
   auto search_val = map_val_.find(map.ident);
@@ -237,12 +237,12 @@ int SemanticAnalyser::create_maps()
     std::string map_name = map_val.first;
     Type type = map_val.second;
 
-    auto search_args = map_args_.find(map_name);
-    if (search_args == map_args_.end())
+    auto search_args = map_key_.find(map_name);
+    if (search_args == map_key_.end())
       abort();
-    auto &args = search_args->second;
+    auto &key = search_args->second;
 
-    bpftrace_.maps_[map_name] = std::make_unique<bpftrace::Map>(map_name, type, args);
+    bpftrace_.maps_[map_name] = std::make_unique<bpftrace::Map>(map_name, type, key);
   }
 
   return 0;

src/semantic_analyser.h

@@ -47,7 +47,7 @@ private:
   bool is_final_pass() const;
 
   std::map<std::string, Type> map_val_;
-  std::map<std::string, std::vector<Type>> map_args_;
+  std::map<std::string, MapKey> map_key_;
 };
 
 } // namespace ast

src/types.cpp

@@ -10,6 +10,17 @@ std::ostream &operator<<(std::ostream &os, Type type)
   return os;
 }
 
+std::ostream &operator<<(std::ostream &os, MapKeyArgument arg)
+{
+  os << arg.type;
+  return os;
+}
+
+bool MapKeyArgument::operator==(const MapKeyArgument &a) const
+{
+  return type == a.type && size == a.size;
+}
+
 std::string typestr(Type t)
 {
   switch (t)

src/types.h

@@ -3,8 +3,8 @@
 #include <ostream>
 #include <sstream>
 #include <string>
-#include <unistd.h>
 #include <vector>
+#include <unistd.h>
 
 #include "libbpf.h"
 
@@ -20,6 +20,17 @@ enum class Type
 
 std::ostream &operator<<(std::ostream &os, Type type);
 
+class MapKeyArgument
+{
+public:
+  Type type;
+  size_t size;
+
+  bool operator==(const MapKeyArgument &a) const;
+};
+
+std::ostream &operator<<(std::ostream &os, MapKeyArgument arg);
+
 enum class ProbeType
 {
   kprobe,
@@ -38,28 +49,4 @@ public:
   std::string name;
 };
 
-template <typename T>
-std::string argument_list(const std::vector<T> &items, size_t n, bool show_empty=false)
-{
-  if (n == 0)
-  {
-    if (show_empty)
-      return "[]";
-    return "";
-  }
-
-  std::ostringstream list;
-  list << "[";
-  for (size_t i = 0; i < n-1; i++)
-    list << items.at(i) << ", ";
-  list << items.at(n-1) << "]";
-  return list.str();
-}
-
-template <typename T>
-std::string argument_list(const std::vector<T> &items, bool show_empty=false)
-{
-  return argument_list(items, items.size(), show_empty);
-}
-
 } // namespace bpftrace

tests/CMakeLists.txt

@@ -13,6 +13,7 @@ add_executable(bpftrace_test
   ${CMAKE_SOURCE_DIR}/src/attached_probe.cpp
   ${CMAKE_SOURCE_DIR}/src/bpftrace.cpp
   ${CMAKE_SOURCE_DIR}/src/map.cpp
+  ${CMAKE_SOURCE_DIR}/src/mapkey.cpp
   ${CMAKE_SOURCE_DIR}/src/semantic_analyser.cpp
   ${CMAKE_SOURCE_DIR}/src/types.cpp
 )