From 918a4d7089175d6ff10f91b72a2630415c2d3e4a Mon Sep 17 00:00:00 2001
From: Anthony Cui <anthonycui@google.com>
Date: Wed, 7 Aug 2024 12:22:22 -0700
Subject: [PATCH] Have driver_ast_parser report offset of record fields.

PiperOrigin-RevId: 660490627
---
 .../nvidia_driver_differ/driver_ast_parser.cc |  9 +++++--
 .../driver_ast_parser_test.go                 | 24 +++++++++----------
 .../parser/json_definitions.go                |  5 ++--
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/tools/nvidia_driver_differ/driver_ast_parser.cc b/tools/nvidia_driver_differ/driver_ast_parser.cc
index 98dc2c278..572b31912 100644
--- a/tools/nvidia_driver_differ/driver_ast_parser.cc
+++ b/tools/nvidia_driver_differ/driver_ast_parser.cc
@@ -187,9 +187,14 @@ struct DriverStructReporter : public MatchFinder::MatchCallback {
         absl::StrAppend(&field_type_name, "[", array_size, "]");
       }
 
+      // getFieldOffset returns the size in bits, so we divide by 8 to get
+      // bytes.
+      uint64_t offset = ctx->getFieldOffset(field) / 8;
+
       // Add field to json.
-      fields.push_back(json::object(
-          {{"name", field->getNameAsString()}, {"type", field_type_name}}));
+      fields.push_back(json::object({{"name", field->getNameAsString()},
+                                     {"type", field_type_name},
+                                     {"offset", offset}}));
 
       // Recurse on the field type.
       add_type_definition(field_type, base_type_name, ctx);
diff --git a/tools/nvidia_driver_differ/driver_ast_parser_test.go b/tools/nvidia_driver_differ/driver_ast_parser_test.go
index b6326a862..86ce3300e 100644
--- a/tools/nvidia_driver_differ/driver_ast_parser_test.go
+++ b/tools/nvidia_driver_differ/driver_ast_parser_test.go
@@ -79,10 +79,10 @@ func TestParser(t *testing.T) {
 		Records: parser.RecordDefs{
 			"TestStruct": parser.RecordDef{
 				Fields: []parser.RecordField{
-					parser.RecordField{Name: "a", Type: "int"},
-					parser.RecordField{Name: "b", Type: "int"},
-					parser.RecordField{Name: "e", Type: "TestStruct::e_t[4]"},
-					parser.RecordField{Name: "f", Type: "TestUnion"},
+					parser.RecordField{Name: "a", Type: "int", Offset: 0},
+					parser.RecordField{Name: "b", Type: "int", Offset: 4},
+					parser.RecordField{Name: "e", Type: "TestStruct::e_t[4]", Offset: 8},
+					parser.RecordField{Name: "f", Type: "TestUnion", Offset: 40},
 				},
 				Size:    44,
 				IsUnion: false,
@@ -90,10 +90,10 @@ func TestParser(t *testing.T) {
 			},
 			"TestStruct2": parser.RecordDef{
 				Fields: []parser.RecordField{
-					parser.RecordField{Name: "a", Type: "int"},
-					parser.RecordField{Name: "b", Type: "int"},
-					parser.RecordField{Name: "e", Type: "TestStruct::e_t[4]"},
-					parser.RecordField{Name: "f", Type: "TestUnion"},
+					parser.RecordField{Name: "a", Type: "int", Offset: 0},
+					parser.RecordField{Name: "b", Type: "int", Offset: 4},
+					parser.RecordField{Name: "e", Type: "TestStruct::e_t[4]", Offset: 8},
+					parser.RecordField{Name: "f", Type: "TestUnion", Offset: 40},
 				},
 				Size:    44,
 				IsUnion: false,
@@ -101,8 +101,8 @@ func TestParser(t *testing.T) {
 			},
 			"TestStruct::e_t": parser.RecordDef{
 				Fields: []parser.RecordField{
-					parser.RecordField{Name: "c", Type: "OtherInt"},
-					parser.RecordField{Name: "d", Type: "OtherInt"},
+					parser.RecordField{Name: "c", Type: "OtherInt", Offset: 0},
+					parser.RecordField{Name: "d", Type: "OtherInt", Offset: 4},
 				},
 				Size:    8,
 				IsUnion: false,
@@ -110,8 +110,8 @@ func TestParser(t *testing.T) {
 			},
 			"TestUnion": parser.RecordDef{
 				Fields: []parser.RecordField{
-					parser.RecordField{Name: "u_a", Type: "int"},
-					parser.RecordField{Name: "u_b", Type: "int"},
+					parser.RecordField{Name: "u_a", Type: "int", Offset: 0},
+					parser.RecordField{Name: "u_b", Type: "int", Offset: 0},
 				},
 				Size:    4,
 				IsUnion: true,
diff --git a/tools/nvidia_driver_differ/parser/json_definitions.go b/tools/nvidia_driver_differ/parser/json_definitions.go
index e8e2d7eaa..26eaec2c1 100644
--- a/tools/nvidia_driver_differ/parser/json_definitions.go
+++ b/tools/nvidia_driver_differ/parser/json_definitions.go
@@ -50,8 +50,9 @@ func (a *OutputJSON) Merge(b OutputJSON) {
 
 // RecordField represents a field in a record (struct or union).
 type RecordField struct {
-	Name string
-	Type string
+	Name   string
+	Type   string
+	Offset uint64
 }
 
 func (s RecordField) String() string {