Add reg() function and remove sp builtin

README.md

@@ -192,7 +192,6 @@ Variables:
 - `arg0`, `arg1`, ... etc. - Arguments to the function being traced
 - `retval` - Return value from function being traced
 - `func` - Name of the function currently being traced
-- `sp` - Stack pointer
 
 Functions:
 - `quantize(int n)` - Produce a log2 histogram of values of `n`
@@ -202,3 +201,4 @@ Functions:
 - `printf(char *fmt, ...)` - Write to stdout
 - `sym(void *p)` - Resolve kernel address
 - `usym(void *p)` - Resolve user space address (incomplete)
+- `reg(char *name)` - Returns the value stored in the named register

src/arch/arch.h

@@ -5,12 +5,11 @@
 namespace bpftrace {
 namespace arch {
 
-// Offsets are based off the pt_regs struct from the Linux kernel
+int offset(std::string reg_name);
 int max_arg();
 int arg_offset(int arg_num);
 int ret_offset();
 int pc_offset();
-int sp_offset();
 std::string name();
 
 } // namespace arch

src/arch/x86_64.cpp

 #include "arch.h"
 
+#include <algorithm>
 #include <array>
 
 namespace bpftrace {
 namespace arch {
 
-static std::array<int, 6> arg_offsets = {
-  14, // di
-  13, // si
-  12, // dx
-  11, // cx
-   9, // r8
-   8, // r9
+static std::array<std::string, 27> registers = {
+  "r15",
+  "r14",
+  "r13",
+  "r12",
+  "bp",
+  "bx",
+  "r11",
+  "r10",
+  "r9",
+  "r8",
+  "ax",
+  "cx",
+  "dx",
+  "si",
+  "di",
+  "orig_ax",
+  "ip",
+  "cs",
+  "flags",
+  "sp",
+  "ss",
+  "fs_base",
+  "gs_base",
+  "ds",
+  "es",
+  "fs",
+  "gs",
 };
 
+static std::array<std::string, 6> arg_registers = {
+  "di",
+  "si",
+  "dx",
+  "cx",
+  "r8",
+  "r9",
+};
+
+int offset(std::string reg_name)
+{
+  auto it = find(registers.begin(), registers.end(), reg_name);
+  if (it == registers.end())
+    return -1;
+  return distance(registers.begin(), it);
+}
+
 int max_arg()
 {
-  return arg_offsets.size() - 1;
+  return arg_registers.size() - 1;
 }
 
 int arg_offset(int arg_num)
 {
-  return arg_offsets.at(arg_num);
+  return offset(arg_registers.at(arg_num));
 }
 
 int ret_offset()
 {
-  return 10; // ax
+  return offset("ax");
 }
 
 int pc_offset()
 {
-  return 16; // ip
-}
-
-int sp_offset()
-{
-  return 19; // sp
+  return offset("ip");
 }
 
 std::string name()

src/ast/codegen_llvm.cpp

@@ -77,24 +77,21 @@ void CodegenLLVM::visit(Builtin &builtin)
   else if (!builtin.ident.compare(0, 3, "arg") && builtin.ident.size() == 4 &&
       builtin.ident.at(3) >= '0' && builtin.ident.at(3) <= '9' ||
       builtin.ident == "retval" ||
-      builtin.ident == "func" ||
-      builtin.ident == "sp")
+      builtin.ident == "func")
   {
     int offset;
     if (builtin.ident == "retval")
-      offset = arch::ret_offset() * sizeof(uintptr_t);
+      offset = arch::ret_offset();
     else if (builtin.ident == "func")
-      offset = arch::pc_offset() * sizeof(uintptr_t);
-    else if (builtin.ident == "sp")
-      offset = arch::sp_offset() * sizeof(uintptr_t);
+      offset = arch::pc_offset();
     else // argX
     {
       int arg_num = atoi(builtin.ident.substr(3).c_str());
-      offset = arch::arg_offset(arg_num) * sizeof(uintptr_t);
+      offset = arch::arg_offset(arg_num);
     }
 
     AllocaInst *dst = b_.CreateAllocaBPF(builtin.type, builtin.ident);
-    Value *src = b_.CreateGEP(ctx_, b_.getInt64(offset));
+    Value *src = b_.CreateGEP(ctx_, b_.getInt64(offset * sizeof(uintptr_t)));
     b_.CreateProbeRead(dst, 8, src);
     expr_ = b_.CreateLoad(dst);
     b_.CreateLifetimeEnd(dst);
@@ -157,8 +154,22 @@ void CodegenLLVM::visit(Call &call)
   }
   else if (call.func == "sym" || call.func == "usym")
   {
+    // We want expr_ to just pass through from the child node - don't set it here
     call.vargs->front()->accept(*this);
   }
+  else if (call.func == "reg")
+  {
+    auto &reg_name = static_cast<String&>(*call.vargs->at(0)).str;
+    int offset = arch::offset(reg_name);
+    if (offset == -1)
+      abort();
+
+    AllocaInst *dst = b_.CreateAllocaBPF(call.type, call.func+"_"+reg_name);
+    Value *src = b_.CreateGEP(ctx_, b_.getInt64(offset * sizeof(uintptr_t)));
+    b_.CreateProbeRead(dst, 8, src);
+    expr_ = b_.CreateLoad(dst);
+    b_.CreateLifetimeEnd(dst);
+  }
   else if (call.func == "printf")
   {
     ArrayType *string_type = ArrayType::get(b_.getInt8Ty(), STRING_SIZE);

src/ast/semantic_analyser.cpp

@@ -33,7 +33,6 @@ void SemanticAnalyser::visit(Builtin &builtin)
       builtin.ident == "uid" ||
       builtin.ident == "gid" ||
       builtin.ident == "cpu" ||
-      builtin.ident == "sp" ||
       builtin.ident == "retval") {
     builtin.type = SizedType(Type::integer, 8);
   }
@@ -122,7 +121,7 @@ void SemanticAnalyser::visit(Call &call)
   }
   else if (call.func == "str" || call.func == "sym" || call.func == "usym") {
     if (nargs != 1) {
-      err_ << call.func << "() should take 1 arguments (";
+      err_ << call.func << "() should take 1 argument (";
       err_ << nargs << " provided)" << std::endl;
     }
     if (is_final_pass() && call.vargs->at(0)->type.type != Type::integer) {
@@ -137,6 +136,29 @@ void SemanticAnalyser::visit(Call &call)
     else if (call.func == "usym")
       call.type = SizedType(Type::usym, 8);
   }
+  else if (call.func == "reg") {
+    if (nargs != 1) {
+      err_ << call.func << "() should take 1 argument (";
+      err_ << nargs << " provided)" << std::endl;
+    }
+    else {
+      auto &arg = *call.vargs->at(0);
+      if (arg.type.type != Type::string || !arg.is_literal) {
+        err_ << "reg() expects a string literal";
+        err_ << " (" << arg.type.type << " provided)" << std::endl;
+      }
+      else {
+        auto &reg_name = static_cast<String&>(arg).str;
+        int offset = arch::offset(reg_name);;
+        if (offset == -1) {
+          err_ << "'" << reg_name << "' is not a valid register on this architecture";
+          err_ << " (" << arch::name() << ")" << std::endl;
+        }
+      }
+    }
+
+    call.type = SizedType(Type::integer, 8);
+  }
   else if (call.func == "printf") {
     if (call.map) {
       err_ << "printf() should not be assigned to a map" << std::endl;
@@ -149,20 +171,22 @@ void SemanticAnalyser::visit(Call &call)
       err_ << "printf() can only take up to 7 arguments (";
       err_ << nargs << " provided)" << std::endl;
     }
-    Expression &fmt_arg = *call.vargs->at(0);
-    if (fmt_arg.type.type != Type::string || !fmt_arg.is_literal) {
-      err_ << "The first argument to printf() must be a string literal";
-      err_ << " (" << fmt_arg.type.type << " provided)" << std::endl;
-    }
-    if (is_final_pass()) {
-      String &fmt = static_cast<String&>(fmt_arg);
-      std::vector<SizedType> args;
-      for (auto iter = call.vargs->begin()+1; iter != call.vargs->end(); iter++) {
-        args.push_back((*iter)->type);
+    if (nargs > 0) {
+      Expression &fmt_arg = *call.vargs->at(0);
+      if (fmt_arg.type.type != Type::string || !fmt_arg.is_literal) {
+        err_ << "The first argument to printf() must be a string literal";
+        err_ << " (" << fmt_arg.type.type << " provided)" << std::endl;
+      }
+      if (is_final_pass()) {
+        String &fmt = static_cast<String&>(fmt_arg);
+        std::vector<SizedType> args;
+        for (auto iter = call.vargs->begin()+1; iter != call.vargs->end(); iter++) {
+          args.push_back((*iter)->type);
+        }
+        err_ << verify_format_string(fmt.str, args);
+
+        bpftrace_.printf_args_.push_back(std::make_tuple(fmt.str, args));
       }
-      err_ << verify_format_string(fmt.str, args);
-
-      bpftrace_.printf_args_.push_back(std::make_tuple(fmt.str, args));
     }
     call.type = SizedType(Type::string, STRING_SIZE);
   }

tests/codegen.cpp

@@ -648,9 +648,9 @@ attributes #1 = { argmemonly nounwind }
 )EXPECTED");
 }
 
-TEST(codegen, builtin_sp)
+TEST(codegen, call_reg) // Identical to builtin_func apart from variable names
 {
-  test("kprobe:f { @x = sp }",
+  test("kprobe:f { @x = reg(\"ip\") }",
 
 R"EXPECTED(; Function Attrs: nounwind
 declare i64 @llvm.bpf.pseudo(i64, i64) #0
@@ -662,12 +662,12 @@ define i64 @"kprobe:f"(i8*) local_unnamed_addr section "s_kprobe:f" {
 entry:
   %"@x_val" = alloca i64, align 8
   %"@x_key" = alloca i64, align 8
-  %sp = alloca i64, align 8
-  %1 = bitcast i64* %sp to i8*
+  %reg_ip = alloca i64, align 8
+  %1 = bitcast i64* %reg_ip to i8*
   call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
-  %2 = getelementptr i8, i8* %0, i64 152
-  %probe_read = call i64 inttoptr (i64 4 to i64 (i8*, i64, i8*)*)(i64* nonnull %sp, i64 8, i8* %2)
-  %3 = load i64, i64* %sp, align 8
+  %2 = getelementptr i8, i8* %0, i64 128
+  %probe_read = call i64 inttoptr (i64 4 to i64 (i8*, i64, i8*)*)(i64* nonnull %reg_ip, i64 8, i8* %2)
+  %3 = load i64, i64* %reg_ip, align 8
   call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
   %4 = bitcast i64* %"@x_key" to i8*
   call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %4)

tests/semantic_analyser.cpp

@@ -59,7 +59,6 @@ TEST(semantic_analyser, builtin_variables)
   test("kprobe:f { arg0 }", 0);
   test("kprobe:f { retval }", 0);
   test("kprobe:f { func }", 0);
-  test("kprobe:f { sp }", 0);
   test("kprobe:f { fake }", 1);
 }
 
@@ -72,6 +71,7 @@ TEST(semantic_analyser, builtin_functions)
   test("kprobe:f { printf(\"hello\\n\") }", 0);
   test("kprobe:f { sym(0xffff) }", 0);
   test("kprobe:f { usym(0xffff) }", 0);
+  test("kprobe:f { reg(\"ip\") }", 0);
   test("kprobe:f { fake() }", 1);
 }
 
@@ -150,6 +150,15 @@ TEST(semantic_analyser, call_usym)
   test("kprobe:f { usym(\"hello\"); }", 10);
 }
 
+TEST(semantic_analyser, call_reg)
+{
+  test("kprobe:f { reg(\"ip\"); }", 0);
+  test("kprobe:f { @x = reg(\"ip\"); }", 0);
+  test("kprobe:f { reg(\"blah\"); }", 1);
+  test("kprobe:f { reg(); }", 1);
+  test("kprobe:f { reg(123); }", 1);
+}
+
 TEST(semantic_analyser, map_reassignment)
 {
   test("kprobe:f { @x = 1; @x = 2; }", 0);
@@ -195,6 +204,7 @@ TEST(semantic_analyser, printf)
 {
   test("kprobe:f { printf(\"hi\") }", 0);
   test("kprobe:f { printf(1234) }", 1);
+  test("kprobe:f { printf() }", 1);
   test("kprobe:f { $fmt = \"mystring\"; printf($fmt) }", 1);
 }