TypeGraph: Add Prune pass

This lets us remove fields from types when they are no longer needed,
speeding up later passes.

A secondary benefit of pruning unused types means that we sometimes
remove types for which we can't generate correct C++ code. This can
allow us to CodeGen for complex types which reference these broken types
without actually requiring them (e.g. as template parameters).

Add a new feature flag "prune-type-graph" to control this pass. It makes
sense to prune most of the time, but for testing CodeGen functionality
on a wider range of types, it will be useful to have the option to not
prune.

oi/CodeGen.cpp

@@ -31,6 +31,7 @@
 #include "type_graph/DrgnParser.h"
 #include "type_graph/Flattener.h"
 #include "type_graph/NameGen.h"
+#include "type_graph/Prune.h"
 #include "type_graph/RemoveMembers.h"
 #include "type_graph/RemoveTopLevelPointer.h"
 #include "type_graph/TopoSorter.h"
@@ -38,13 +39,24 @@
 #include "type_graph/TypeIdentifier.h"
 #include "type_graph/Types.h"
 
+using type_graph::AddChildren;
+using type_graph::AddPadding;
+using type_graph::AlignmentCalc;
 using type_graph::Class;
 using type_graph::Container;
+using type_graph::DrgnParser;
 using type_graph::Enum;
+using type_graph::Flattener;
 using type_graph::Member;
+using type_graph::NameGen;
+using type_graph::Prune;
+using type_graph::RemoveMembers;
+using type_graph::RemoveTopLevelPointer;
+using type_graph::TopoSorter;
 using type_graph::Type;
 using type_graph::Typedef;
 using type_graph::TypeGraph;
+using type_graph::TypeIdentifier;
 
 template <typename T>
 using ref = std::reference_wrapper<T>;
@@ -407,7 +419,7 @@ void CodeGen::getClassSizeFuncConcrete(std::string_view funcName,
 
   for (size_t i = 0; i < c.members.size(); i++) {
     const auto& member = c.members[i];
-    if (member.name.starts_with(type_graph::AddPadding::MemberPrefix))
+    if (member.name.starts_with(AddPadding::MemberPrefix))
       continue;
 
     if (thriftIssetMember && thriftIssetMember != &member) {
@@ -610,7 +622,7 @@ void CodeGen::getClassTypeHandler(const Class& c, std::string& code) {
   size_t lastNonPaddingElement = -1;
   for (size_t i = 0; i < c.members.size(); i++) {
     const auto& el = c.members[i];
-    if (!el.name.starts_with(type_graph::AddPadding::MemberPrefix)) {
+    if (!el.name.starts_with(AddPadding::MemberPrefix)) {
       lastNonPaddingElement = i;
     }
   }
@@ -627,7 +639,7 @@ void CodeGen::getClassTypeHandler(const Class& c, std::string& code) {
 
     for (size_t i = 0; i < lastNonPaddingElement + 1; i++) {
       const auto& member = c.members[i];
-      if (member.name.starts_with(type_graph::AddPadding::MemberPrefix)) {
+      if (member.name.starts_with(AddPadding::MemberPrefix)) {
         continue;
       }
 
@@ -671,7 +683,7 @@ void CodeGen::getClassTypeHandler(const Class& c, std::string& code) {
   {
     for (size_t i = 0; i < lastNonPaddingElement + 1; i++) {
       const auto& member = c.members[i];
-      if (member.name.starts_with(type_graph::AddPadding::MemberPrefix)) {
+      if (member.name.starts_with(AddPadding::MemberPrefix)) {
         continue;
       }
 
@@ -761,7 +773,7 @@ bool CodeGen::codegenFromDrgn(struct drgn_type* drgnType, std::string& code) {
     return false;
   }
 
-  type_graph::TypeGraph typeGraph;
+  TypeGraph typeGraph;
   try {
     addDrgnRoot(drgnType, typeGraph);
   } catch (const type_graph::DrgnParserError& err) {
@@ -779,46 +791,47 @@ void CodeGen::registerContainer(const fs::path& path) {
   VLOG(1) << "Registered container: " << info.typeName;
 }
 
-void CodeGen::addDrgnRoot(struct drgn_type* drgnType,
-                          type_graph::TypeGraph& typeGraph) {
-  type_graph::DrgnParser drgnParser{
-      typeGraph, containerInfos_, config_.features[Feature::ChaseRawPointers]};
+void CodeGen::addDrgnRoot(struct drgn_type* drgnType, TypeGraph& typeGraph) {
+  DrgnParser drgnParser{typeGraph, containerInfos_,
+                        config_.features[Feature::ChaseRawPointers]};
   Type& parsedRoot = drgnParser.parse(drgnType);
   typeGraph.addRoot(parsedRoot);
 }
 
-void CodeGen::transform(type_graph::TypeGraph& typeGraph) {
+void CodeGen::transform(TypeGraph& typeGraph) {
   type_graph::PassManager pm;
 
   // Simplify the type graph first so there is less work for later passes
-  pm.addPass(type_graph::RemoveTopLevelPointer::createPass());
-  pm.addPass(type_graph::Flattener::createPass());
-  pm.addPass(type_graph::TypeIdentifier::createPass(config_.passThroughTypes));
+  pm.addPass(RemoveTopLevelPointer::createPass());
+  pm.addPass(Flattener::createPass());
+  pm.addPass(TypeIdentifier::createPass(config_.passThroughTypes));
+  if (config_.features[Feature::PruneTypeGraph])
+    pm.addPass(Prune::createPass());
 
   if (config_.features[Feature::PolymorphicInheritance]) {
     // Parse new children nodes
-    type_graph::DrgnParser drgnParser{
-        typeGraph, containerInfos_,
-        config_.features[Feature::ChaseRawPointers]};
-    pm.addPass(type_graph::AddChildren::createPass(drgnParser, symbols_));
+    DrgnParser drgnParser{typeGraph, containerInfos_,
+                          config_.features[Feature::ChaseRawPointers]};
+    pm.addPass(AddChildren::createPass(drgnParser, symbols_));
 
     // Re-run passes over newly added children
-    pm.addPass(type_graph::Flattener::createPass());
-    pm.addPass(
-        type_graph::TypeIdentifier::createPass(config_.passThroughTypes));
+    pm.addPass(Flattener::createPass());
+    pm.addPass(TypeIdentifier::createPass(config_.passThroughTypes));
+    if (config_.features[Feature::PruneTypeGraph])
+      pm.addPass(Prune::createPass());
   }
 
   // Calculate alignment before removing members, as those members may have an
   // influence on the class' overall alignment.
-  pm.addPass(type_graph::AlignmentCalc::createPass());
-  pm.addPass(type_graph::RemoveMembers::createPass(config_.membersToStub));
+  pm.addPass(AlignmentCalc::createPass());
+  pm.addPass(RemoveMembers::createPass(config_.membersToStub));
 
   // Add padding to fill in the gaps of removed members and ensure their
   // alignments
-  pm.addPass(type_graph::AddPadding::createPass(config_.features));
+  pm.addPass(AddPadding::createPass(config_.features));
 
-  pm.addPass(type_graph::NameGen::createPass());
-  pm.addPass(type_graph::TopoSorter::createPass());
+  pm.addPass(NameGen::createPass());
+  pm.addPass(TopoSorter::createPass());
 
   pm.run(typeGraph);
 
@@ -829,7 +842,7 @@ void CodeGen::transform(type_graph::TypeGraph& typeGraph) {
 }
 
 void CodeGen::generate(
-    type_graph::TypeGraph& typeGraph,
+    TypeGraph& typeGraph,
     std::string& code,
     struct drgn_type* drgnType /* TODO: this argument should not be required */
 ) {

oi/Features.cpp

@@ -34,6 +34,8 @@ std::string_view featureHelp(Feature f) {
       return "Capture isset data for Thrift object.";
     case Feature::TypeGraph:
       return "Use Type Graph for code generation (CodeGen v2).";
+    case Feature::PruneTypeGraph:
+      return "Prune unreachable nodes from the type graph";
     case Feature::TypedDataSegment:
       return "Use Typed Data Segment in generated code.";
     case Feature::TreeBuilderTypeChecking:

oi/Features.h

@@ -27,6 +27,7 @@
   X(GenPaddingStats, "gen-padding-stats")                  \
   X(CaptureThriftIsset, "capture-thrift-isset")            \
   X(TypeGraph, "type-graph")                               \
+  X(PruneTypeGraph, "prune-type-graph")                    \
   X(TypedDataSegment, "typed-data-segment")                \
   X(TreeBuilderTypeChecking, "tree-builder-type-checking") \
   X(TreeBuilderV2, "tree-builder-v2")                      \

oi/OID.cpp

@@ -468,6 +468,7 @@ int main(int argc, char* argv[]) {
   std::map<Feature, bool> features = {
       {Feature::PackStructs, true},
       {Feature::GenPaddingStats, true},
+      {Feature::PruneTypeGraph, true},
   };
 
   bool logAllStructs = true;

oi/OIGenerator.cpp

@@ -185,6 +185,7 @@ int OIGenerator::generate(fs::path& primaryObject, SymbolService& symbols) {
 
   std::map<Feature, bool> featuresMap = {
       {Feature::PackStructs, true},
+      {Feature::PruneTypeGraph, true},
   };
 
   OICodeGen::Config generatorConfig{};

oi/OILibraryImpl.cpp

@@ -89,6 +89,7 @@ bool OILibraryImpl::processConfigFile() {
       {
           {Feature::ChaseRawPointers, _self->opts.chaseRawPointers},
           {Feature::PackStructs, true},
+          {Feature::PruneTypeGraph, true},
           {Feature::GenJitDebug, _self->opts.generateJitDebugInfo},
       },
       compilerConfig, generatorConfig);

oi/type_graph/CMakeLists.txt

@@ -7,6 +7,7 @@ add_library(type_graph
   NameGen.cpp
   PassManager.cpp
   Printer.cpp
+  Prune.cpp
   RemoveMembers.cpp
   RemoveTopLevelPointer.cpp
   TopoSorter.cpp

oi/type_graph/Prune.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 "Prune.h"
+
+#include "TypeGraph.h"
+#include "TypeIdentifier.h"
+
+namespace type_graph {
+
+Pass Prune::createPass() {
+  auto fn = [](TypeGraph& typeGraph) {
+    Prune pass{typeGraph.resetTracker()};
+    for (auto& type : typeGraph.rootTypes()) {
+      pass.accept(type);
+    }
+  };
+
+  return Pass("Prune", fn);
+}
+
+void Prune::accept(Type& type) {
+  if (tracker_.visit(type))
+    return;
+
+  type.accept(*this);
+}
+
+void Prune::visit(Class& c) {
+  for (const auto& param : c.templateParams) {
+    accept(param.type());
+  }
+  for (const auto& parent : c.parents) {
+    accept(parent.type());
+  }
+  for (const auto& member : c.members) {
+    accept(member.type());
+  }
+  for (const auto& child : c.children) {
+    accept(child);
+  }
+
+  c.templateParams.clear();
+  c.parents.clear();
+  c.functions.clear();
+
+  // Should we bother with this shrinking? It saves memory but costs CPU
+  c.templateParams.shrink_to_fit();
+  c.parents.shrink_to_fit();
+  c.functions.shrink_to_fit();
+}
+
+}  // namespace type_graph

oi/type_graph/Prune.h

@@ -0,0 +1,51 @@
+/*
+ * 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 <string>
+#include <vector>
+
+#include "NodeTracker.h"
+#include "PassManager.h"
+#include "Types.h"
+#include "Visitor.h"
+
+namespace type_graph {
+
+/*
+ * Prune
+ *
+ * Removes unnecessary information from a type graph and releases memory where
+ * possible.
+ */
+class Prune : public RecursiveVisitor {
+ public:
+  static Pass createPass();
+
+  Prune(NodeTracker& tracker) : tracker_(tracker) {
+  }
+
+  using RecursiveVisitor::accept;
+  using RecursiveVisitor::visit;
+
+  void accept(Type& type) override;
+  void visit(Class& c) override;
+
+ private:
+  NodeTracker& tracker_;
+};
+
+}  // namespace type_graph

test/CMakeLists.txt

@@ -44,6 +44,7 @@ add_executable(test_type_graph
   test_flattener.cpp
   test_name_gen.cpp
   test_node_tracker.cpp
+  test_prune.cpp
   test_remove_members.cpp
   test_remove_top_level_pointer.cpp
   test_topo_sorter.cpp

test/test_prune.cpp

@@ -0,0 +1,61 @@
+#include <gtest/gtest.h>
+
+#include "oi/type_graph/Prune.h"
+#include "test/type_graph_utils.h"
+
+using type_graph::Prune;
+
+TEST(PruneTest, PruneClass) {
+  test(Prune::createPass(), R"(
+[0] Class: MyClass (size: 8)
+      Param
+        Primitive: int32_t
+      Param
+        Value: "123"
+      Parent (offset: 0)
+[1]     Class: MyParent (size: 4)
+          Member: a (offset: 0)
+            Primitive: int32_t
+      Member: a (offset: 0)
+        Primitive: int32_t
+      Member: b (offset: 4)
+        Primitive: int32_t
+      Function: foo
+      Function: bar
+)",
+       R"(
+[0] Class: MyClass (size: 8)
+      Member: a (offset: 0)
+        Primitive: int32_t
+      Member: b (offset: 4)
+        Primitive: int32_t
+)");
+}
+
+TEST(PruneTest, RecurseClassMember) {
+  test(Prune::createPass(), R"(
+[0] Class: MyClass (size: 0)
+      Member: xxx (offset: 0)
+[1]     Class: ClassA (size: 12)
+          Function: foo
+)",
+       R"(
+[0] Class: MyClass (size: 0)
+      Member: xxx (offset: 0)
+[1]     Class: ClassA (size: 12)
+)");
+}
+
+TEST(PruneTest, RecurseClassChild) {
+  test(Prune::createPass(), R"(
+[0] Class: MyClass (size: 0)
+      Child
+[1]     Class: ClassA (size: 12)
+          Function: foo
+)",
+       R"(
+[0] Class: MyClass (size: 0)
+      Child
+[1]     Class: ClassA (size: 12)
+)");
+}