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]"]