Make KeyCapture work with nested typedefs (#473)

oi/CodeGen.cpp

@@ -83,6 +83,10 @@ namespace {
 std::vector<std::string_view> enumerateTypeNames(Type& type) {
   std::vector<std::string_view> names;
   Type* t = &type;
+
+  if (const auto* ck = dynamic_cast<CaptureKeys*>(t))
+    t = &ck->underlyingType();
+
   while (const Typedef* td = dynamic_cast<Typedef*>(t)) {
     names.emplace_back(t->inputName());
     t = &td->underlyingType();
@@ -1152,9 +1156,14 @@ void CodeGen::addTypeHandlers(const TypeGraph& typeGraph, std::string& code) {
       genContainerTypeHandler(
           definedContainers_, con->containerInfo_, con->templateParams, code);
     } else if (const auto* cap = dynamic_cast<const CaptureKeys*>(&t)) {
+      auto* container =
+          dynamic_cast<Container*>(&(stripTypedefs(cap->underlyingType())));
+      if (!container)
+        throw std::runtime_error("KaptureKeys requires a container");
+
       genContainerTypeHandler(definedContainers_,
                               cap->containerInfo(),
-                              cap->container().templateParams,
+                              container->templateParams,
                               code);
     }
   }

oi/type_graph/KeyCapture.cpp

@@ -99,7 +99,7 @@ void KeyCapture::visit(Class& c) {
  *         [VAL]
  */
 Type& KeyCapture::captureKey(Type& type) {
-  auto* container = dynamic_cast<Container*>(&type);
+  auto* container = dynamic_cast<Container*>(&(stripTypedefs(type)));
   if (!container)  // We only want to capture keys from containers
     return type;
 
@@ -114,7 +114,9 @@ Type& KeyCapture::captureKey(Type& type) {
   auto infoPtr = std::make_unique<ContainerInfo>(std::move(newContainerInfo));
   const auto& info = containerInfos_.emplace_back(std::move(infoPtr));
 
-  auto& captureKeysNode = typeGraph_.makeType<CaptureKeys>(*container, *info);
+  // auto& captureKeysNode = typeGraph_.makeType<CaptureKeys>(*container,
+  // *info);
+  auto& captureKeysNode = typeGraph_.makeType<CaptureKeys>(type, *info);
   return captureKeysNode;
 }
 

oi/type_graph/Printer.cpp

@@ -183,7 +183,7 @@ void Printer::visit(const DummyAllocator& d) {
 void Printer::visit(const CaptureKeys& d) {
   prefix();
   out_ << "CaptureKeys" << std::endl;
-  print(d.container());
+  print(d.underlyingType());
 }
 
 void Printer::prefix() {

oi/type_graph/TopoSorter.cpp

@@ -150,7 +150,7 @@ void TopoSorter::visit(Reference& r) {
 }
 
 void TopoSorter::visit(CaptureKeys& c) {
-  accept(c.container());
+  accept(c.underlyingType());
   sortedTypes_.push_back(c);
 }
 

oi/type_graph/Types.h

@@ -79,6 +79,9 @@ class ConstVisitor;
   Type& accept(Visitor<Type&>& m) override;           \
   void accept(ConstVisitor& v) const override;
 
+class Type;
+Type& stripTypedefs(Type& type);
+
 // TODO delete copy and move ctors
 
 /*
@@ -934,8 +937,8 @@ class DummyAllocator : public Type {
  */
 class CaptureKeys : public Type {
  public:
-  explicit CaptureKeys(Container& c, const ContainerInfo& info)
+  explicit CaptureKeys(Type& t, const ContainerInfo& info)
-      : container_(c), containerInfo_(info) {
+      : underlyingType_(t), containerInfo_(info) {
     regenerateName();
   }
 
@@ -948,19 +951,19 @@ class CaptureKeys : public Type {
   }
 
   virtual std::string_view inputName() const override {
-    return container_.get().inputName();
+    return underlyingType_.get().inputName();
   }
 
   void regenerateName() {
-    name_ = "OICaptureKeys<" + container_.get().name() + ">";
+    name_ = "OICaptureKeys<" + underlyingType_.get().name() + ">";
   }
 
   virtual size_t size() const override {
-    return container_.get().size();
+    return underlyingType_.get().size();
   }
 
   virtual uint64_t align() const override {
-    return container_.get().align();
+    return underlyingType_.get().align();
   }
 
   virtual NodeId id() const override {
@@ -971,22 +974,20 @@ class CaptureKeys : public Type {
     return containerInfo_;
   }
 
-  Container& container() const {
+  Type& underlyingType() const {
-    return container_;
+    return underlyingType_;
   }
 
-  void setContainer(Container& container) {
+  void setUnderlyingType(Type& t) {
-    container_ = container;
+    underlyingType_ = t;
   }
 
  private:
-  std::reference_wrapper<Container> container_;
+  std::reference_wrapper<Type> underlyingType_;
   const ContainerInfo& containerInfo_;
   std::string name_;
 };
 
-Type& stripTypedefs(Type& type);
-
 }  // namespace oi::detail::type_graph
 
 #undef DECLARE_ACCEPT

oi/type_graph/Visitor.h

@@ -117,7 +117,7 @@ class RecursiveVisitor : public Visitor<void> {
     accept(d.allocType());
   }
   virtual void visit(CaptureKeys& c) {
-    accept(c.container());
+    accept(c.underlyingType());
   }
 };
 
@@ -190,8 +190,7 @@ class RecursiveMutator : public Visitor<Type&> {
     return d;
   }
   virtual Type& visit(CaptureKeys& c) {
-    auto& newContainer = dynamic_cast<Container&>(mutate(c.container()));
+    c.setUnderlyingType(mutate(c.underlyingType()));
-    c.setContainer(newContainer);
     return c;
   }
 };

test/integration/capture_keys.toml

@@ -8,6 +8,11 @@ struct MapHolder {
   std::map<int, int> dontCaptureKeysHere;
 };
 
+using MyIntMap = std::map<int,int>;
+struct MyTypedefIntMapHolder {
+  MyIntMap captureMyKeys;
+};
+
 enum class MyEnum {
   One = 1,
   Two = 2,
@@ -320,6 +325,48 @@ class FixedAllocator {
       ]
     }]'''
 
+
+  [cases.typedef]
+    oid_skip = "Requires TreeBuilderV2"
+    param_types = ["const MyTypedefIntMapHolder&"]
+    setup = '''return {MyTypedefIntMapHolder{
+      .captureMyKeys{ {1,2},{3,4} },
+    }};'''
+    config_suffix = '''
+      [[codegen.capture_keys]]
+      type = "MyTypedefIntMapHolder"
+      members = ["captureMyKeys"]
+    '''
+    expect_json_v2 = '''[{
+      "members": [
+        {
+          "name": "captureMyKeys",
+          "typePath": ["a0","captureMyKeys"],
+          "typeNames": ["MyIntMap", "std::map<int32_t, int32_t, less<int>, std::allocator<std::pair<int32_t const, int32_t>>>"],
+          "members": [
+            {
+              "name": "[1]",
+              "typePath": ["a0","captureMyKeys","[1]"],
+              "data": 1,
+              "members": [
+                {"name": "key", "typePath": ["a0","captureMyKeys","[1]","key"]},
+                {"name": "value", "typePath": ["a0","captureMyKeys","[1]","value"]}
+              ]
+            },
+            {
+              "name": "[3]",
+              "typePath": ["a0","captureMyKeys","[3]"],
+              "data": 3,
+              "members": [
+                {"name": "key", "typePath": ["a0","captureMyKeys","[3]","key"]},
+                {"name": "value", "typePath": ["a0","captureMyKeys","[3]","value"]}
+              ]
+            }
+          ]
+        }
+      ]
+    }]'''
+
   [cases.std_unordered_map]
     oid_skip = "Requires TreeBuilderV2"
     param_types = ["const std::unordered_map<int, int>&"]

test/test_key_capture.cpp

@@ -76,6 +76,35 @@ TEST(KeyCaptureTest, MapInMap) {
 )");
 }
 
+TEST(KeyCaptureTest, Typedef) {
+  std::vector<OICodeGen::Config::KeyToCapture> keysToCapture = {
+      {"MyClass", "a"},
+  };
+  std::vector<std::unique_ptr<ContainerInfo>> containerInfos;
+  test(KeyCapture::createPass(keysToCapture, containerInfos),
+       R"(
+[0] Class: MyClass (size: 12)
+      Member: a (offset: 8)
+[1]     Typedef: MyTypeDef 
+[2]       Container: std::map (size: 24)
+            Param
+              Primitive: int32_t
+            Param
+              Primitive: int32_t
+)",
+       R"(
+[0] Class: MyClass (size: 12)
+      Member: a (offset: 8)
+        CaptureKeys
+[1]       Typedef: MyTypeDef 
+[2]         Container: std::map (size: 24)
+              Param
+                Primitive: int32_t
+              Param
+                Primitive: int32_t
+)");
+}
+
 TEST(KeyCaptureTest, TopLevel) {
   std::vector<OICodeGen::Config::KeyToCapture> keysToCapture = {
       {{}, {}, true},