TypeGraph: Fix handling for classes which inherit from containers

We previously moved container identification later in CodeGen in order
to preserve information for AlignmentCalc.

However, Flattener needs to know if a class is a container in order to
apply its special handling for this case.

This new approach moves container identification back into DrgnParser,
but has Container own a type node, representing its layout. This
underlying type node can be used for calculating a container's
alignment.

oi/CodeGen.cpp

@@ -32,7 +32,6 @@
 #include "type_graph/DrgnParser.h"
 #include "type_graph/EnforceCompatibility.h"
 #include "type_graph/Flattener.h"
-#include "type_graph/IdentifyContainers.h"
 #include "type_graph/KeyCapture.h"
 #include "type_graph/NameGen.h"
 #include "type_graph/Prune.h"
@@ -56,7 +55,6 @@ using type_graph::DrgnParserOptions;
 using type_graph::EnforceCompatibility;
 using type_graph::Enum;
 using type_graph::Flattener;
-using type_graph::IdentifyContainers;
 using type_graph::KeyCapture;
 using type_graph::Member;
 using type_graph::NameGen;
@@ -1147,7 +1145,7 @@ void CodeGen::addDrgnRoot(struct drgn_type* drgnType, TypeGraph& typeGraph) {
   DrgnParserOptions options{
       .chaseRawPointers = config_.features[Feature::ChaseRawPointers],
   };
-  DrgnParser drgnParser{typeGraph, options};
+  DrgnParser drgnParser{typeGraph, containerInfos_, options};
   Type& parsedRoot = drgnParser.parse(drgnType);
   typeGraph.addRoot(parsedRoot);
 }
@@ -1159,7 +1157,6 @@ void CodeGen::transform(TypeGraph& typeGraph) {
   pm.addPass(RemoveTopLevelPointer::createPass());
   pm.addPass(Flattener::createPass());
   pm.addPass(AlignmentCalc::createPass());
-  pm.addPass(IdentifyContainers::createPass(containerInfos_));
   pm.addPass(TypeIdentifier::createPass(config_.passThroughTypes));
   if (config_.features[Feature::PruneTypeGraph])
     pm.addPass(Prune::createPass());
@@ -1169,13 +1166,12 @@ void CodeGen::transform(TypeGraph& typeGraph) {
     DrgnParserOptions options{
         .chaseRawPointers = config_.features[Feature::ChaseRawPointers],
     };
-    DrgnParser drgnParser{typeGraph, options};
+    DrgnParser drgnParser{typeGraph, containerInfos_, options};
     pm.addPass(AddChildren::createPass(drgnParser, symbols_));
 
     // Re-run passes over newly added children
     pm.addPass(Flattener::createPass());
     pm.addPass(AlignmentCalc::createPass());
-    pm.addPass(IdentifyContainers::createPass(containerInfos_));
     pm.addPass(TypeIdentifier::createPass(config_.passThroughTypes));
     if (config_.features[Feature::PruneTypeGraph])
       pm.addPass(Prune::createPass());

oi/type_graph/AlignmentCalc.cpp

@@ -48,15 +48,7 @@ void AlignmentCalc::accept(Type& type) {
 }
 
 void AlignmentCalc::visit(Class& c) {
-  for (const auto& param : c.templateParams) {
+  RecursiveVisitor::visit(c);
-    accept(param.type());
-  }
-  for (const auto& parent : c.parents) {
-    accept(parent.type());
-  }
-  for (const auto& child : c.children) {
-    accept(child);
-  }
 
   uint64_t alignment = 1;
   for (auto& member : c.members) {
@@ -82,4 +74,12 @@ void AlignmentCalc::visit(Class& c) {
   }
 }
 
+void AlignmentCalc::visit(Container& c) {
+  RecursiveVisitor::visit(c);
+
+  if (c.underlying()) {
+    c.setAlign(c.underlying()->align());
+  }
+}
+
 }  // namespace oi::detail::type_graph

oi/type_graph/AlignmentCalc.h

@@ -41,6 +41,7 @@ class AlignmentCalc final : public RecursiveVisitor {
 
   void accept(Type& type) override;
   void visit(Class& c) override;
+  void visit(Container& c) override;
 
  private:
   std::unordered_set<Type*> visited_;

oi/type_graph/CMakeLists.txt

@@ -5,7 +5,6 @@ add_library(type_graph
   DrgnParser.cpp
   EnforceCompatibility.cpp
   Flattener.cpp
-  IdentifyContainers.cpp
   KeyCapture.cpp
   NameGen.cpp
   PassManager.cpp

oi/type_graph/DrgnParser.cpp

@@ -25,6 +25,8 @@ extern "C" {
 #include <drgn.h>
 }
 
+#include <regex>
+
 namespace oi::detail::type_graph {
 namespace {
 
@@ -146,7 +148,7 @@ Type& DrgnParser::enumerateType(struct drgn_type* type) {
   return *t;
 }
 
-Class& DrgnParser::enumerateClass(struct drgn_type* type) {
+Type& DrgnParser::enumerateClass(struct drgn_type* type) {
   std::string fqName;
   char* nameStr = nullptr;
   size_t length = 0;
@@ -180,14 +182,20 @@ Class& DrgnParser::enumerateClass(struct drgn_type* type) {
                             std::to_string(drgn_type_kind(type))};
   }
 
-  auto& c = makeType<Class>(
+  Class& c =
-      type, kind, std::move(name), std::move(fqName), size, virtuality);
+      makeType<Class>(type, kind, std::move(name), fqName, size, virtuality);
 
   enumerateClassTemplateParams(type, c.templateParams);
   enumerateClassParents(type, c.parents);
   enumerateClassMembers(type, c.members);
   enumerateClassFunctions(type, c.functions);
 
+  if (auto* info = getContainerInfo(fqName)) {
+    auto& container = makeType<Container>(type, *info, size, &c);
+    container.templateParams = c.templateParams;
+    return container;
+  }
+
   return c;
 }
 
@@ -505,6 +513,15 @@ bool DrgnParser::chasePointer() const {
   return options_.chaseRawPointers;
 }
 
+ContainerInfo* DrgnParser::getContainerInfo(const std::string& fqName) const {
+  for (const auto& containerInfo : containers_) {
+    if (std::regex_search(fqName, containerInfo->matcher)) {
+      return containerInfo.get();
+    }
+  }
+  return nullptr;
+}
+
 DrgnParserError::DrgnParserError(const std::string& msg, struct drgn_error* err)
     : std::runtime_error{msg + ": " + std::to_string(err->code) + " " +
                          err->message},

oi/type_graph/DrgnParser.h

@@ -51,14 +51,16 @@ struct DrgnParserOptions {
  */
 class DrgnParser {
  public:
-  DrgnParser(TypeGraph& typeGraph, DrgnParserOptions options)
+  DrgnParser(TypeGraph& typeGraph,
-      : typeGraph_(typeGraph), options_(options) {
+             const std::vector<std::unique_ptr<ContainerInfo>>& containers,
+             DrgnParserOptions options)
+      : typeGraph_(typeGraph), containers_(containers), options_(options) {
   }
   Type& parse(struct drgn_type* root);
 
  private:
   Type& enumerateType(struct drgn_type* type);
-  Class& enumerateClass(struct drgn_type* type);
+  Type& enumerateClass(struct drgn_type* type);
   Enum& enumerateEnum(struct drgn_type* type);
   Typedef& enumerateTypedef(struct drgn_type* type);
   Type& enumeratePointer(struct drgn_type* type);
@@ -81,10 +83,11 @@ class DrgnParser {
   template <typename T, typename... Args>
   T& makeType(struct drgn_type* drgnType, Args&&... args) {
     auto& newType = typeGraph_.makeType<T>(std::forward<Args>(args)...);
-    drgn_types_.insert({drgnType, newType});
+    drgn_types_.insert_or_assign(drgnType, newType);
     return newType;
   }
   bool chasePointer() const;
+  ContainerInfo* getContainerInfo(const std::string& fqName) const;
 
   // Store a mapping of drgn types to type graph nodes for deduplication during
   // parsing. This stops us getting caught in cycles.
@@ -92,6 +95,7 @@ class DrgnParser {
       drgn_types_;
 
   TypeGraph& typeGraph_;
+  const std::vector<std::unique_ptr<ContainerInfo>>& containers_;
   int depth_;
   DrgnParserOptions options_;
 };

oi/type_graph/Flattener.cpp

@@ -134,8 +134,6 @@ void Flattener::visit(Class& c) {
   //   };
   //    TODO comment about virtual inheritance
 
-  // TODO alignment of parent classes
-
   // Flatten types referenced by template params, parents and members
   for (const auto& param : c.templateParams) {
     accept(param.type());
@@ -207,12 +205,4 @@ void Flattener::visit(Class& c) {
   }
 }
 
-void Flattener::visit(Container& c) {
-  // Containers themselves don't need to be flattened, but their template
-  // parameters might need to be
-  for (const auto& templateParam : c.templateParams) {
-    accept(templateParam.type());
-  }
-}
-
 }  // namespace oi::detail::type_graph

oi/type_graph/Flattener.h

@@ -42,7 +42,6 @@ class Flattener : public RecursiveVisitor {
 
   void accept(Type& type) override;
   void visit(Class& c) override;
-  void visit(Container& c) override;
 
   static const inline std::string ParentPrefix = "__oi_parent";
 

oi/type_graph/IdentifyContainers.cpp

@@ -1,73 +0,0 @@
-/*
- * 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 "IdentifyContainers.h"
-
-#include <regex>
-
-#include "TypeGraph.h"
-#include "oi/ContainerInfo.h"
-
-namespace oi::detail::type_graph {
-
-Pass IdentifyContainers::createPass(
-    const std::vector<std::unique_ptr<ContainerInfo>>& containers) {
-  auto fn = [&containers](TypeGraph& typeGraph, NodeTracker&) {
-    IdentifyContainers typeId{typeGraph, containers};
-    for (auto& type : typeGraph.rootTypes()) {
-      type = typeId.mutate(type);
-    }
-  };
-
-  return Pass("IdentifyContainers", fn);
-}
-
-IdentifyContainers::IdentifyContainers(
-    TypeGraph& typeGraph,
-    const std::vector<std::unique_ptr<ContainerInfo>>& containers)
-    : tracker_(typeGraph.size()),
-      typeGraph_(typeGraph),
-      containers_(containers) {
-}
-
-Type& IdentifyContainers::mutate(Type& type) {
-  if (Type* mutated = tracker_.get(type))
-    return *mutated;
-
-  Type& mutated = type.accept(*this);
-  tracker_.set(type, mutated);
-  return mutated;
-}
-
-Type& IdentifyContainers::visit(Class& c) {
-  for (const auto& containerInfo : containers_) {
-    if (!std::regex_search(c.fqName(), containerInfo->matcher)) {
-      continue;
-    }
-
-    auto& container = typeGraph_.makeType<Container>(*containerInfo, c.size());
-    container.templateParams = c.templateParams;
-
-    tracker_.set(c, container);
-    RecursiveMutator::visit(container);
-    return container;
-  }
-
-  tracker_.set(c, c);
-  RecursiveMutator::visit(c);
-  return c;
-}
-
-}  // namespace oi::detail::type_graph

oi/type_graph/IdentifyContainers.h

@@ -1,58 +0,0 @@
-/*
- * 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 <functional>
-#include <unordered_map>
-#include <vector>
-
-#include "NodeTracker.h"
-#include "PassManager.h"
-#include "Types.h"
-#include "Visitor.h"
-#include "oi/ContainerInfo.h"
-
-namespace oi::detail::type_graph {
-
-class TypeGraph;
-
-/*
- * IdentifyContainers
- *
- * Walks a flattened type graph and replaces type nodes based on container
- * definition TOML files.
- */
-class IdentifyContainers : public RecursiveMutator {
- public:
-  static Pass createPass(
-      const std::vector<std::unique_ptr<ContainerInfo>>& containers);
-
-  IdentifyContainers(
-      TypeGraph& typeGraph,
-      const std::vector<std::unique_ptr<ContainerInfo>>& containers);
-
-  using RecursiveMutator::mutate;
-
-  Type& mutate(Type& type) override;
-  Type& visit(Class& c) override;
-
- private:
-  MutationTracker tracker_;
-  TypeGraph& typeGraph_;
-  const std::vector<std::unique_ptr<ContainerInfo>>& containers_;
-};
-
-}  // namespace oi::detail::type_graph

oi/type_graph/Printer.cpp

@@ -85,7 +85,7 @@ void Printer::visit(const Class& c) {
     print_function(function);
   }
   for (auto& child : c.children) {
-    print_child(child);
+    print_type("Child", child);
   }
 }
 
@@ -93,11 +93,14 @@ void Printer::visit(const Container& c) {
   if (prefix(c))
     return;
 
-  out_ << "Container: " << c.name() << " (size: " << c.size() << ")"
+  out_ << "Container: " << c.name() << " (size: " << c.size()
-       << std::endl;
+       << align_str(c.align()) << ")" << std::endl;
   for (const auto& param : c.templateParams) {
     print_param(param);
   }
+  if (c.underlying()) {
+    print_type("Underlying", *c.underlying());
+  }
 }
 
 void Printer::visit(const Primitive& p) {
@@ -240,10 +243,10 @@ void Printer::print_function(const Function& function) {
   depth_--;
 }
 
-void Printer::print_child(const Type& child) {
+void Printer::print_type(std::string_view header, const Type& child) {
   depth_++;
   prefix();
-  out_ << "Child" << std::endl;
+  out_ << header << std::endl;
   print(child);
   depth_--;
 }

oi/type_graph/Printer.h

@@ -51,7 +51,7 @@ class Printer : public ConstVisitor {
   void print_parent(const Parent& parent);
   void print_member(const Member& member);
   void print_function(const Function& function);
-  void print_child(const Type& child);
+  void print_type(std::string_view header, const Type& child);
   void print_value(const std::string& value);
   void print_qualifiers(const QualifierSet& qualifiers);
   void print_enumerator(int64_t val, const std::string& name);

oi/type_graph/Prune.cpp

@@ -39,18 +39,7 @@ void Prune::accept(Type& type) {
 }
 
 void Prune::visit(Class& c) {
-  for (const auto& param : c.templateParams) {
+  RecursiveVisitor::visit(c);
-    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();
@@ -62,4 +51,10 @@ void Prune::visit(Class& c) {
   c.functions.shrink_to_fit();
 }
 
+void Prune::visit(Container& c) {
+  RecursiveVisitor::visit(c);
+
+  c.setUnderlying(nullptr);
+}
+
 }  // namespace oi::detail::type_graph

oi/type_graph/Prune.h

@@ -43,6 +43,7 @@ class Prune : public RecursiveVisitor {
 
   void accept(Type& type) override;
   void visit(Class& c) override;
+  void visit(Container& c) override;
 
  private:
   NodeTracker& tracker_;

oi/type_graph/TypeIdentifier.cpp

@@ -79,7 +79,8 @@ void TypeIdentifier::visit(Container& c) {
               it != passThroughTypeDummys_.end()) {
             dummy = &it->second.get();
           } else {
-            dummy = &typeGraph_.makeType<Container>(info, param.type().size());
+            dummy = &typeGraph_.makeType<Container>(
+                info, param.type().size(), paramClass);
             dummy->templateParams = paramClass->templateParams;
             passThroughTypeDummys_.insert(it,
                                           {paramClass->id(), std::ref(*dummy)});

oi/type_graph/Types.h

@@ -150,7 +150,6 @@ struct Function {
   int virtuality;
 };
 
-class Class;
 struct Parent {
   Parent(Type& type, uint64_t bitOffset) : type_(type), bitOffset(bitOffset) {
   }
@@ -298,10 +297,6 @@ class Class : public Type {
 
   DECLARE_ACCEPT
 
-  Kind kind() const {
-    return kind_;
-  }
-
   virtual const std::string& name() const override {
     return name_;
   }
@@ -326,12 +321,16 @@ class Class : public Type {
     return align_;
   }
 
+  void setAlign(uint64_t alignment) {
+    align_ = alignment;
+  }
+
   virtual NodeId id() const override {
     return id_;
   }
 
-  void setAlign(uint64_t alignment) {
+  Kind kind() const {
-    align_ = alignment;
+    return kind_;
   }
 
   int virtuality() const {
@@ -371,10 +370,19 @@ class Class : public Type {
   bool packed_ = false;
 };
 
+/*
+ * Container
+ *
+ * A type of class for which we can do special processing.
+ */
 class Container : public Type {
  public:
-  Container(NodeId id, const ContainerInfo& containerInfo, size_t size)
+  Container(NodeId id,
+            const ContainerInfo& containerInfo,
+            size_t size,
+            Type* underlying)
       : containerInfo_(containerInfo),
+        underlying_(underlying),
         name_(containerInfo.typeName),
         inputName_(containerInfo.typeName),
         size_(size),
@@ -386,6 +394,7 @@ class Container : public Type {
             const ContainerInfo& containerInfo)
       : templateParams(other.templateParams),
         containerInfo_(containerInfo),
+        underlying_(other.underlying_),
         name_(other.name_),
         inputName_(other.inputName_),
         size_(other.size_),
@@ -396,22 +405,18 @@ class Container : public Type {
 
   DECLARE_ACCEPT
 
-  const std::string& containerName() const {
-    return containerInfo_.typeName;
-  }
-
   virtual const std::string& name() const override {
     return name_;
   }
 
-  void setName(std::string name) {
-    name_ = std::move(name);
-  }
-
   virtual std::string_view inputName() const override {
     return inputName_;
   }
 
+  void setName(std::string name) {
+    name_ = std::move(name);
+  }
+
   void setInputName(std::string name) {
     inputName_ = std::move(name);
   }
@@ -424,18 +429,31 @@ class Container : public Type {
     return align_;
   }
 
+  void setAlign(uint64_t alignment) {
+    align_ = alignment;
+  }
+
   virtual NodeId id() const override {
     return id_;
   }
 
-  void setAlign(uint64_t alignment) {
+  const std::string& containerName() const {
-    align_ = alignment;
+    return containerInfo_.typeName;
+  }
+
+  Type* underlying() const {
+    return underlying_;
+  }
+
+  void setUnderlying(Type* underlying) {
+    underlying_ = underlying;
   }
 
   std::vector<TemplateParam> templateParams;
   const ContainerInfo& containerInfo_;
 
  private:
+  Type* underlying_;
   std::string name_;
   std::string inputName_;
   size_t size_;

oi/type_graph/Visitor.h

@@ -107,6 +107,7 @@ class RecursiveVisitor : public Visitor {
     for (const auto& param : c.templateParams) {
       accept(param.type());
     }
+    accept(c.underlying());
   }
   virtual void visit(Primitive&) {
   }
@@ -140,6 +141,11 @@ class RecursiveMutator : public Mutator {
  public:
   virtual ~RecursiveMutator() = default;
   virtual Type& mutate(Type&) = 0;
+  virtual Type* mutate(Type* type) {
+    if (type)
+      return &mutate(*type);
+    return nullptr;
+  }
   virtual Type& visit(Incomplete& i) {
     return i;
   }
@@ -162,6 +168,7 @@ class RecursiveMutator : public Mutator {
     for (auto& param : c.templateParams) {
       param.setType(mutate(param.type()));
     }
+    c.setUnderlying(mutate(c.underlying()));
     return c;
   }
   virtual Type& visit(Primitive& p) {

test/CMakeLists.txt

@@ -43,7 +43,6 @@ add_executable(test_type_graph
   test_drgn_parser.cpp
   test_enforce_compatibility.cpp
   test_flattener.cpp
-  test_identify_containers.cpp
   test_key_capture.cpp
   test_name_gen.cpp
   test_node_tracker.cpp

test/TypeGraphParser.cpp

@@ -265,10 +265,11 @@ Type& TypeGraphParser::parseType(std::string_view& input, size_t rootIndent) {
 
     auto size = parseNumericAttribute(line, nodeTypeName, "size: ");
 
-    Container& c = typeGraph_.makeType<Container>(id, info, size);
+    Container& c = typeGraph_.makeType<Container>(id, info, size, nullptr);
     nodesById_.insert({id, c});
 
     parseParams(c, input, indent + 2);
+    parseUnderlying(c, input, indent + 2);
 
     type = &c;
   } else if (nodeTypeName == "Primitive") {
@@ -452,3 +453,26 @@ void TypeGraphParser::parseChildren(Class& c,
   // No more children for us - put back the line we just read
   input = origInput;
 }
+
+void TypeGraphParser::parseUnderlying(Container& c,
+                                      std::string_view& input,
+                                      size_t rootIndent) {
+  std::string_view origInput = input;
+  std::string_view line;
+  getline(input, line);
+
+  size_t indent = stripIndent(line);
+  if (indent != rootIndent) {
+    input = origInput;
+    return;
+  }
+
+  // Format: "Underlying"
+  if (!tryRemovePrefix(line, "Underlying")) {
+    input = origInput;
+    return;
+  }
+
+  Type& type = parseType(input, rootIndent + 2);
+  c.setUnderlying(&type);
+}

test/TypeGraphParser.h

@@ -34,6 +34,9 @@ class TypeGraphParser {
   void parseMembers(Class& c, std::string_view& input, size_t rootIndent);
   void parseFunctions(Class& c, std::string_view& input, size_t rootIndent);
   void parseChildren(Class& c, std::string_view& input, size_t rootIndent);
+  void parseUnderlying(Container& c,
+                       std::string_view& input,
+                       size_t rootIndent);
 };
 
 class TypeGraphParserError : public std::runtime_error {

test/test_add_children.cpp

@@ -79,11 +79,11 @@ TEST_F(AddChildrenTest, InheritancePolymorphic) {
          Function: A
          Function: A
          Child
-[17]       Class: B (size: 40)
+[18]       Class: B (size: 40)
              Parent (offset: 0)
                [0]
              Member: vec_b (offset: 16)
-[4]            Class: vector<int, std::allocator<int> > (size: 24)
+[17]           Container: std::vector (size: 24)
                  Param
                    Primitive: int32_t
                  Param
@@ -106,15 +106,17 @@ TEST_F(AddChildrenTest, InheritancePolymorphic) {
                      Function: ~allocator
                      Function: allocate
                      Function: deallocate
-                 *
+                 Underlying
+[4]                Class: vector<int, std::allocator<int> > (size: 24)
+*
              Function: ~B (virtual)
              Function: myfunc (virtual)
              Function: B
              Function: B
              Child
-[19]           Class: C (size: 48)
+[20]           Class: C (size: 48)
                  Parent (offset: 0)
-                   [17]
+                   [18]
                  Member: int_c (offset: 40)
                    Primitive: int32_t
                  Function: ~C (virtual)
@@ -135,11 +137,11 @@ TEST_F(AddChildrenTest, InheritancePolymorphic) {
          Function: ~A (virtual)
          Function: myfunc (virtual)
          Child
-[13]       Class: B (size: 40)
+[14]       Class: B (size: 40)
              Parent (offset: 0)
                [0]
              Member: vec_b (offset: 16)
-[4]            Class: vector<int, std::allocator<int> > (size: 24)
+[13]           Container: std::vector (size: 24)
                  Param
                    Primitive: int32_t
                  Param
@@ -159,6 +161,8 @@ TEST_F(AddChildrenTest, InheritancePolymorphic) {
                      Function: ~allocator
                      Function: allocate
                      Function: deallocate
+                 Underlying
+[4]                Class: vector<int, std::allocator<int> > (size: 24)
                  *
              Function: operator=
              Function: B
@@ -166,9 +170,9 @@ TEST_F(AddChildrenTest, InheritancePolymorphic) {
              Function: ~B (virtual)
              Function: myfunc (virtual)
              Child
-[15]           Class: C (size: 48)
+[16]           Class: C (size: 48)
                  Parent (offset: 0)
-                   [13]
+                   [14]
                  Member: int_c (offset: 40)
                    Primitive: int32_t
                  Function: operator=

test/test_alignment_calc.cpp

@@ -58,16 +58,16 @@ TEST(AlignmentCalcTest, StructInContainer) {
           Member: n (offset: 0)
             Primitive: int8_t
           Member: n (offset: 8)
-            Primitive: int64_t
+            Primitive: int32_t
 )",
        R"(
 [0] Container: std::vector (size: 8)
       Param
-[1]     Class: MyClass (size: 16, align: 8)
+[1]     Class: MyClass (size: 16, align: 4)
           Member: n (offset: 0, align: 1)
             Primitive: int8_t
-          Member: n (offset: 8, align: 8)
+          Member: n (offset: 8, align: 4)
-            Primitive: int64_t
+            Primitive: int32_t
 )");
 }
 
@@ -264,3 +264,33 @@ TEST(AlignmentCalcTest, Typedef) {
               Primitive: int8_t
 )");
 }
+
+TEST(AlignmentCalcTest, Container) {
+  test(AlignmentCalc::createPass(),
+       R"(
+[0] Container: std::vector (size: 24)
+      Underlying
+[1]     Class: vector (size: 24)
+          Member: n (offset: 0)
+            Primitive: int8_t
+          Member: s (offset: 4)
+[2]         Struct: MyStruct (size: 8)
+              Member: n1 (offset: 0)
+                Primitive: int32_t
+              Member: n2 (offset: 4)
+                Primitive: int32_t
+)",
+       R"(
+[0] Container: std::vector (size: 24, align: 4)
+      Underlying
+[1]     Class: vector (size: 24, align: 4)
+          Member: n (offset: 0, align: 1)
+            Primitive: int8_t
+          Member: s (offset: 4, align: 4)
+[2]         Struct: MyStruct (size: 8, align: 4)
+              Member: n1 (offset: 0, align: 4)
+                Primitive: int32_t
+              Member: n2 (offset: 4, align: 4)
+                Primitive: int32_t
+)");
+}

test/test_codegen.cpp

@@ -45,13 +45,15 @@ void testTransform(OICodeGen::Config& config,
 
 void testTransform(std::string_view input, std::string_view expectedAfter) {
   OICodeGen::Config config;
+  config.features[Feature::PruneTypeGraph] = true;
+  config.features[Feature::TreeBuilderV2] = true;
   testTransform(config, input, expectedAfter);
 }
 }  // namespace
 
 TEST(CodeGenTest, TransformContainerAllocator) {
   testTransform(R"(
-[0] Class: std::vector (size: 24)
+[0] Container: std::vector (size: 24)
       Param
         Primitive: int32_t
       Param
@@ -62,18 +64,18 @@ TEST(CodeGenTest, TransformContainerAllocator) {
           Function: deallocate
 )",
                 R"(
-[2] Container: std::vector<int32_t, DummyAllocator<int32_t, 8, 1, 3>> (size: 24)
+[0] Container: std::vector<int32_t, DummyAllocator<int32_t, 8, 1, 2>> (size: 24)
       Param
         Primitive: int32_t
       Param
-[3]     DummyAllocator [MyAlloc] (size: 8, align: 1)
+[2]     DummyAllocator [MyAlloc] (size: 8, align: 1)
           Primitive: int32_t
 )");
 }
 
 TEST(CodeGenTest, TransformContainerAllocatorParamInParent) {
   testTransform(R"(
-[0] Class: std::map (size: 24)
+[0] Container: std::map (size: 24)
       Param
         Primitive: int32_t
       Param
@@ -83,7 +85,7 @@ TEST(CodeGenTest, TransformContainerAllocatorParamInParent) {
           Parent (offset: 0)
 [2]         Struct: MyAllocBase<std::pair<const int, int>> (size: 1)
               Param
-[3]             Class: std::pair (size: 8)
+[3]             Container: std::pair (size: 8)
                   Param
                     Primitive: int32_t
                     Qualifiers: const
@@ -95,14 +97,14 @@ TEST(CodeGenTest, TransformContainerAllocatorParamInParent) {
           Function: deallocate
 )",
                 R"(
-[4] Container: std::map<int32_t, int32_t, DummyAllocator<std::pair<int32_t const, int32_t>, 0, 1, 6>> (size: 24)
+[0] Container: std::map<int32_t, int32_t, DummyAllocator<std::pair<int32_t const, int32_t>, 0, 1, 4>> (size: 24)
       Param
         Primitive: int32_t
       Param
         Primitive: int32_t
       Param
-[6]     DummyAllocator [MyAlloc<std::pair<const int, int>>] (size: 0, align: 1)
+[4]     DummyAllocator [MyAlloc<std::pair<const int, int>>] (size: 0, align: 1)
-[5]       Container: std::pair<int32_t const, int32_t> (size: 8)
+[3]       Container: std::pair<int32_t const, int32_t> (size: 8)
             Param
               Primitive: int32_t
               Qualifiers: const
@@ -157,7 +159,7 @@ TEST(CodeGenTest, UnionMembersAlignment) {
 
 TEST(CodeGenTest, ReplaceContainersAndDummies) {
   testTransform(R"(
-[0] Class: std::vector (size: 24)
+[0] Container: std::vector (size: 24)
       Param
         Primitive: uint32_t
       Param
@@ -168,11 +170,87 @@ TEST(CodeGenTest, ReplaceContainersAndDummies) {
           Function: deallocate
 )",
                 R"(
-[2] Container: std::vector<uint32_t, DummyAllocator<uint32_t, 0, 1, 3>> (size: 24)
+[0] Container: std::vector<uint32_t, DummyAllocator<uint32_t, 0, 1, 2>> (size: 24)
       Param
         Primitive: uint32_t
       Param
-[3]     DummyAllocator [allocator<int>] (size: 0, align: 1)
+[2]     DummyAllocator [allocator<int>] (size: 0, align: 1)
           Primitive: uint32_t
 )");
 }
+
+TEST(CodeGenTest, ContainerAlignment) {
+  testTransform(R"(
+[0] Class: MyClass (size: 24)
+      Member: container (offset: 0)
+[1]     Container: std::vector (size: 24)
+          Param
+            Primitive: int32_t
+          Underlying
+[2]         Class: vector (size: 24)
+              Member: __impl__ (offset: 0)
+                Primitive: StubbedPointer
+              Member: __impl__ (offset: 8)
+                Primitive: StubbedPointer
+              Member: __impl__ (offset: 16)
+                Primitive: StubbedPointer
+)",
+                R"(
+[0] Class: MyClass_0 [MyClass] (size: 24, align: 8)
+      Member: container_0 [container] (offset: 0, align: 8)
+[1]     Container: std::vector<int32_t> (size: 24, align: 8)
+          Param
+            Primitive: int32_t
+)");
+}
+
+TEST(CodeGenTest, InheritFromContainer) {
+  testTransform(R"(
+[0] Class: MyClass (size: 24)
+      Parent (offset: 0)
+[1]     Container: std::vector (size: 24)
+          Param
+            Primitive: int32_t
+          Underlying
+[2]         Class: vector (size: 24)
+              Member: __impl__ (offset: 0)
+                Primitive: StubbedPointer
+              Member: __impl__ (offset: 8)
+                Primitive: StubbedPointer
+              Member: __impl__ (offset: 16)
+                Primitive: StubbedPointer
+)",
+                R"(
+[0] Class: MyClass_0 [MyClass] (size: 24, align: 8)
+      Member: __oi_parent_0 [__oi_parent] (offset: 0, align: 8)
+[1]     Container: std::vector<int32_t> (size: 24, align: 8)
+          Param
+            Primitive: int32_t
+)");
+}
+
+TEST(CodeGenTest, InheritFromContainerCompat) {
+  OICodeGen::Config config;
+  testTransform(config,
+                R"(
+[0] Class: MyClass (size: 24)
+      Parent (offset: 0)
+[1]     Container: std::vector (size: 24)
+          Param
+            Primitive: int32_t
+          Underlying
+[2]         Class: vector (size: 24)
+              Member: __impl__ (offset: 0)
+                Primitive: StubbedPointer
+              Member: __impl__ (offset: 8)
+                Primitive: StubbedPointer
+              Member: __impl__ (offset: 16)
+                Primitive: StubbedPointer
+)",
+                R"(
+[0] Class: MyClass_0 [MyClass] (size: 24, align: 8)
+      Member: __oi_padding_0 (offset: 0)
+[3]     Array: [int8_t[24]] (length: 24)
+          Primitive: int8_t
+)");
+}

test/test_drgn_parser.cpp

@@ -10,6 +10,7 @@
 #include "oi/type_graph/Printer.h"
 #include "oi/type_graph/TypeGraph.h"
 #include "oi/type_graph/Types.h"
+#include "test/type_graph_utils.h"
 
 using namespace type_graph;
 
@@ -28,7 +29,8 @@ void DrgnParserTest::TearDownTestSuite() {
 
 DrgnParser DrgnParserTest::getDrgnParser(TypeGraph& typeGraph,
                                          DrgnParserOptions options) {
-  DrgnParser drgnParser{typeGraph, options};
+  static auto infos = getContainerInfos();
+  DrgnParser drgnParser{typeGraph, infos, options};
   return drgnParser;
 }
 
@@ -255,8 +257,8 @@ TEST_F(DrgnParserTest, InheritanceMultiple) {
 TEST_F(DrgnParserTest, Container) {
   testMultiCompilerGlob("oid_test_case_std_vector_int_empty",
                         R"(
-[13] Pointer
+[14] Pointer
-[0]    Class: vector<int, std::allocator<int> > (size: 24)
+[13]   Container: std::vector (size: 24)
          Param
            Primitive: int32_t
          Param
@@ -279,12 +281,18 @@ TEST_F(DrgnParserTest, Container) {
              Function: ~allocator
              Function: allocate
              Function: deallocate
-         Parent (offset: 0)
+         Underlying
+[0]        Class: vector<int, std::allocator<int> > (size: 24)
+             Param
+               Primitive: int32_t
+             Param
+               [1]
+             Parent (offset: 0)
 *
 )",
                         R"(
-[9]  Pointer
+[10] Pointer
-[0]    Class: vector<int, std::allocator<int> > (size: 24)
+[9]    Container: std::vector (size: 24)
          Param
            Primitive: int32_t
          Param
@@ -304,7 +312,13 @@ TEST_F(DrgnParserTest, Container) {
              Function: ~allocator
              Function: allocate
              Function: deallocate
-         Parent (offset: 0)
+         Underlying
+[0]        Class: vector<int, std::allocator<int> > (size: 24)
+             Param
+               Primitive: int32_t
+             Param
+               [1]
+             Parent (offset: 0)
 *
 )");
 }

test/test_flattener.cpp

@@ -799,6 +799,27 @@ TEST(FlattenerTest, ParentClassAndContainer) {
 )");
 }
 
+TEST(FlattenerTest, ContainerWithParent) {
+  // This is necessary to correctly calculate container alignment
+  test(Flattener::createPass(),
+       R"(
+[0] Container: std::vector (size: 24)
+      Underlying
+[1]     Class: vector (size: 24)
+          Parent (offset: 0)
+[2]         Class: Parent (size: 4)
+              Member: x (offset: 0)
+                Primitive: int32_t
+)",
+       R"(
+[0] Container: std::vector (size: 24)
+      Underlying
+[1]     Class: vector (size: 24)
+          Member: x (offset: 0)
+            Primitive: int32_t
+)");
+}
+
 TEST(FlattenerTest, AllocatorParamInParent) {
   test(Flattener::createPass(),
        R"(
@@ -857,8 +878,7 @@ TEST(FlattenerTest, AllocatorUnfixableNoParent) {
 )");
 }
 
-TEST(FlattenerTest, AllocatorUnfixableParentNotClass) {
+TEST(FlattenerTest, AllocatorParamInParentContainer) {
-  // This could be supported if need-be, we just don't do it yet
   test(Flattener::createPass(),
        R"(
 [0] Container: std::vector (size: 24)

test/test_identify_containers.cpp

@@ -1,194 +0,0 @@
-#include <gtest/gtest.h>
-
-#include "oi/ContainerInfo.h"
-#include "oi/type_graph/IdentifyContainers.h"
-#include "oi/type_graph/Types.h"
-#include "test/type_graph_utils.h"
-
-using namespace type_graph;
-
-namespace {
-void test(std::string_view input, std::string_view expectedAfter) {
-  ::test(IdentifyContainers::createPass(getContainerInfos()),
-         input,
-         expectedAfter);
-}
-};  // namespace
-
-TEST(IdentifyContainers, Container) {
-  test(R"(
-[0] Class: std::vector (size: 24)
-      Param
-        Primitive: int32_t
-      Member: a (offset: 0)
-        Primitive: int32_t
-)",
-       R"(
-[1] Container: std::vector (size: 24)
-      Param
-        Primitive: int32_t
-)");
-}
-
-TEST(IdentifyContainers, ContainerInClass) {
-  test(R"(
-[0] Class: MyClass (size: 0)
-      Param
-[1]     Class: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-      Parent (offset: 0)
-[2]     Class: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-      Member: a (offset: 0)
-[3]     Class: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-)",
-       R"(
-[0] Class: MyClass (size: 0)
-      Param
-[4]     Container: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-      Parent (offset: 0)
-[5]     Container: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-      Member: a (offset: 0)
-[6]     Container: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-)");
-}
-
-TEST(IdentifyContainers, ContainerInContainer) {
-  test(R"(
-[0] Class: std::vector (size: 24)
-      Param
-[1]     Class: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-)",
-       R"(
-[2] Container: std::vector (size: 24)
-      Param
-[3]     Container: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-)");
-}
-
-TEST(IdentifyContainers, ContainerInContainer2) {
-  test(R"(
-[0] Container: std::vector (size: 24)
-      Param
-[1]     Class: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-)",
-       R"(
-[0] Container: std::vector (size: 24)
-      Param
-[2]     Container: std::vector (size: 24)
-          Param
-            Primitive: int32_t
-)");
-}
-
-TEST(IdentifyContainers, ContainerInArray) {
-  test(R"(
-[0] Array: (length: 2)
-[1]   Class: std::vector (size: 24)
-        Param
-          Primitive: int32_t
-)",
-       R"(
-[0] Array: (length: 2)
-[2]   Container: std::vector (size: 24)
-        Param
-          Primitive: int32_t
-)");
-}
-
-TEST(IdentifyContainers, ContainerInTypedef) {
-  test(R"(
-[0] Typedef: MyAlias
-[1]   Class: std::vector (size: 24)
-        Param
-          Primitive: int32_t
-)",
-       R"(
-[0] Typedef: MyAlias
-[2]   Container: std::vector (size: 24)
-        Param
-          Primitive: int32_t
-)");
-}
-
-TEST(IdentifyContainers, ContainerInPointer) {
-  test(R"(
-[0] Pointer
-[1]   Class: std::vector (size: 24)
-        Param
-          Primitive: int32_t
-)",
-       R"(
-[0] Pointer
-[2]   Container: std::vector (size: 24)
-        Param
-          Primitive: int32_t
-)");
-}
-
-TEST(IdentifyContainers, ContainerDuplicate) {
-  test(R"(
-[0] Class: std::vector (size: 24)
-      Param
-        Primitive: int32_t
-      Member: a (offset: 0)
-        Primitive: int32_t
-    [0]
-)",
-       R"(
-[1] Container: std::vector (size: 24)
-      Param
-        Primitive: int32_t
-    [1]
-)");
-}
-
-TEST(IdentifyContainers, CycleClass) {
-  test(R"(
-[0] Class: ClassA (size: 0)
-      Member: x (offset: 0)
-[1]     Class: ClassB (size: 0)
-          Param
-            [0]
-)",
-       R"(
-[0] Class: ClassA (size: 0)
-      Member: x (offset: 0)
-[1]     Class: ClassB (size: 0)
-          Param
-            [0]
-)");
-}
-
-TEST(IdentifyContainers, CycleContainer) {
-  test(R"(
-[0] Class: ClassA (size: 0)
-      Member: x (offset: 0)
-[1]     Class: std::vector (size: 0)
-          Param
-            [0]
-)",
-       R"(
-[0] Class: ClassA (size: 0)
-      Member: x (offset: 0)
-[2]     Container: std::vector (size: 0)
-          Param
-            [0]
-)");
-}

test/test_prune.cpp

@@ -63,3 +63,33 @@ TEST(PruneTest, RecurseClassChild) {
 [1]     Class: ClassA (size: 12)
 )");
 }
+
+TEST(PruneTest, PruneContainer) {
+  test(Prune::createPass(),
+       R"(
+[0] Container: std::vector (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+        Value: "123"
+        Primitive: int32_t
+      Underlying
+[1]     Class: vector<int32_t> (size: 24)
+          Parent (offset: 0)
+[2]         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
+)",
+       R"(
+[0] Container: std::vector (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+        Value: "123"
+        Primitive: int32_t
+)");
+}

test/test_type_identifier.cpp

@@ -108,6 +108,12 @@ TEST(TypeIdentifierTest, PassThroughTypes) {
 [2]     Container: std::allocator (size: 1)
           Param
             Primitive: int32_t
+          Underlying
+[1]         Class: std::allocator (size: 1)
+              Param
+                Primitive: int32_t
+              Function: allocate
+              Function: deallocate
 )");
 }
 
@@ -137,6 +143,12 @@ TEST(TypeIdentifierTest, PassThroughSameType) {
 [2]     Container: std::allocator (size: 1)
           Param
             Primitive: int32_t
+          Underlying
+[1]         Class: std::allocator (size: 1)
+              Param
+                Primitive: int32_t
+              Function: allocate
+              Function: deallocate
       Param
         [2]
 )");

test/type_graph_utils.cpp

@@ -86,22 +86,23 @@ std::vector<std::unique_ptr<ContainerInfo>> getContainerInfos() {
 Container getVector(NodeId id) {
   static ContainerInfo info{"std::vector", SEQ_TYPE, "vector"};
   info.stubTemplateParams = {1};
-  return Container{id, info, 24};
+  return Container{id, info, 24, nullptr};
 }
 
 Container getMap(NodeId id) {
   static ContainerInfo info{"std::map", STD_MAP_TYPE, "map"};
   info.stubTemplateParams = {2, 3};
-  return Container{id, info, 48};
+  return Container{id, info, 48, nullptr};
 }
 
 Container getList(NodeId id) {
   static ContainerInfo info{"std::list", LIST_TYPE, "list"};
   info.stubTemplateParams = {1};
-  return Container{id, info, 24};
+  return Container{id, info, 24, nullptr};
 }
 
 Container getPair(NodeId id) {
   static ContainerInfo info{"std::pair", PAIR_TYPE, "utility"};
-  return Container{id, info, 8};  // Nonsense size, shouldn't matter for tests
+  return Container{
+      id, info, 8, nullptr};  // Nonsense size, shouldn't matter for tests
 }