Add ast_canopy macro define extraction

Author: isVoid

ast_canopy/ast_canopy/api.py

@@ -57,6 +57,7 @@ class Declarations:
     class_template_specializations: list[ClassTemplateSpecialization]
     typedefs: list[bindings.Typedef]
     enums: list[bindings.Enum]
+    macro_defines: dict[str, str]
 
 
 def paths_to_include_flags(paths: list[str]) -> list[str]:
@@ -543,6 +544,7 @@ def parse_declarations_from_source(
         class_template_specializations,
         decls.typedefs,
         decls.enums,
+        decls.macro_defines,
     )
 
 

ast_canopy/ast_canopy/pylibastcanopy.cpp

@@ -303,7 +303,8 @@ PYBIND11_MODULE(pylibastcanopy, m) {
       .def_readwrite("class_template_specializations",
                      &Declarations::class_template_specializations)
       .def_readwrite("typedefs", &Declarations::typedefs)
-      .def_readwrite("enums", &Declarations::enums);
+      .def_readwrite("enums", &Declarations::enums)
+      .def_readwrite("macro_defines", &Declarations::macro_defines);
 
   m.def("parse_declarations_from_command_line",
         &parse_declarations_from_command_line,

ast_canopy/ast_canopy/pylibastcanopy.pyi

@@ -30,6 +30,7 @@ class Declarations:
     enums: list[Enum]
     function_templates: list[FunctionTemplate]
     functions: list[Function]
+    macro_defines: dict[str, str]
     records: list[Record]
     typedefs: list[Typedef]
     def __init__(self, *args, **kwargs) -> None: ...

ast_canopy/cpp/include/ast_canopy/ast_canopy.hpp

@@ -270,6 +270,7 @@ struct Declarations {
   std::vector<ClassTemplateSpecialization> class_template_specializations;
   std::vector<Typedef> typedefs;
   std::vector<Enum> enums;
+  std::unordered_map<std::string, std::string> macro_defines;
 };
 
 Declarations

ast_canopy/cpp/src/ast_canopy.cpp

@@ -9,7 +9,10 @@
 #include <clang/Frontend/ASTUnit.h>
 #include <clang/Frontend/CompilerInstance.h>
 #include <clang/Frontend/TextDiagnosticPrinter.h>
+#include <clang/Lex/MacroInfo.h>
+#include <clang/Lex/Preprocessor.h>
 
+#include <algorithm>
 #include <filesystem>
 #include <utility>
 
@@ -66,6 +69,69 @@ class AstCanopyDiagnosticsConsumer : public DiagnosticConsumer {
   }
 };
 
+bool filename_is_retained(const std::string &file_name,
+                          const std::vector<std::string> &files_to_retain) {
+  return std::any_of(files_to_retain.begin(), files_to_retain.end(),
+                     [&file_name](const std::string &file_to_retain) {
+                       return file_name == file_to_retain;
+                     });
+}
+
+std::string replacement_text_from_macro_info(const MacroInfo &macro_info,
+                                             const Preprocessor &preprocessor) {
+  std::string replacement_text;
+  for (const Token &token : macro_info.tokens()) {
+    bool invalid = false;
+    std::string spelling = preprocessor.getSpelling(token, &invalid);
+    if (invalid) {
+      continue;
+    }
+
+    if (!replacement_text.empty()) {
+      replacement_text += " ";
+    }
+    replacement_text += spelling;
+  }
+  return replacement_text;
+}
+
+void collect_macro_defines_from_ast(
+    ASTUnit *ast, const std::vector<std::string> &files_to_retain,
+    Declarations *decls) {
+  decls->macro_defines.clear();
+
+  Preprocessor &preprocessor = ast->getPreprocessor();
+  const SourceManager &source_manager = ast->getSourceManager();
+
+  for (const auto &entry : preprocessor.getIdentifierTable()) {
+    const IdentifierInfo *identifier_info = entry.second;
+    if (!identifier_info) {
+      continue;
+    }
+
+    MacroDefinition macro_definition =
+        preprocessor.getMacroDefinition(identifier_info);
+    const MacroInfo *macro_info = macro_definition.getMacroInfo();
+    if (!macro_info || macro_info->isFunctionLike()) {
+      continue;
+    }
+
+    SourceLocation spelling_location =
+        source_manager.getSpellingLoc(macro_info->getDefinitionLoc());
+    if (!spelling_location.isValid()) {
+      continue;
+    }
+
+    std::string file_name = source_manager.getFilename(spelling_location).str();
+    if (!filename_is_retained(file_name, files_to_retain)) {
+      continue;
+    }
+
+    decls->macro_defines[identifier_info->getName().str()] =
+        replacement_text_from_macro_info(*macro_info, preprocessor);
+  }
+}
+
 /**
  * @brief Return the source filename of the declaration.
  */
@@ -215,6 +281,7 @@ parse_declarations_from_command_line(std::vector<std::string> options,
                     &ctsd_callback);
 
   finder.matchAST(ast->getASTContext());
+  detail::collect_macro_defines_from_ast(ast.get(), files_to_retain, &decls);
 
 #ifndef NDEBUG
   std::cout << "Records: " << decls.records.size() << std::endl;

ast_canopy/tests/data/sample_macro_defines.cu

@@ -0,0 +1,7 @@
+#include "sample_macro_defines_include.cuh"
+
+#define foo 123
+#define FLAG
+#define MAKE_VALUE(x) ((x) + 1)
+
+__device__ int use_macro_define() { return foo; }

ast_canopy/tests/data/sample_macro_defines_include.cuh

@@ -0,0 +1 @@
+#define INCLUDED_VALUE 456

ast_canopy/tests/test_parse_macro.py

@@ -32,3 +32,32 @@ def test_macro_expansions(sample_macro_source):
     assert forty_two_int.return_type.name == "int"
     assert forty_two_float.return_type.name == "float"
     assert forty_two_double.return_type.name == "double"
+
+
+def test_macro_defines(data_folder):
+    srcstr = str(data_folder / "sample_macro_defines.cu")
+
+    decls = parse_declarations_from_source(
+        srcstr,
+        [srcstr],
+        "sm_80",
+    )
+
+    assert decls.macro_defines["foo"] == "123"
+    assert decls.macro_defines["FLAG"] == ""
+    assert "MAKE_VALUE" not in decls.macro_defines
+    assert "INCLUDED_VALUE" not in decls.macro_defines
+
+
+def test_macro_defines_retained_include(data_folder):
+    srcstr = str(data_folder / "sample_macro_defines.cu")
+    include = str(data_folder / "sample_macro_defines_include.cuh")
+
+    decls = parse_declarations_from_source(
+        srcstr,
+        [srcstr, include],
+        "sm_80",
+    )
+
+    assert decls.macro_defines["foo"] == "123"
+    assert decls.macro_defines["INCLUDED_VALUE"] == "456"