From b96718ae17b010f8788886becce82cf3369064ac Mon Sep 17 00:00:00 2001
From: Anthony Cui <anthonycui@google.com>
Date: Wed, 7 Aug 2024 11:57:22 -0700
Subject: [PATCH] Have driver_ast_parser track if a record is a union type.

PiperOrigin-RevId: 660480745
---
 .../nvidia_driver_differ/driver_ast_parser.cc |  7 +++--
 .../driver_ast_parser_test.go                 | 20 ++++++++------
 .../parser/json_definitions.go                | 26 ++++++++++++++-----
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/tools/nvidia_driver_differ/driver_ast_parser.cc b/tools/nvidia_driver_differ/driver_ast_parser.cc
index b89202ca5..98dc2c278 100644
--- a/tools/nvidia_driver_differ/driver_ast_parser.cc
+++ b/tools/nvidia_driver_differ/driver_ast_parser.cc
@@ -199,8 +199,11 @@ struct DriverStructReporter : public MatchFinder::MatchCallback {
         record_decl->getLocation().printToString(ctx->getSourceManager());
     // getTypeSize returns the size in bits, so we divide by 8 to get bytes.
     uint64_t size = ctx->getTypeSize(record_decl->getTypeForDecl()) / 8;
-    RecordDefinitions[name] =
-        json::object({{"source", source}, {"fields", fields}, {"size", size}});
+    bool is_union = record_decl->isUnion();
+    RecordDefinitions[name] = json::object({{"source", source},
+                                            {"fields", fields},
+                                            {"size", size},
+                                            {"is_union", is_union}});
   }
 };
 
diff --git a/tools/nvidia_driver_differ/driver_ast_parser_test.go b/tools/nvidia_driver_differ/driver_ast_parser_test.go
index 225a53258..b6326a862 100644
--- a/tools/nvidia_driver_differ/driver_ast_parser_test.go
+++ b/tools/nvidia_driver_differ/driver_ast_parser_test.go
@@ -84,8 +84,9 @@ func TestParser(t *testing.T) {
 					parser.RecordField{Name: "e", Type: "TestStruct::e_t[4]"},
 					parser.RecordField{Name: "f", Type: "TestUnion"},
 				},
-				Size:   44,
-				Source: "test_struct.cc:25:16",
+				Size:    44,
+				IsUnion: false,
+				Source:  "test_struct.cc:25:16",
 			},
 			"TestStruct2": parser.RecordDef{
 				Fields: []parser.RecordField{
@@ -94,24 +95,27 @@ func TestParser(t *testing.T) {
 					parser.RecordField{Name: "e", Type: "TestStruct::e_t[4]"},
 					parser.RecordField{Name: "f", Type: "TestUnion"},
 				},
-				Size:   44,
-				Source: "test_struct.cc:25:16",
+				Size:    44,
+				IsUnion: false,
+				Source:  "test_struct.cc:25:16",
 			},
 			"TestStruct::e_t": parser.RecordDef{
 				Fields: []parser.RecordField{
 					parser.RecordField{Name: "c", Type: "OtherInt"},
 					parser.RecordField{Name: "d", Type: "OtherInt"},
 				},
-				Size:   8,
-				Source: "test_struct.cc:28:3",
+				Size:    8,
+				IsUnion: false,
+				Source:  "test_struct.cc:28:3",
 			},
 			"TestUnion": parser.RecordDef{
 				Fields: []parser.RecordField{
 					parser.RecordField{Name: "u_a", Type: "int"},
 					parser.RecordField{Name: "u_b", Type: "int"},
 				},
-				Size:   4,
-				Source: "test_struct.cc:20:9",
+				Size:    4,
+				IsUnion: true,
+				Source:  "test_struct.cc:20:9",
 			},
 		},
 		Aliases: parser.TypeAliases{
diff --git a/tools/nvidia_driver_differ/parser/json_definitions.go b/tools/nvidia_driver_differ/parser/json_definitions.go
index 7686ba7c1..e8e2d7eaa 100644
--- a/tools/nvidia_driver_differ/parser/json_definitions.go
+++ b/tools/nvidia_driver_differ/parser/json_definitions.go
@@ -32,8 +32,8 @@ type InputJSON struct {
 
 // OutputJSON is the format for the output of driver_ast_parser.
 type OutputJSON struct {
-	Records RecordDefs  `json:"records"`
-	Aliases TypeAliases `json:"aliases"`
+	Records RecordDefs
+	Aliases TypeAliases
 }
 
 // Merge merges the struct definitions from b into this OutputJSON.
@@ -60,14 +60,15 @@ func (s RecordField) String() string {
 
 // RecordDef represents the definition of a record (struct or union).
 type RecordDef struct {
-	Fields []RecordField
-	Size   uint64
-	Source string
+	Fields  []RecordField
+	Size    uint64
+	IsUnion bool `json:"is_union"`
+	Source  string
 }
 
 // Equals returns true if the two record definitions are equal. We ignore the source of the records.
 func (s RecordDef) Equals(other RecordDef) bool {
-	return s.Size == other.Size && slices.Equal(s.Fields, other.Fields)
+	return s.IsUnion == other.IsUnion && s.Size == other.Size && slices.Equal(s.Fields, other.Fields)
 }
 
 // TypeDef represents the definition of a type.
@@ -88,7 +89,18 @@ func GetRecordDiff(name nvproxy.DriverStructName, a, b RecordDef) string {
 	fmt.Fprintf(&sb, "--- A: %s\n", a.Source)
 	fmt.Fprintf(&sb, "+++ B: %s\n", b.Source)
 
-	fmt.Fprintf(&sb, "struct %s\n", name)
+	switch {
+	case a.IsUnion && !b.IsUnion:
+		fmt.Fprintf(&sb, "- union %s\n", name)
+		fmt.Fprintf(&sb, "+ struct %s\n", name)
+	case !a.IsUnion && b.IsUnion:
+		fmt.Fprintf(&sb, "- struct %s\n", name)
+		fmt.Fprintf(&sb, "+ union %s\n", name)
+	case a.IsUnion && b.IsUnion:
+		fmt.Fprintf(&sb, "union %s\n", name)
+	case !a.IsUnion && !b.IsUnion:
+		fmt.Fprintf(&sb, "struct %s\n", name)
+	}
 	if a.Size != b.Size {
 		fmt.Fprintf(&sb, "  size: %d -> %d (bytes)\n", a.Size, b.Size)
 	}