From b88dabae5c2bda074f16de79b57ffaf5be1d275e Mon Sep 17 00:00:00 2001
From: Jake Hillion <jakehillion@meta.com>
Date: Fri, 14 Jul 2023 09:59:01 -0700
Subject: [PATCH] CycleFinder: support containers

---
 .circleci/config.yml          |  2 +-
 oi/type_graph/CycleFinder.cpp |  6 +++++-
 oi/type_graph/Types.h         |  4 ++++
 test/TypeGraphParser.cpp      |  4 ++++
 test/test_cycle_finder.cpp    | 34 ++++++++++++++++++++++++++++++++++
 5 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index da0d037..82e10f2 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -27,7 +27,7 @@ workflows:
             - build-gcc
           oid_test_args: "-ftyped-data-segment"
           tests_regex: "OidIntegration\\..*"
-          exclude_regex: ".*inheritance_polymorphic.*|.*pointers.*|.*arrays_member_int0|OidIntegration.unions.*|.*thrift_unions_dynamic_.*|.*cycles_unique_ptr|.*cycles_shared_ptr"
+          exclude_regex: ".*inheritance_polymorphic.*|.*pointers.*|.*arrays_member_int0|OidIntegration.unions.*|.*thrift_unions_dynamic_.*"
       - test:
           name: test-tree-builder-type-checking-gcc
           requires:
diff --git a/oi/type_graph/CycleFinder.cpp b/oi/type_graph/CycleFinder.cpp
index a741c59..8dfedf1 100644
--- a/oi/type_graph/CycleFinder.cpp
+++ b/oi/type_graph/CycleFinder.cpp
@@ -38,7 +38,11 @@ void CycleFinder::accept(Type& type) {
     if (auto* p = dynamic_cast<Pointer*>(&from)) {
       p->setPointeeType(edge);
     } else if (auto* c = dynamic_cast<Container*>(&from)) {
-      throw std::runtime_error("TODO: cycle is from a Container");
+      for (auto& param : c->templateParams) {
+        if (param.type() == &type) {
+          param.setType(&edge);
+        }
+      }
     } else {
       throw std::logic_error(
           "all cycles are expected to come from a Pointer or Container");
diff --git a/oi/type_graph/Types.h b/oi/type_graph/Types.h
index 574b49f..e70cbd7 100644
--- a/oi/type_graph/Types.h
+++ b/oi/type_graph/Types.h
@@ -158,6 +158,10 @@ class TemplateParam {
   TemplateParam(std::string value) : value(std::move(value)) {
   }
 
+  void setType(Type* type) {
+    type_ = type;
+  }
+
   Type* type() const {
     return type_;
   }
diff --git a/test/TypeGraphParser.cpp b/test/TypeGraphParser.cpp
index c9a78b1..058c709 100644
--- a/test/TypeGraphParser.cpp
+++ b/test/TypeGraphParser.cpp
@@ -72,6 +72,10 @@ ContainerInfo& getContainerInfo(std::string_view name) {
     static ContainerInfo info{"std::allocator", DUMMY_TYPE, "memory"};
     return info;
   }
+  if (name == "std::unique_ptr") {
+    static ContainerInfo info{"std::unique_ptr", UNIQ_PTR_TYPE, "memory"};
+    return info;
+  }
   throw TypeGraphParserError{"Unsupported container: " + std::string{name}};
 }
 
diff --git a/test/test_cycle_finder.cpp b/test/test_cycle_finder.cpp
index 2fcf900..c490f44 100644
--- a/test/test_cycle_finder.cpp
+++ b/test/test_cycle_finder.cpp
@@ -59,3 +59,37 @@ TEST(CycleFinderTest, RawPointer) {
             [0]
 )");
 }
+
+TEST(CycleFinderTest, UniquePointer) {
+  // Original:
+  //   class UniqueNode {
+  //     int value;
+  //     std::unique_ptr<struct UniqueNode> next;
+  //   };
+  //
+  // After:
+  //   class std__unique_ptr_UniqueNode_UniqueNode;
+  //   class UniqueNode {
+  //     int value;
+  //     std::unique_ptr<std__unique_ptr_UniqueNode_UniqueNode> next;
+  //   };
+  test(CycleFinder::createPass(), R"(
+[0] Struct: UniqueNode (size: 16)
+      Member: value (offset: 0)
+        Primitive: int32_t
+      Member: next (offset: 8)
+[1]     Container: std::unique_ptr (size: 8)
+          Param
+            [0]
+)",
+       R"(
+[0] Struct: UniqueNode (size: 16)
+      Member: value (offset: 0)
+        Primitive: int32_t
+      Member: next (offset: 8)
+[1]     Container: std::unique_ptr (size: 8)
+          Param
+[2]         CycleBreaker
+              [0]
+)");
+}