Create Clang tool to find struct definitions in Nvidia source files.

PiperOrigin-RevId: 655338427

WORKSPACE

@@ -196,6 +196,56 @@ cc_crosstool(
 
 register_toolchains("//:cc_toolchain_k8", "//:cc_toolchain_aarch64")
 
+# Load LLVM dependencies.
+LLVM_COMMIT = "926f85db98aae66ab8f57b9981f47ddddb868c51"
+LLVM_SHA256 = "c78c94b2a03b2cf6ef1ba035c31a6f1b0bb7913da8af5aa8d5c2061f6499d589"
+
+http_archive(
+    name = "llvm-raw",
+    build_file_content = "# empty",
+    sha256 = LLVM_SHA256,
+    strip_prefix = "llvm-project-" + LLVM_COMMIT,
+    urls = ["https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT)],
+)
+
+load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
+
+llvm_configure(name = "llvm-project")
+
+load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
+
+maybe(
+    http_archive,
+    name = "llvm_zlib",
+    build_file = "@llvm-raw//utils/bazel/third_party_build:zlib-ng.BUILD",
+    sha256 = "e36bb346c00472a1f9ff2a0a4643e590a254be6379da7cddd9daeb9a7f296731",
+    strip_prefix = "zlib-ng-2.0.7",
+    urls = [
+        "https://github.com/zlib-ng/zlib-ng/archive/refs/tags/2.0.7.zip",
+    ],
+)
+
+maybe(
+    http_archive,
+    name = "llvm_zstd",
+    build_file = "@llvm-raw//utils/bazel/third_party_build:zstd.BUILD",
+    sha256 = "7c42d56fac126929a6a85dbc73ff1db2411d04f104fae9bdea51305663a83fd0",
+    strip_prefix = "zstd-1.5.2",
+    urls = [
+        "https://github.com/facebook/zstd/releases/download/v1.5.2/zstd-1.5.2.tar.gz",
+    ],
+)
+
+# Load other C++ dependencies.
+http_archive(
+  name = "nlohmann_json",
+  sha256 = "ba6e7817353793d13e5214ed819ea5b0defc0ffb2a348f4e34b10ac6f1c50154",
+  strip_prefix = "json-960b763ecd144f156d05ec61f577b04107290137",
+  urls = [
+    "https://github.com/nlohmann/json/archive/960b763ecd144f156d05ec61f577b04107290137.tar.gz"
+  ]
+)
+
 http_archive(
     name = "com_google_protobuf",
     sha256 = "c968404387c9cccd18676c6e1d83a1dcc39d162a7f468dace4b243c274de1f02",

tools/nvidia_driver_differ/BUILD

@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "build_test", "cc_binary")
+load("//tools:defs.bzl", "cc_binary", "go_test")
 
 package(
     default_applicable_licenses = ["//:license"],
@@ -11,9 +11,25 @@ cc_binary(
         "driver_ast_parser.cc",
         "driver_ast_parser.h",
     ],
+    deps = [
+        "@llvm-project//clang:ast",
+        "@llvm-project//clang:ast_matchers",
+        "@llvm-project//clang:tooling",
+        "@llvm-project//llvm:Support",
+        "@nlohmann_json//:json",
+    ],
 )
 
-build_test(
+go_test(
     name = "driver_ast_parser_test",
-    targets = [":driver_ast_parser"],
+    srcs = ["driver_ast_parser_test.go"],
+    data = [
+        "test_struct.cc",
+        ":driver_ast_parser",
+    ],
+    deps = [
+        "//pkg/test/testutil",
+        "@com_github_google_go_cmp//cmp:go_default_library",
+        "@com_github_google_go_cmp//cmp/cmpopts:go_default_library",
+    ],
 )

tools/nvidia_driver_differ/driver_ast_parser.cc

@@ -13,4 +13,149 @@
 // limitations under the License.
 
 #include "tools/nvidia_driver_differ/driver_ast_parser.h"
-int main(int argc, const char **argv) {}
+
+#include <fstream>
+#include <iostream>
+#include <string>
+
+#include "nlohmann/json.hpp"
+#include "clang/include/clang/AST/Decl.h"
+#include "clang/include/clang/ASTMatchers/ASTMatchFinder.h"
+#include "clang/include/clang/ASTMatchers/ASTMatchers.h"
+#include "clang/include/clang/Tooling/CommonOptionsParser.h"
+#include "clang/include/clang/Tooling/Tooling.h"
+#include "llvm/include/llvm/Support/CommandLine.h"
+#include "llvm/include/llvm/Support/raw_ostream.h"
+
+using clang::ast_matchers::allOf;
+using clang::ast_matchers::elaboratedType;
+using clang::ast_matchers::hasDeclaration;
+using clang::ast_matchers::hasName;
+using clang::ast_matchers::hasType;
+using clang::ast_matchers::recordDecl;
+using clang::ast_matchers::typedefDecl;
+
+using clang::ast_matchers::MatchFinder;
+
+using json = nlohmann::json;
+
+struct DriverStructReporter : public MatchFinder::MatchCallback {
+  json StructDefinitions;
+
+  auto get_struct_matcher(std::string struct_name) {
+    // Nvidia's driver typedefs all their struct. We search for the
+    // typedef declaration, and go from there to find the struct definition.
+    return typedefDecl(
+               allOf(hasName(struct_name),
+                     // Match and bind to the struct declaration.
+                     hasType(
+                         // Need to specify elaboratedType, otherwise hasType
+                         // will complain that the type is ambiguous.
+                         elaboratedType(hasDeclaration(
+                             recordDecl().bind("struct_decl"))))))
+        .bind("typedef_decl");
+  }
+
+  void run(const MatchFinder::MatchResult &result) override {
+    const auto *typedef_decl =
+        result.Nodes.getNodeAs<clang::TypedefDecl>("typedef_decl");
+    if (typedef_decl == nullptr) {
+      std::cerr << "Unable to find typedef decl\n";
+      return;
+    }
+
+    const auto *struct_decl =
+        result.Nodes.getNodeAs<clang::RecordDecl>("struct_decl");
+    if (struct_decl == nullptr) {
+      std::cerr << "Unable to find struct decl for "
+                << typedef_decl->getNameAsString() << "\n";
+      return;
+    }
+
+    // Add struct definition to json.
+    // TODO(b/347796680): Consider improvements:
+    //  Store alignment attributes as well?
+    //  Make recursive? Relevant for recursive structs not defined in nvproxy,
+    //    or unnamed structs/unions.
+    //  Handle anonymous names?
+    json fields = json::array();
+    for (const auto *field : struct_decl->fields()) {
+      fields.push_back(
+          json::object({{"name", field->getNameAsString()},
+                        {"type", field->getType().getAsString()}}));
+    }
+
+    std::string name = typedef_decl->getNameAsString();
+    std::string source = typedef_decl->getLocation().printToString(
+        result.Context->getSourceManager());
+
+    StructDefinitions[name] =
+        json::object({{"fields", fields}, {"source", source}});
+  }
+};
+
+static llvm::cl::OptionCategory DriverASTParserCategory("Driver AST Parser");
+
+static llvm::cl::extrahelp CommonHelp(
+    clang::tooling::CommonOptionsParser::HelpMessage);
+static llvm::cl::extrahelp MoreHelp(ToolHelpDescription);
+
+static llvm::cl::opt<std::string> StructNames(
+    "structs",
+    llvm::cl::desc(
+        "Path to the input file containing the struct names to parse."),
+    llvm::cl::cat(DriverASTParserCategory), llvm::cl::Required);
+
+static llvm::cl::opt<std::string> OutputFile(
+    "output", "o",
+    llvm::cl::desc("Path to the output file for the parsed structs. "
+                   "By default, will print to stdout."),
+    llvm::cl::cat(DriverASTParserCategory));
+
+int main(int argc, const char **argv) {
+  auto ExpectedParser = clang::tooling::CommonOptionsParser::create(
+      argc, argv, DriverASTParserCategory);
+  if (!ExpectedParser) {
+    // Fail gracefully for unsupported options.
+    llvm::errs() << ExpectedParser.takeError();
+    return 1;
+  }
+
+  clang::tooling::CommonOptionsParser &OptionsParser = ExpectedParser.get();
+  clang::tooling::ClangTool Tool(OptionsParser.getCompilations(),
+                                 OptionsParser.getSourcePathList());
+
+  DriverStructReporter reporter;
+  MatchFinder finder;
+
+  // Read from StructNames file.
+  std::ifstream StructNamesIS(StructNames);
+  if (!StructNamesIS) {
+    std::cerr << "Unable to open struct names file: " << StructNames << "\n";
+    return 1;
+  }
+  json StructNamesJSON;
+  StructNamesIS >> StructNamesJSON;
+  for (json::iterator it = StructNamesJSON["structs"].begin();
+       it != StructNamesJSON["structs"].end(); ++it) {
+    finder.addMatcher(reporter.get_struct_matcher(*it), &reporter);
+  }
+
+  // Run tool
+  int ret = Tool.run(clang::tooling::newFrontendActionFactory(&finder).get());
+
+  // Print output.
+  json output = json::object({{"structs", reporter.StructDefinitions}});
+  if (OutputFile.empty()) {
+    std::cout << output.dump() << "\n";
+  } else {
+    std::ofstream OutputFileOS(OutputFile);
+    if (!OutputFileOS) {
+      std::cerr << "Unable to open output file: " << OutputFile << "\n";
+      return 1;
+    }
+    OutputFileOS << output.dump() << "\n";
+  }
+
+  return ret;
+}

tools/nvidia_driver_differ/driver_ast_parser.h

@@ -14,4 +14,34 @@
 
 #ifndef TOOLS_NVIDIA_DRIVER_DIFFER_DRIVER_AST_PARSER_H_
 #define TOOLS_NVIDIA_DRIVER_DIFFER_DRIVER_AST_PARSER_H_
+
+const char ToolHelpDescription[] =
+    R"a(This tool parses a given C++ source file and outputs the struct definitions
+for a list of provided struct names. To parse structs defined in multiple files,
+it is easier to create a C++ file that includes all the files to be parsed. You
+will also need a compile_commands.json file that contains a compile command with
+the relevant include directories.
+
+The struct names should be specified in a JSON file containing a JSON object,
+which has a "structs" key that maps to a list of strings. The tool will search
+for the struct definition in the given source files, and output the struct
+definition to the specified output file.
+
+This output file will contain a JSON object with a "structs" field mapping each
+struct name to its struct definition. Each struct definition will be a JSON
+array of fields, where each field is a JSON object with a "name" and a "type"
+key. The fields will be ordered in the same order as they appear in the struct
+
+Example usage:
+    driver_ast_parser --structs=structs.json -o=output.json driver_source_files.h
+
+The structs.json file should contain an array of struct names to parse:
+    {
+        "structs": [
+            "TestStruct",
+            "TestStruct2"
+        ]
+    }
+)a";
+
 #endif  // TOOLS_NVIDIA_DRIVER_DIFFER_DRIVER_AST_PARSER_H_