TypeGraph: Add EnforceCompatibility pass

This extracts the compatibility logic from AddPadding, which allows for it to be
simplified and will make it easier to extend and eventually remove in the
future. No functional changes.

oi/CodeGen.cpp

@@ -29,6 +29,7 @@
 #include "type_graph/AddPadding.h"
 #include "type_graph/AlignmentCalc.h"
 #include "type_graph/DrgnParser.h"
+#include "type_graph/EnforceCompatibility.h"
 #include "type_graph/Flattener.h"
 #include "type_graph/NameGen.h"
 #include "type_graph/Prune.h"
@@ -45,6 +46,7 @@ using type_graph::AlignmentCalc;
 using type_graph::Class;
 using type_graph::Container;
 using type_graph::DrgnParser;
+using type_graph::EnforceCompatibility;
 using type_graph::Enum;
 using type_graph::Flattener;
 using type_graph::Member;
@@ -825,10 +827,13 @@ void CodeGen::transform(TypeGraph& typeGraph) {
   // influence on the class' overall alignment.
   pm.addPass(AlignmentCalc::createPass());
   pm.addPass(RemoveMembers::createPass(config_.membersToStub));
+  if (!config_.features[Feature::TreeBuilderV2]) {
+    pm.addPass(EnforceCompatibility::createPass());
+  }
 
   // Add padding to fill in the gaps of removed members and ensure their
   // alignments
-  pm.addPass(AddPadding::createPass(config_.features));
+  pm.addPass(AddPadding::createPass());
 
   pm.addPass(NameGen::createPass());
   pm.addPass(TopoSorter::createPass());

oi/type_graph/AddPadding.cpp

@@ -17,20 +17,16 @@
 
 #include <cassert>
 
-#include "Flattener.h"
 #include "TypeGraph.h"
 
 template <typename T>
 using ref = std::reference_wrapper<T>;
 
-using ObjectIntrospection::Feature;
-using ObjectIntrospection::FeatureSet;
-
 namespace type_graph {
 
-Pass AddPadding::createPass(FeatureSet features) {
-  auto fn = [features](TypeGraph& typeGraph) {
-    AddPadding pass(typeGraph, features);
+Pass AddPadding::createPass() {
+  auto fn = [](TypeGraph& typeGraph) {
+    AddPadding pass(typeGraph);
     for (auto& type : typeGraph.rootTypes()) {
       pass.accept(type);
     }
@@ -79,18 +75,6 @@ void AddPadding::visit(Class& c) {
       addPadding(c.members[i - 1], c.members[i].bitOffset, paddedMembers);
     }
 
-    if (!features_[Feature::TreeBuilderV2] &&
-        c.members[i].name.starts_with(Flattener::ParentPrefix)) {
-      // CodeGen v1 can't handle parent containers. It replaces them with
-      // padding.
-      auto& primitive = typeGraph_.makeType<Primitive>(Primitive::Kind::Int8);
-      auto& paddingArray =
-          typeGraph_.makeType<Array>(primitive, c.members[i].type().size());
-      paddedMembers.emplace_back(paddingArray, MemberPrefix,
-                                 c.members[i].bitOffset);
-      continue;
-    }
-
     paddedMembers.push_back(c.members[i]);
   }
 

oi/type_graph/AddPadding.h

@@ -21,7 +21,6 @@
 #include "PassManager.h"
 #include "Types.h"
 #include "Visitor.h"
-#include "oi/Features.h"
 
 namespace type_graph {
 
@@ -36,11 +35,9 @@ class TypeGraph;
  */
 class AddPadding final : public RecursiveVisitor {
  public:
-  static Pass createPass(ObjectIntrospection::FeatureSet features);
+  static Pass createPass();
 
-  explicit AddPadding(TypeGraph& typeGraph,
-                      ObjectIntrospection::FeatureSet features)
-      : typeGraph_(typeGraph), features_(features) {
+  explicit AddPadding(TypeGraph& typeGraph) : typeGraph_(typeGraph) {
   }
 
   using RecursiveVisitor::accept;
@@ -53,7 +50,6 @@ class AddPadding final : public RecursiveVisitor {
  private:
   std::unordered_set<Type*> visited_;
   TypeGraph& typeGraph_;
-  ObjectIntrospection::FeatureSet features_;
 
   void addPadding(const Member& prevMember,
                   uint64_t paddingEndBits,

oi/type_graph/CMakeLists.txt

@@ -3,6 +3,7 @@ add_library(type_graph
   AddPadding.cpp
   AlignmentCalc.cpp
   DrgnParser.cpp
+  EnforceCompatibility.cpp
   Flattener.cpp
   NameGen.cpp
   PassManager.cpp

oi/type_graph/EnforceCompatibility.cpp

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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.
+ */
+#include "EnforceCompatibility.h"
+
+#include <cassert>
+
+#include "AddPadding.h"
+#include "Flattener.h"
+#include "TypeGraph.h"
+#include "TypeIdentifier.h"
+
+namespace type_graph {
+
+Pass EnforceCompatibility::createPass() {
+  auto fn = [](TypeGraph& typeGraph) {
+    EnforceCompatibility pass{typeGraph.resetTracker()};
+    for (auto& type : typeGraph.rootTypes()) {
+      pass.accept(type);
+    }
+  };
+
+  return Pass("EnforceCompatibility", fn);
+}
+
+void EnforceCompatibility::accept(Type& type) {
+  if (tracker_.visit(type))
+    return;
+
+  type.accept(*this);
+}
+
+void EnforceCompatibility::visit(Class& c) {
+  for (auto& param : c.templateParams) {
+    accept(param.type());
+  }
+  for (auto& parent : c.parents) {
+    accept(parent.type());
+  }
+  for (auto& member : c.members) {
+    accept(member.type());
+  }
+  for (auto& child : c.children) {
+    accept(child);
+  }
+
+  // CodeGen v1 replaces parent containers with padding
+  std::erase_if(c.members, [](Member member) {
+    return member.name.starts_with(Flattener::ParentPrefix);
+  });
+}
+
+}  // namespace type_graph

oi/type_graph/EnforceCompatibility.h

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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.
+ */
+#pragma once
+
+#include <vector>
+
+#include "NodeTracker.h"
+#include "PassManager.h"
+#include "Types.h"
+#include "Visitor.h"
+
+namespace type_graph {
+
+/*
+ * EnforceCompatibility
+ *
+ * Transforms the type graph so that CodeGen produces code which is compatible
+ * with OICodeGen.
+ */
+class EnforceCompatibility : public RecursiveVisitor {
+ public:
+  static Pass createPass();
+
+  EnforceCompatibility(NodeTracker& tracker) : tracker_(tracker) {
+  }
+
+  using RecursiveVisitor::accept;
+
+  void accept(Type& type) override;
+  void visit(Class& c) override;
+
+ private:
+  NodeTracker& tracker_;
+};
+
+}  // namespace type_graph

test/CMakeLists.txt

@@ -41,6 +41,7 @@ add_executable(test_type_graph
   test_alignment_calc.cpp
   test_codegen.cpp
   test_drgn_parser.cpp
+  test_enforce_compatibility.cpp
   test_flattener.cpp
   test_name_gen.cpp
   test_node_tracker.cpp

test/test_add_padding.cpp

@@ -5,17 +5,6 @@
 #include "test/type_graph_utils.h"
 
 using namespace type_graph;
-using ObjectIntrospection::Feature;
-using ObjectIntrospection::FeatureSet;
-
-namespace {
-void test(std::vector<std::reference_wrapper<type_graph::Type>> rootTypes,
-          std::string_view expectedAfter) {
-  FeatureSet features;
-  features[Feature::TreeBuilderV2] = true;
-  ::test(AddPadding::createPass({}), rootTypes, expectedAfter);
-}
-}  // namespace
 
 TEST(AddPaddingTest, BetweenMembers) {
   auto myclass = Class{0, Class::Kind::Class, "MyClass", 16};
@@ -24,7 +13,7 @@ TEST(AddPaddingTest, BetweenMembers) {
   myclass.members.push_back(Member{myint8, "n1", 0});
   myclass.members.push_back(Member{myint64, "n2", 8 * 8});
 
-  test({myclass}, R"(
+  test(AddPadding::createPass(), {myclass}, R"(
 [0] Class: MyClass (size: 16)
       Member: n1 (offset: 0)
         Primitive: int8_t
@@ -43,7 +32,7 @@ TEST(AddPaddingTest, AtEnd) {
   myclass.members.push_back(Member{myint64, "n1", 0});
   myclass.members.push_back(Member{myint8, "n2", 8 * 8});
 
-  test({myclass}, R"(
+  test(AddPadding::createPass(), {myclass}, R"(
 [0] Struct: MyStruct (size: 16)
       Member: n1 (offset: 0)
         Primitive: int64_t
@@ -62,7 +51,7 @@ TEST(AddPaddingTest, UnionBetweenMembers) {
   myclass.members.push_back(Member{myint64, "n1", 0});
   myclass.members.push_back(Member{myint8, "n2", 0});
 
-  test({myclass}, R"(
+  test(AddPadding::createPass(), {myclass}, R"(
 [0] Union: MyUnion (size: 8)
       Member: n1 (offset: 0)
         Primitive: int64_t
@@ -72,7 +61,7 @@ TEST(AddPaddingTest, UnionBetweenMembers) {
 }
 
 TEST(AddPaddingTest, UnionAtEnd) {
-  test(AddPadding::createPass({}), R"(
+  test(AddPadding::createPass(), R"(
 [0] Union: MyUnion (size: 16)
       Member: n1 (offset: 0)
         Primitive: int64_t
@@ -117,7 +106,7 @@ TEST(AddPaddingTest, Bitfields) {
   myclass.members.push_back(b4);
   myclass.members.push_back(n);
 
-  test({myclass}, R"(
+  test(AddPadding::createPass(), {myclass}, R"(
 [0] Class: MyClass (size: 16)
       Member: b1 (offset: 0, bitsize: 3)
         Primitive: int64_t
@@ -145,39 +134,14 @@ TEST(AddPaddingTest, Bitfields) {
 )");
 }
 
-TEST(AddPaddingTest, CodeGenCompatibility) {
-  auto myint = Primitive{Primitive::Kind::Int32};
-  auto vector = getVector(1);
-  vector.templateParams.push_back(TemplateParam{myint});
-
-  auto myclass = Class{0, Class::Kind::Class, "MyClass", 24};
-  myclass.members.push_back(Member{vector, "__oi_parent", 0});
-
-  FeatureSet features;
-  features[Feature::TreeBuilderV2] = false;
-  test(AddPadding::createPass(features), {myclass}, R"(
-[0] Class: MyClass (size: 24)
-      Member: __oi_parent (offset: 0)
-[1]     Container: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-)",
-       R"(
-[0] Class: MyClass (size: 24)
-      Member: __oi_padding (offset: 0)
-[1]     Array: (length: 24)
-          Primitive: int8_t
-)");
-}
-
 TEST(AddPaddingTest, EmptyClass) {
-  testNoChange(AddPadding::createPass({}), R"(
+  testNoChange(AddPadding::createPass(), R"(
 [0] Class: MyClass (size: 0)
 )");
 }
 
 TEST(AddPaddingTest, MemberlessClass) {
-  test(AddPadding::createPass({}), R"(
+  test(AddPadding::createPass(), R"(
 [0] Class: MyClass (size: 12)
 )",
        R"(
@@ -189,7 +153,7 @@ TEST(AddPaddingTest, MemberlessClass) {
 }
 
 TEST(AddPaddingTest, MemberlessUnion) {
-  test(AddPadding::createPass({}), R"(
+  test(AddPadding::createPass(), R"(
 [0] Union: MyUnion (size: 16)
 )",
        R"(

test/test_enforce_compatibility.cpp

@@ -0,0 +1,19 @@
+#include <gtest/gtest.h>
+
+#include "oi/type_graph/EnforceCompatibility.h"
+#include "test/type_graph_utils.h"
+
+using namespace type_graph;
+
+TEST(EnforceCompatibilityTest, ParentContainers) {
+  test(EnforceCompatibility::createPass(), R"(
+[0] Class: MyClass (size: 24)
+      Member: __oi_parent (offset: 0)
+[1]     Container: std::vector (size: 24)
+          Param
+            Primitive: int32_t
+)",
+       R"(
+[0] Class: MyClass (size: 24)
+)");
+}