Printf: Work with non 64-bit integers

src/ast/codegen_llvm.cpp

@@ -177,19 +177,26 @@ void CodegenLLVM::visit(Call &call)
     std::vector<llvm::Type *> elements = { b_.getInt64Ty() }; // printf ID
     String &fmt = static_cast<String&>(*call.vargs->at(0));
 
-    static int printf_id = 0;
-    auto args = std::get<1>(bpftrace_.printf_args_.at(printf_id));
-    for (SizedType &t : args)
+    auto &args = std::get<1>(bpftrace_.printf_args_.at(printf_id_));
+    for (Field &arg : args)
     {
-      llvm::Type *ty = b_.GetType(t);
+      llvm::Type *ty = b_.GetType(arg.type);
       elements.push_back(ty);
     }
-    StructType *printf_struct = StructType::create(elements, "printf_t", true);
+    StructType *printf_struct = StructType::create(elements, "printf_t", false);
     int struct_size = layout_.getTypeAllocSize(printf_struct);
 
+    auto *struct_layout = layout_.getStructLayout(printf_struct);
+    for (int i=0; i<args.size(); i++)
+    {
+      Field &arg = args[i];
+      arg.offset = struct_layout->getElementOffset(i+1); // +1 for the printf_id field
+    }
+
     AllocaInst *printf_args = b_.CreateAllocaBPF(printf_struct, "printf_args");
+    b_.CreateMemSet(printf_args, b_.getInt8(0), struct_size, 1);
 
-    b_.CreateStore(b_.getInt64(printf_id), printf_args);
+    b_.CreateStore(b_.getInt64(printf_id_), printf_args);
     for (int i=1; i<call.vargs->size(); i++)
     {
       Expression &arg = *call.vargs->at(i);
@@ -201,7 +208,7 @@ void CodegenLLVM::visit(Call &call)
         b_.CreateStore(expr_, offset);
     }
 
-    printf_id++;
+    printf_id_++;
     b_.CreatePerfEventOutput(ctx_, printf_args, struct_size);
     b_.CreateLifetimeEnd(printf_args);
     expr_ = nullptr;

src/ast/codegen_llvm.h

@@ -65,6 +65,7 @@ private:
   BPFtrace &bpftrace_;
 
   std::map<std::string, Value *> variables_;
+  int printf_id_ = 0;
 };
 
 } // namespace ast

src/ast/semantic_analyser.cpp

@@ -137,9 +137,9 @@ void SemanticAnalyser::visit(Call &call)
       if (is_final_pass()) {
         auto &fmt_arg = *call.vargs->at(0);
         String &fmt = static_cast<String&>(fmt_arg);
-        std::vector<SizedType> args;
+        std::vector<Field> args;
         for (auto iter = call.vargs->begin()+1; iter != call.vargs->end(); iter++) {
-          args.push_back((*iter)->type);
+          args.push_back({ .type = (*iter)->type });
         }
         err_ << verify_format_string(fmt.str, args);
 

src/bpftrace.cpp

@@ -128,7 +128,7 @@ void perf_event_printer(void *cb_cookie, void *data, int size)
 {
   auto bpftrace = static_cast<BPFtrace*>(cb_cookie);
   auto printf_id = *static_cast<uint64_t*>(data);
-  auto arg_data = static_cast<uint8_t*>(data) + sizeof(uint64_t);
+  auto arg_data = static_cast<uint8_t*>(data);
 
   auto fmt = std::get<0>(bpftrace->printf_args_[printf_id]).c_str();
   auto args = std::get<1>(bpftrace->printf_args_[printf_id]);
@@ -136,28 +136,43 @@ void perf_event_printer(void *cb_cookie, void *data, int size)
   std::vector<std::unique_ptr<char>> resolved_symbols;
   for (auto arg : args)
   {
-    switch (arg.type)
+    switch (arg.type.type)
     {
       case Type::integer:
-        arg_values.push_back(*(uint64_t*)arg_data);
+        switch (arg.type.size)
+        {
+          case 8:
+            arg_values.push_back(*(uint64_t*)(arg_data+arg.offset));
+            break;
+          case 4:
+            arg_values.push_back(*(uint32_t*)(arg_data+arg.offset));
+            break;
+          case 2:
+            arg_values.push_back(*(uint16_t*)(arg_data+arg.offset));
+            break;
+          case 1:
+            arg_values.push_back(*(uint8_t*)(arg_data+arg.offset));
+            break;
+          default:
+            abort();
+        }
         break;
       case Type::string:
-        arg_values.push_back((uint64_t)arg_data);
+        arg_values.push_back((uint64_t)(arg_data+arg.offset));
         break;
       case Type::sym:
         resolved_symbols.emplace_back(strdup(
-              bpftrace->resolve_sym(*(uint64_t*)arg_data).c_str()));
+              bpftrace->resolve_sym(*(uint64_t*)(arg_data+arg.offset)).c_str()));
         arg_values.push_back((uint64_t)resolved_symbols.back().get());
         break;
       case Type::usym:
         resolved_symbols.emplace_back(strdup(
-              bpftrace->resolve_usym(*(uint64_t*)arg_data).c_str()));
+              bpftrace->resolve_usym(*(uint64_t*)(arg_data+arg.offset)).c_str()));
         arg_values.push_back((uint64_t)resolved_symbols.back().get());
         break;
       default:
         abort();
     }
-    arg_data +=  arg.size;
   }
 
   switch (args.size())

src/bpftrace.h

@@ -33,7 +33,7 @@ public:
 
   std::map<std::string, std::unique_ptr<IMap>> maps_;
   std::map<std::string, Struct> structs_;
-  std::vector<std::tuple<std::string, std::vector<SizedType>>> printf_args_;
+  std::vector<std::tuple<std::string, std::vector<Field>>> printf_args_;
   std::unique_ptr<IMap> stackid_map_;
   std::unique_ptr<IMap> perf_event_map_;
 

src/printf.cpp

@@ -2,10 +2,11 @@
 
 #include "printf.h"
 #include "printf_format_types.h"
+#include "struct.h"
 
 namespace bpftrace {
 
-std::string verify_format_string(const std::string &fmt, std::vector<SizedType> args)
+std::string verify_format_string(const std::string &fmt, std::vector<Field> args)
 {
   std::stringstream message;
   const std::regex re("%-?[0-9]*[a-zA-Z]+");
@@ -31,7 +32,7 @@ std::string verify_format_string(const std::string &fmt, std::vector<SizedType>
   auto token_iter = tokens_begin;
   for (int i=0; i<num_args; i++, token_iter++)
   {
-    Type arg_type = args.at(i).type;
+    Type arg_type = args.at(i).type.type;
     if (arg_type == Type::sym || arg_type == Type::usym)
       arg_type = Type::string; // Symbols should be printed as strings
     int offset = 1;

src/printf.h

@@ -5,6 +5,8 @@
 
 namespace bpftrace {
 
-std::string verify_format_string(const std::string &fmt, std::vector<SizedType> args);
+struct Field;
+
+std::string verify_format_string(const std::string &fmt, std::vector<Field> args);
 
 } // namespace bpftrace

tests/codegen.cpp

@@ -24,12 +24,44 @@ TEST(codegen, populate_sections)
   ASSERT_EQ(semantics.analyse(), 0);
   std::stringstream out;
   ast::CodegenLLVM codegen(driver.root_, bpftrace);
-  auto bpforc = codegen.compile(true, out);
+  auto bpforc = codegen.compile();
 
   // Check sections are populated
-  ASSERT_EQ(bpforc->sections_.size(), 2);
-  ASSERT_EQ(bpforc->sections_.count("s_kprobe:foo"), 1);
-  ASSERT_EQ(bpforc->sections_.count("s_kprobe:bar"), 1);
+  EXPECT_EQ(bpforc->sections_.size(), 2);
+  EXPECT_EQ(bpforc->sections_.count("s_kprobe:foo"), 1);
+  EXPECT_EQ(bpforc->sections_.count("s_kprobe:bar"), 1);
+}
+
+TEST(codegen, printf_offsets)
+{
+  BPFtrace bpftrace;
+  Driver driver;
+
+  ASSERT_EQ(driver.parse_str("struct Foo { char c; int i; } kprobe:f { $foo = (Foo*)0; printf(\"%c %u\\n\", $foo->c, $foo->i) }"), 0);
+  ClangParser clang;
+  clang.parse(driver.root_, bpftrace.structs_);
+  ast::SemanticAnalyser semantics(driver.root_, bpftrace);
+  ASSERT_EQ(semantics.analyse(), 0);
+  ASSERT_EQ(semantics.create_maps(true), 0);
+  std::stringstream out;
+  ast::CodegenLLVM codegen(driver.root_, bpftrace);
+  auto bpforc = codegen.compile();
+
+  EXPECT_EQ(bpftrace.printf_args_.size(), 1);
+  auto &fmt = std::get<0>(bpftrace.printf_args_[0]);
+  auto &args = std::get<1>(bpftrace.printf_args_[0]);
+
+  EXPECT_EQ(fmt, "%c %u\n");
+
+  EXPECT_EQ(args.size(), 2);
+
+  EXPECT_EQ(args[0].type.type, Type::integer);
+  EXPECT_EQ(args[0].type.size, 1);
+  EXPECT_EQ(args[0].offset, 8);
+
+  EXPECT_EQ(args[1].type.type, Type::integer);
+  EXPECT_EQ(args[1].type.size, 4);
+  EXPECT_EQ(args[1].offset, 12);
 }
 
 std::string header = R"HEAD(; ModuleID = 'bpftrace'
@@ -921,9 +953,9 @@ attributes #1 = { argmemonly nounwind }
 
 TEST(codegen, call_printf)
 {
-  test("kprobe:f { printf(\"hello\\n\") }",
+  test("struct Foo { char c; long l; } kprobe:f { $foo = (Foo*)0; printf(\"%c %lu\\n\", $foo->c, $foo->l) }",
 
-R"EXPECTED(%printf_t = type { i64 }
+R"EXPECTED(%printf_t = type { i64, i8, i64 }
 
 ; Function Attrs: nounwind
 declare i64 @llvm.bpf.pseudo(i64, i64) #0
@@ -933,17 +965,36 @@ declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
 
 define i64 @"kprobe:f"(i8*) local_unnamed_addr section "s_kprobe:f" {
 entry:
+  %Foo.l = alloca i64, align 8
+  %Foo.c = alloca i8, align 1
   %printf_args = alloca %printf_t, align 8
   %1 = bitcast %printf_t* %printf_args to i8*
   call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
-  store i64 0, %printf_t* %printf_args, align 8
-  %pseudo = tail call i64 @llvm.bpf.pseudo(i64 1, i64 1)
-  %get_cpu_id = tail call i64 inttoptr (i64 8 to i64 ()*)()
-  %perf_event_output = call i64 inttoptr (i64 25 to i64 (i8*, i8*, i64, i8*, i64)*)(i8* %0, i64 %pseudo, i64 %get_cpu_id, %printf_t* nonnull %printf_args, i64 8)
+  %2 = bitcast %printf_t* %printf_args to i8*
+  call void @llvm.memset.p0i8.i64(i8* nonnull %2, i8 0, i64 16, i32 8, i1 false)
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %Foo.c)
+  %probe_read = call i64 inttoptr (i64 4 to i64 (i8*, i64, i8*)*)(i8* nonnull %Foo.c, i64 1, i64 0)
+  %3 = load i8, i8* %Foo.c, align 1
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %Foo.c)
+  %4 = getelementptr inbounds %printf_t, %printf_t* %printf_args, i64 0, i32 1
+  store i8 %3, i8* %4, align 8
+  %5 = bitcast i64* %Foo.l to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %5)
+  %probe_read1 = call i64 inttoptr (i64 4 to i64 (i8*, i64, i8*)*)(i64* nonnull %Foo.l, i64 8, i64 8)
+  %6 = load i64, i64* %Foo.l, align 8
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %5)
+  %7 = getelementptr inbounds %printf_t, %printf_t* %printf_args, i64 0, i32 2
+  store i64 %6, i64* %7, align 8
+  %pseudo = call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %get_cpu_id = call i64 inttoptr (i64 8 to i64 ()*)()
+  %perf_event_output = call i64 inttoptr (i64 25 to i64 (i8*, i8*, i64, i8*, i64)*)(i8* %0, i64 %pseudo, i64 %get_cpu_id, %printf_t* nonnull %printf_args, i64 20)
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
   ret i64 0
 }
 
+; Function Attrs: argmemonly nounwind
+declare void @llvm.memset.p0i8.i64(i8* nocapture writeonly, i8, i64, i32, i1) #1
+
 ; Function Attrs: argmemonly nounwind
 declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1