Add helper code to set up and run the parser on a given driver version.

PiperOrigin-RevId: 657727841

tools/nvidia_driver_differ/parser/BUILD

@@ -8,10 +8,15 @@ package(
 go_library(
     name = "parser",
     srcs = [
+        "auxiliary_files.go",
         "clang_config.go",
         "json_definitions.go",
+        "runner.go",
         "sources.go",
     ],
     visibility = ["//tools/nvidia_driver_differ:__subpackages__"],
-    deps = ["@com_github_google_go_cmp//cmp:go_default_library"],
+    deps = [
+        "//pkg/sentry/devices/nvproxy",
+        "@com_github_google_go_cmp//cmp:go_default_library",
+    ],
 )

tools/nvidia_driver_differ/parser/auxiliary_files.go

@@ -0,0 +1,92 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package parser
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path"
+
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy"
+)
+
+// GitRepoURL is the URL for the NVIDIA open-gpu-kernel-modules repo.
+const GitRepoURL = "https://github.com/NVIDIA/open-gpu-kernel-modules.git"
+
+// CloneDriverSource clones the given driver version into the given directory.
+func CloneDriverSource(dir string, version nvproxy.DriverVersion) (*DriverSourceDir, error) {
+	// git clone -b $VERSION --depth 1 https://github.com/NVIDIA/open-gpu-kernel-modules.git $PATH
+	args := []string{
+		"clone",
+		"-b",
+		version.String(),
+		"--depth",
+		"1",
+		GitRepoURL,
+		path.Join(dir, version.String()),
+	}
+	cmd := exec.Command("git", args...)
+	if out, err := cmd.CombinedOutput(); err != nil {
+		return nil, fmt.Errorf("failed to clone %s: %w\n%s", version, err, string(out))
+	}
+	return &DriverSourceDir{
+		ParentDirectory: dir,
+		Version:         version,
+	}, nil
+}
+
+// CreateIncludeFiles creates the necessary include files for the given driver version, and returns
+// the config options for the files.
+func CreateIncludeFiles(dir string, driverSource DriverSourceDir) ([]ClangASTConfig, error) {
+	// Create include file for non-uvm sources
+	nonUVMFile, err := os.CreateTemp(dir, "include_non_uvm_*.cc")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create temporary file: %w", err)
+	}
+	defer nonUVMFile.Close()
+
+	includeSources, err := driverSource.GetNonUVMSourcePaths()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get non-uvm include paths: %w", err)
+	}
+	if err := WriteIncludeFile(includeSources, nonUVMFile); err != nil {
+		return nil, fmt.Errorf("failed to write include file: %w", err)
+	}
+	configNonUVM := NewParserConfig(
+		dir,
+		nonUVMFile.Name(),
+		driverSource.GetNonUVMIncludePaths(),
+	)
+
+	// Create include file for uvm sources
+	UVMFile, err := os.CreateTemp(dir, "include_uvm_*.cc")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create temporary file: %w", err)
+	}
+	defer UVMFile.Close()
+
+	includeSources = driverSource.GetUVMSourcePaths()
+	if err := WriteIncludeFile(includeSources, UVMFile); err != nil {
+		return nil, fmt.Errorf("failed to write include file: %w", err)
+	}
+	configUVM := NewParserConfig(
+		dir,
+		UVMFile.Name(),
+		driverSource.GetUVMIncludePaths(),
+	)
+
+	return []ClangASTConfig{configNonUVM, configUVM}, nil
+}

tools/nvidia_driver_differ/parser/json_definitions.go

@@ -17,6 +17,7 @@ package parser
 
 import (
 	"fmt"
+	"maps"
 	"slices"
 	"strings"
 
@@ -34,6 +35,18 @@ type OutputJSON struct {
 	Aliases TypeAliases `json:"aliases"`
 }
 
+// Merge merges the struct definitions from b into this OutputJSON.
+func (a *OutputJSON) Merge(b OutputJSON) {
+	if a.Records == nil {
+		a.Records = make(RecordDefs)
+	}
+	if a.Aliases == nil {
+		a.Aliases = make(TypeAliases)
+	}
+	maps.Copy(a.Records, b.Records)
+	maps.Copy(a.Aliases, b.Aliases)
+}
+
 // RecordField represents a field in a record (struct or union).
 type RecordField struct {
 	Name string

tools/nvidia_driver_differ/parser/runner.go

@@ -0,0 +1,152 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package parser
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"os/exec"
+
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy"
+)
+
+// ParserFile is a wrapper around the driver_ast_parser binary.
+type ParserFile os.File
+
+// StructsFile is a wrapper around the structs list file.
+type StructsFile os.File
+
+// Runner is a helper for running the driver_ast_parser on a given set of structs.
+type Runner struct {
+	// Working directory for the runner.
+	dir         string
+	parserFile  *ParserFile
+	structsFile *StructsFile
+}
+
+// NewRunner creates a new Runner around a given parser file and a temporary working directory.
+func NewRunner(parserFile *ParserFile) (*Runner, error) {
+	// Create a temp directory for the runner.
+	dir, err := os.MkdirTemp(os.TempDir(), "run_differ_*")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create temporary directory: %w", err)
+	}
+
+	return &Runner{
+		dir:        dir,
+		parserFile: parserFile,
+	}, nil
+}
+
+// Cleanup removes the working directory for the runner.
+func (r *Runner) Cleanup() error {
+	return os.RemoveAll(r.dir)
+}
+
+// CreateStructsFile saves a list of structs for the runner to parse.
+func (r *Runner) CreateStructsFile(structs []nvproxy.DriverStructName) error {
+	inputJSON := InputJSON{
+		Structs: structs,
+	}
+
+	f, err := os.CreateTemp(r.dir, "structs_list_*.json")
+	if err != nil {
+		return fmt.Errorf("failed to create temporary structs list: %w", err)
+	}
+	defer f.Close()
+
+	if err := json.NewEncoder(f).Encode(inputJSON); err != nil {
+		return fmt.Errorf("failed to write structs list to file: %w", err)
+	}
+
+	r.structsFile = (*StructsFile)(f)
+	return nil
+}
+
+// parseSourceFile runs driver_ast_parser on sourceFile for the structs listed in structsFile,
+// and returns the parsed JSON output.
+func (r *Runner) parseSourceFile(sourcePath string) (*OutputJSON, error) {
+	if r.structsFile == nil {
+		return nil, fmt.Errorf("structs file not created")
+	}
+
+	// Run driver_ast_parser on the .cc file
+	cmd := exec.Command((*os.File)(r.parserFile).Name(), "--structs", (*os.File)(r.structsFile).Name(), sourcePath)
+	out, err := cmd.Output()
+	if err != nil {
+		return nil, fmt.Errorf("failed to run driver_ast_parser: %v", err)
+	}
+
+	// Unmarshal the output
+	var defs OutputJSON
+	if err := json.Unmarshal(out, &defs); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal output file: %w", err)
+	}
+
+	return &defs, nil
+}
+
+// runParserConfig runs the driver_ast_parser on the given config options and merges all the
+// JSON outputs into a single OutputJSON.
+func (r *Runner) runParserConfig(config []ClangASTConfig) (*OutputJSON, error) {
+	var allDefs *OutputJSON = nil
+	for _, config := range config {
+		defs, err := r.parseSourceFile(config.Filename)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse source file: %w", err)
+		}
+
+		if allDefs == nil {
+			allDefs = defs
+		} else {
+			allDefs.Merge(*defs)
+		}
+	}
+	return allDefs, nil
+}
+
+// ParseDriver checks out the git repo for the given version, and runs the driver_ast_parser on the
+// source code.
+func (r *Runner) ParseDriver(version nvproxy.DriverVersion) (*OutputJSON, error) {
+	// Create a temp directory to run the parser in.
+	// This is needed to set up compile_commands.json, since it needs to be named that exactly.
+	dir, err := os.MkdirTemp(r.dir, "run_differ_*")
+	if err != nil {
+		return nil, fmt.Errorf("failed to create temporary directory: %w", err)
+	}
+	defer os.RemoveAll(dir)
+
+	source, err := CloneDriverSource(dir, version)
+	if err != nil {
+		return nil, fmt.Errorf("failed to clone git repo: %w", err)
+	}
+
+	config, err := CreateIncludeFiles(dir, *source)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create include files: %w", err)
+	}
+
+	if err := CreateCompileCommandsFile(dir, config); err != nil {
+		return nil, fmt.Errorf("failed to create compile_commands.json: %w", err)
+	}
+
+	defs, err := r.runParserConfig(config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to run driver_ast_parser: %w", err)
+	}
+
+	return defs, nil
+}

tools/nvidia_driver_differ/parser/sources.go

@@ -19,17 +19,24 @@ import (
 	"fmt"
 	"io"
 	"path/filepath"
+
+	"gvisor.dev/gvisor/pkg/sentry/devices/nvproxy"
 )
 
 // DriverSourceDir represents a directory containing the source code for a given driver version.
 type DriverSourceDir struct {
-	Path    string
-	Version string
+	ParentDirectory string
+	Version         nvproxy.DriverVersion
+}
+
+// Name returns the name of the driver source directory.
+func (d DriverSourceDir) Name() string {
+	return d.Version.String()
 }
 
 // GlobDriverFiles returns all files in the given driver directory that match the given pattern.
 func (d *DriverSourceDir) GlobDriverFiles(pattern string) ([]string, error) {
-	files, err := filepath.Glob(fmt.Sprintf("%s/%s", d.Path, pattern))
+	files, err := filepath.Glob(fmt.Sprintf("%s/%s/%s", d.ParentDirectory, d.Name(), pattern))
 	if err != nil {
 		return nil, fmt.Errorf("failed to glob files: %w", err)
 	}
@@ -61,24 +68,24 @@ func (d *DriverSourceDir) GetNonUVMSourcePaths() ([]string, error) {
 // GetUVMSourcePaths returns the list of paths for uvm source files.
 func (d *DriverSourceDir) GetUVMSourcePaths() []string {
 	return []string{
-		fmt.Sprintf("%s/kernel-open/nvidia-uvm/uvm_ioctl.h", d.Path),
-		fmt.Sprintf("%s/kernel-open/nvidia-uvm/uvm_linux_ioctl.h", d.Path),
+		fmt.Sprintf("%s/kernel-open/nvidia-uvm/uvm_ioctl.h", d.Name()),
+		fmt.Sprintf("%s/kernel-open/nvidia-uvm/uvm_linux_ioctl.h", d.Name()),
 	}
 }
 
 // GetNonUVMIncludePaths returns the list of paths for non-uvm include files.
 func (d *DriverSourceDir) GetNonUVMIncludePaths() []string {
 	return []string{
-		fmt.Sprintf("%s/src/common/sdk/nvidia/inc", d.Path),
-		fmt.Sprintf("%s/src/common/shared/inc", d.Path),
-		fmt.Sprintf("%s/src/nvidia/arch/nvalloc/unix/include", d.Path),
+		fmt.Sprintf("%s/src/common/sdk/nvidia/inc", d.Name()),
+		fmt.Sprintf("%s/src/common/shared/inc", d.Name()),
+		fmt.Sprintf("%s/src/nvidia/arch/nvalloc/unix/include", d.Name()),
 	}
 }
 
 // GetUVMIncludePaths returns the list of paths for uvm include files.
 func (d *DriverSourceDir) GetUVMIncludePaths() []string {
 	return []string{
-		fmt.Sprintf("%s/kernel-open/common/inc", d.Path),
+		fmt.Sprintf("%s/kernel-open/common/inc", d.Name()),
 	}
 }