From 2d28b20d4696cf391b7b5fa6d035f4f8f5c292a3 Mon Sep 17 00:00:00 2001
From: Alastair Robertson <ajor@meta.com>
Date: Wed, 31 May 2023 08:40:21 -0700
Subject: [PATCH] TypeGraph: Fix multi dimensional arrays

Multi dimensional arrays are not flattened into 1-D arrays when using
TypeGraph. Update TreeBuilder to account for this.

By not flattening arrays, we are able to produce more descriptive
results.

The disadvantage is that we must now recurse inside arrays
containing only primitives. A better solution to requiring flattening
would be the planned work to not recurse into any static types (not just
primitives). This would also apply to multi-dimensional arrays of
primtivies.
---
 oi/CodeGen.cpp               |  3 +++
 oi/TreeBuilder.cpp           |  9 +++++++--
 test/integration/arrays.toml | 13 ++++++++++++-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/oi/CodeGen.cpp b/oi/CodeGen.cpp
index 3b0d54e..d8d9946 100644
--- a/oi/CodeGen.cpp
+++ b/oi/CodeGen.cpp
@@ -81,6 +81,9 @@ struct OIArray {
 }
 
 void addIncludes(const TypeGraph& typeGraph, std::string& code) {
+  // Required for the offsetof() macro
+  code += "#include <cstddef>\n";
+
   // TODO deduplicate containers
   for (const Type& t : typeGraph.finalTypes) {
     if (const auto* c = dynamic_cast<const Container*>(&t)) {
diff --git a/oi/TreeBuilder.cpp b/oi/TreeBuilder.cpp
index 33b97c3..0b719a0 100644
--- a/oi/TreeBuilder.cpp
+++ b/oi/TreeBuilder.cpp
@@ -571,8 +571,13 @@ void TreeBuilder::processContainer(const Variable& variable, Node& node) {
     kind = ARRAY_TYPE;
     struct drgn_type* arrayElementType = nullptr;
     size_t numElems = 0;
-    drgn_utils::getDrgnArrayElementType(variable.type, &arrayElementType,
-                                        numElems);
+    if (config.features[Feature::TypeGraph]) {
+      arrayElementType = drgn_type_type(variable.type).type;
+      numElems = drgn_type_length(variable.type);
+    } else {
+      drgn_utils::getDrgnArrayElementType(variable.type, &arrayElementType,
+                                          numElems);
+    }
     assert(numElems > 0);
     elementTypes.push_back(
         drgn_qualified_type{arrayElementType, (enum drgn_qualifiers)(0)});
diff --git a/test/integration/arrays.toml b/test/integration/arrays.toml
index 6836594..8f3a7ae 100644
--- a/test/integration/arrays.toml
+++ b/test/integration/arrays.toml
@@ -43,7 +43,8 @@ definitions = '''
         "staticSize":0,
         "dynamicSize":0
       }]}]'''
-  [cases.multidim]
+  [cases.multidim_legacy] # Test for legacy behaviour. Remove with OICodeGen
+    cli_options = ["-Ftype-graph"]
     param_types = ["const MultiDim&"]
     setup = "return {};"
     expect_json = '''[{
@@ -56,6 +57,16 @@ definitions = '''
         "capacity":6,
         "elementStaticSize":4
       }]}]'''
+  [cases.multidim]
+    cli_options = ["-ftype-graph"]
+    param_types = ["const MultiDim&"]
+    setup = "return {};"
+    expect_json = '''[
+      {"staticSize":24, "dynamicSize":0, "exclusiveSize":0, "members":[
+        {"staticSize":24, "dynamicSize":0, "exclusiveSize":0, "length":2, "capacity":2, "elementStaticSize":12, "members":[
+          {"staticSize":12, "dynamicSize":0, "exclusiveSize":12, "length":3, "capacity":3, "elementStaticSize":4},
+          {"staticSize":12, "dynamicSize":0, "exclusiveSize":12, "length":3, "capacity":3, "elementStaticSize":4}]
+      }]}]'''
   [cases.direct_int10]
     skip = "Direct array arguments don't work"
     param_types = ["int[10]"]