add macro definition support

Add macro definition support as a bpftrace "preprocessor"-like feature. This is implemented by running our Bison parser stage to collect preprocessor stuff, structs, enums, etc., then we run our clang parser stage, and then we run our parser stage again, replacing all occurrences of macros with their literal values. This will work for multi-level macro definitions, because Bison will keep parsing and replacing until there's no macro left in its cursor. Macro expressions (`#define MACRO()`) are not supported, but we can consider adding support for those in the future. The downside (maybe it is not a downside) is that macros defining things that bpftrace can't parse will throw. Fixes: bpftrace#153
caringi · May 15, 2019 · 8587411 · 8587411
1 parent 5143209
commit 8587411
Show file tree

Hide file tree

Showing 16 changed files with 158 additions and 25 deletions.
diff --git a/src/bpftrace.h b/src/bpftrace.h
@@ -82,6 +82,7 @@ class BPFtrace
 
   std::map<std::string, std::unique_ptr<IMap>> maps_;
   std::map<std::string, Struct> structs_;
+  std::map<std::string, std::string> macros_;
   std::vector<std::tuple<std::string, std::vector<Field>>> printf_args_;
   std::vector<std::tuple<std::string, std::vector<Field>>> system_args_;
   std::vector<std::string> time_args_;

diff --git a/src/clang_parser.cpp b/src/clang_parser.cpp
@@ -53,6 +53,39 @@ static int get_indirect_field_offset(CXCursor c)
   return offset;
 }
 
+// NOTE(mmarchini): as suggested in http://clang-developers.42468.n3.nabble.com/Extracting-macro-information-using-libclang-the-C-Interface-to-Clang-td4042648.html#message4042666
+static bool translateMacro(CXCursor cursor, std::string &name, std::string &value)
+{
+  CXToken* tokens    = nullptr;
+  unsigned numTokens = 0;
+  CXTranslationUnit transUnit = clang_Cursor_getTranslationUnit(cursor);
+  CXSourceRange srcRange  = clang_getCursorExtent(cursor);
+  clang_tokenize(transUnit, srcRange, &tokens, &numTokens);
+  for (unsigned n=0; n<numTokens; n++)
+  {
+    auto tokenText = clang_getTokenSpelling(transUnit, tokens[n]);
+    if (n == 0)
+    {
+      value.clear();
+      name = clang_getCString(tokenText);
+      if (name[0] == '_')
+        break;
+    }
+    else
+    {
+      CXTokenKind tokenKind = clang_getTokenKind(tokens[n]);
+      if (tokenKind != CXToken_Comment)
+      {
+        const char* text = clang_getCString(tokenText);
+        if (text)
+          value += text;
+      }
+    }
+  }
+  clang_disposeTokens(transUnit, tokens, numTokens);
+  return value.length() != 0;
+}
+
 static SizedType get_sized_type(CXType clang_type)
 {
   auto size = clang_Type_getSizeOf(clang_type);
@@ -245,7 +278,7 @@ void ClangParser::parse(ast::Program *program, BPFtrace &bpftrace)
       "definitions.h",
       &args[0], args.size(),
       unsaved_files, sizeof(unsaved_files)/sizeof(CXUnsavedFile),
-      CXTranslationUnit_None,
+      CXTranslationUnit_DetailedPreprocessingRecord,
       &translation_unit);
   if (error)
   {
@@ -260,6 +293,17 @@ void ClangParser::parse(ast::Program *program, BPFtrace &bpftrace)
       [](CXCursor c, CXCursor parent, CXClientData client_data)
       {
 
+        if (clang_getCursorKind(c) == CXCursor_MacroDefinition)
+        {
+          std::string macro_name;
+          std::string macro_value;
+          if (translateMacro(c, macro_name, macro_value)) {
+            auto &macros = static_cast<BPFtrace*>(client_data)->macros_;
+            macros[macro_name] = macro_value;
+          }
+          return CXChildVisit_Recurse;
+        }
+
         if (clang_getCursorKind(parent) != CXCursor_StructDecl &&
             clang_getCursorKind(parent) != CXCursor_UnionDecl)
           return CXChildVisit_Recurse;

diff --git a/src/driver.cpp b/src/driver.cpp
@@ -10,7 +10,7 @@ extern int yylex_destroy (yyscan_t yyscanner);
 
 namespace bpftrace {
 
-Driver::Driver()
+Driver::Driver(BPFtrace &bpftrace) : bpftrace_(bpftrace)
 {
   yylex_init(&scanner_);
   parser_ = std::make_unique<Parser>(*this, scanner_);

diff --git a/src/driver.h b/src/driver.h
@@ -3,6 +3,7 @@
 #include <memory>
 
 #include "ast.h"
+#include "bpftrace.h"
 #include "parser.tab.hh"
 
 typedef void* yyscan_t;
@@ -14,7 +15,7 @@ namespace bpftrace {
 class Driver
 {
 public:
-  Driver();
+  explicit Driver(BPFtrace &bpftrace);
   ~Driver();
 
   int parse_stdin();
@@ -25,6 +26,7 @@ class Driver
 
   ast::Program *root_;
 
+  BPFtrace &bpftrace_;
 private:
   std::unique_ptr<Parser> parser_;
   yyscan_t scanner_;

diff --git a/src/lexer.l b/src/lexer.l
@@ -144,7 +144,16 @@ struct|union            { BEGIN(STRUCT); buffer = ""; struct_type = std::string(
                           buffer += std::string(yytext);
                         }
 
-{ident}                 { return Parser::make_IDENT(yytext, loc); }
+{ident}                 {
+                          if (driver.bpftrace_.macros_.count(yytext) != 0) {
+                            const char *s = driver.bpftrace_.macros_[yytext].c_str();
+                            int z;
+                            for (z=strlen(s) - 1; z >= 0; z--)
+                              unput(s[z]);
+                          } else {
+                            return Parser::make_IDENT(yytext, loc);
+                          }
+                        }
 .                       { driver.error(loc, std::string("invalid character '") +
                                             std::string(yytext) + std::string("'")); }
 

diff --git a/src/main.cpp b/src/main.cpp
@@ -73,7 +73,6 @@ bool is_root()
 int main(int argc, char *argv[])
 {
   int err;
-  Driver driver;
   char *pid_str = nullptr;
   char *cmd_str = nullptr;
   bool listing = false;
@@ -84,7 +83,7 @@ int main(int argc, char *argv[])
     return 0;
   }
 
-  std::string script, search;
+  std::string script, search, file_name;
   int c;
   while ((c = getopt(argc, argv, "dB:e:hlp:vc:")) != -1)
   {
@@ -153,6 +152,7 @@ int main(int argc, char *argv[])
   }
 
   BPFtrace bpftrace;
+  Driver driver(bpftrace);
 
   // PID is currently only used for USDT probes that need enabling. Future work:
   // - make PID a filter for all probe types: pass to perf_event_open(), etc.
@@ -181,8 +181,8 @@ int main(int argc, char *argv[])
   if (script.empty())
   {
     // Script file
-    char *file_name = argv[optind];
-    if (!file_name)
+    file_name = std::string(argv[optind]);
+    if (file_name.empty())
     {
       std::cerr << "USAGE: filename or -e 'program' required." << std::endl;
       return 1;
@@ -274,6 +274,18 @@ int main(int argc, char *argv[])
   ClangParser clang;
   clang.parse(driver.root_, bpftrace);
 
+  if (script.empty())
+  {
+    err = driver.parse_file(file_name);
+  }
+  else
+  {
+    err = driver.parse_str(script);
+  }
+
+  if (err)
+    return err;
+
   ast::SemanticAnalyser semantics(driver.root_, bpftrace);
   err = semantics.analyse();
   if (err)

diff --git a/tests/clang_parser.cpp b/tests/clang_parser.cpp
@@ -10,7 +10,7 @@ namespace clang_parser {
 void parse(const std::string &input, BPFtrace &bpftrace)
 {
   auto extended_input = input + "kprobe:sys_read { 1 }";
-  Driver driver;
+  Driver driver(bpftrace);
   ASSERT_EQ(driver.parse_str(extended_input), 0);
 
   ClangParser clang;
@@ -272,6 +272,18 @@ TEST(clang_parser, builtin_headers)
   EXPECT_EQ(structs["Foo"].fields["z"].offset, 16);
 }
 
+TEST(clang_parser, macro_preprocessor)
+{
+  // size_t is definied in stddef.h
+  BPFtrace bpftrace;
+  parse("#define FOO size_t\n k:f { 0 }", bpftrace);
+
+  auto &macros = bpftrace.macros_;
+
+  ASSERT_EQ(macros.count("FOO"), 1U);
+  ASSERT_EQ(macros["FOO"], "size_t");
+}
+
 } // namespace clang_parser
 } // namespace test
 } // namespace bpftrace
diff --git a/tests/codegen.cpp b/tests/codegen.cpp
@@ -56,6 +56,7 @@
 #include "codegen/int_propagation.cpp"
 #include "codegen/logical_and.cpp"
 #include "codegen/logical_or.cpp"
+#include "codegen/macro_definition.cpp"
 #include "codegen/map_assign_int.cpp"
 #include "codegen/map_assign_string.cpp"
 #include "codegen/map_key_int.cpp"

diff --git a/tests/codegen/call_kstack.cpp b/tests/codegen/call_kstack.cpp
@@ -62,7 +62,7 @@ attributes #1 = { argmemonly nounwind }
 TEST(codegen, call_kstack_mapids)
 {
   BPFtrace bpftrace;
-  Driver driver;
+  Driver driver(bpftrace);
   FakeMap::next_mapfd_ = 1;
 
   ASSERT_EQ(driver.parse_str("kprobe:f { @x = kstack(5); @y = kstack(6); @z = kstack(6) }"), 0);
@@ -90,7 +90,7 @@ TEST(codegen, call_kstack_mapids)
 TEST(codegen, call_kstack_modes_mapids)
 {
   BPFtrace bpftrace;
-  Driver driver;
+  Driver driver(bpftrace);
   FakeMap::next_mapfd_ = 1;
 
   ASSERT_EQ(driver.parse_str("kprobe:f { @x = kstack(perf); @y = kstack(bpftrace); @z = kstack() }"), 0);

diff --git a/tests/codegen/call_ustack.cpp b/tests/codegen/call_ustack.cpp
@@ -68,7 +68,7 @@ attributes #1 = { argmemonly nounwind }
 TEST(codegen, call_ustack_mapids)
 {
   BPFtrace bpftrace;
-  Driver driver;
+  Driver driver(bpftrace);
   FakeMap::next_mapfd_ = 1;
 
   ASSERT_EQ(driver.parse_str("kprobe:f { @x = ustack(5); @y = ustack(6); @z = ustack(6) }"), 0);
@@ -96,7 +96,7 @@ TEST(codegen, call_ustack_mapids)
 TEST(codegen, call_ustack_modes_mapids)
 {
   BPFtrace bpftrace;
-  Driver driver;
+  Driver driver(bpftrace);
   FakeMap::next_mapfd_ = 1;
 
   ASSERT_EQ(driver.parse_str("kprobe:f { @x = ustack(perf); @y = ustack(bpftrace); @z = ustack() }"), 0);

diff --git a/tests/codegen/common.h b/tests/codegen/common.h
@@ -24,14 +24,16 @@ target triple = "bpf-pc-linux"
 static void test(const std::string &input, const std::string expected_output)
 {
   BPFtrace bpftrace;
-  Driver driver;
+  Driver driver(bpftrace);
   FakeMap::next_mapfd_ = 1;
 
   ASSERT_EQ(driver.parse_str(input), 0);
 
   ClangParser clang;
   clang.parse(driver.root_, bpftrace);
 
+  ASSERT_EQ(driver.parse_str(input), 0);
+
   ast::SemanticAnalyser semantics(driver.root_, bpftrace);
   ASSERT_EQ(semantics.analyse(), 0);
   ASSERT_EQ(semantics.create_maps(true), 0);

diff --git a/tests/codegen/general.cpp b/tests/codegen/general.cpp
@@ -22,7 +22,7 @@ class MockBPFtrace : public BPFtrace {
 TEST(codegen, populate_sections)
 {
   BPFtrace bpftrace;
-  Driver driver;
+  Driver driver(bpftrace);
 
   ASSERT_EQ(driver.parse_str("kprobe:foo { 1 } kprobe:bar { 1 }"), 0);
   ast::SemanticAnalyser semantics(driver.root_, bpftrace);
@@ -40,7 +40,7 @@ TEST(codegen, populate_sections)
 TEST(codegen, printf_offsets)
 {
   BPFtrace bpftrace;
-  Driver driver;
+  Driver driver(bpftrace);
 
   // TODO (mmarchini): also test printf with a string argument
   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);
@@ -77,7 +77,7 @@ TEST(codegen, probe_count)
   MockBPFtrace bpftrace;
   EXPECT_CALL(bpftrace, add_probe(_)).Times(2);
 
-  Driver driver;
+  Driver driver(bpftrace);
 
   ASSERT_EQ(driver.parse_str("kprobe:f { 1; } kprobe:d { 1; }"), 0);
   ast::SemanticAnalyser semantics(driver.root_, bpftrace);

diff --git a/tests/codegen/macro_definition.cpp b/tests/codegen/macro_definition.cpp
@@ -0,0 +1,45 @@
+#include "common.h"
+
+namespace bpftrace {
+namespace test {
+namespace codegen {
+
+TEST(codegen, macro_definition)
+{
+  test("#define FOO 100\nk:f { @ = FOO }",
+
+R"EXPECTED(; Function Attrs: nounwind
+declare i64 @llvm.bpf.pseudo(i64, i64) #0
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #1
+
+define i64 @"kprobe:f"(i8* nocapture readnone) local_unnamed_addr section "s_kprobe:f_1" {
+entry:
+  %"@_val" = alloca i64, align 8
+  %"@_key" = alloca i64, align 8
+  %1 = bitcast i64* %"@_key" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %1)
+  store i64 0, i64* %"@_key", align 8
+  %2 = bitcast i64* %"@_val" to i8*
+  call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %2)
+  store i64 100, i64* %"@_val", align 8
+  %pseudo = tail call i64 @llvm.bpf.pseudo(i64 1, i64 1)
+  %update_elem = call i64 inttoptr (i64 2 to i64 (i8*, i8*, i8*, i64)*)(i64 %pseudo, i64* nonnull %"@_key", i64* nonnull %"@_val", i64 0)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %1)
+  call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %2)
+  ret i64 0
+}
+
+; Function Attrs: argmemonly nounwind
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #1
+
+attributes #0 = { nounwind }
+attributes #1 = { argmemonly nounwind }
+)EXPECTED");
+}
+
+} // namespace codegen
+} // namespace test
+} // namespace bpftrace
+
diff --git a/tests/parser.cpp b/tests/parser.cpp
@@ -12,7 +12,8 @@ using Printer = ast::Printer;
 
 void test(const std::string &input, const std::string &output)
 {
-  Driver driver;
+  BPFtrace bpftrace;
+  Driver driver(bpftrace);
   ASSERT_EQ(driver.parse_str(input), 0);
 
   std::ostringstream out;
@@ -870,7 +871,8 @@ TEST(Parser, cstruct_nested)
 TEST(Parser, unterminated_string)
 {
   // Make sure parser doesn't get stuck in an infinite loop
-  Driver driver;
+  BPFtrace bpftrace;
+  Driver driver(bpftrace);
   EXPECT_EQ(driver.parse_str("kprobe:f { \"asdf }"), 1);
 }
 

diff --git a/tests/probe.cpp b/tests/probe.cpp
@@ -22,7 +22,7 @@ using bpftrace::ast::Probe;
 void gen_bytecode(const std::string &input, std::stringstream &out)
 {
 	BPFtrace bpftrace;
-	Driver driver;
+	Driver driver(bpftrace);
 	FakeMap::next_mapfd_ = 1;
 
 	ASSERT_EQ(driver.parse_str(input), 0);