From 784b900218822a5de706f5035baf5273a1ff7ac1 Mon Sep 17 00:00:00 2001
From: Alastair Robertson <ajor@meta.com>
Date: Wed, 31 May 2023 06:49:47 -0700
Subject: [PATCH] TypeGraph: Replace allocators with DummyAllocator

When we were previously removing allocators, we were only able to work
with containers whose allocators appeared as their last template
parameter.

Now we can replace allocators in the middle of a parameter list.

This fixes tests for folly::sorted_vector_set.
---
 oi/OITraceCode.cpp               |  27 ++++++++
 oi/type_graph/Printer.cpp        |   5 +-
 oi/type_graph/TypeIdentifier.cpp |  55 ++++++++--------
 oi/type_graph/Types.h            |   6 +-
 test/test_name_gen.cpp           |  13 +++-
 test/test_type_identifier.cpp    | 104 ++++++++++++++++++++++++++++++-
 types/cxx11_string_type.toml     |   8 +--
 types/seq_type.toml              |   8 +--
 types/string_type.toml           |   8 +--
 9 files changed, 186 insertions(+), 48 deletions(-)

diff --git a/oi/OITraceCode.cpp b/oi/OITraceCode.cpp
index b6323b6..ace2ece 100644
--- a/oi/OITraceCode.cpp
+++ b/oi/OITraceCode.cpp
@@ -155,3 +155,30 @@ struct alignas(align) DummySizedOperator {
 // DummySizedOperator<0,0> also collapses to this
 template <>
 struct DummySizedOperator<0> {};
+
+template <template <typename, size_t, size_t> typename DerivedT,
+          typename T,
+          size_t N,
+          size_t Align>
+struct DummyAllocatorBase {
+  using value_type = T;
+  T* allocate(std::size_t n) {
+    return nullptr;
+  }
+  void deallocate(T* p, std::size_t n) noexcept {
+  }
+
+  template <typename U>
+  struct rebind {
+    using other = DerivedT<U, N, Align>;
+  };
+};
+
+template <typename T, size_t N, size_t Align = 0>
+struct alignas(Align) DummyAllocator
+    : DummyAllocatorBase<DummyAllocator, T, N, Align> {
+  char c[N];
+};
+
+template <typename T>
+struct DummyAllocator<T, 0> : DummyAllocatorBase<DummyAllocator, T, 0, 0> {};
diff --git a/oi/type_graph/Printer.cpp b/oi/type_graph/Printer.cpp
index b4050cd..e2bf652 100644
--- a/oi/type_graph/Printer.cpp
+++ b/oi/type_graph/Printer.cpp
@@ -115,8 +115,9 @@ void Printer::visit(const Dummy& d) {
 
 void Printer::visit(const DummyAllocator& d) {
   prefix();
-  out_ << "DummyAllocatorTODO (size: " << d.size() << align_str(d.align())
-       << ")" << std::endl;
+  out_ << "DummyAllocator (size: " << d.size() << align_str(d.align()) << ")"
+       << std::endl;
+  print(d.allocType());
 }
 
 bool Printer::prefix(const Type* type) {
diff --git a/oi/type_graph/TypeIdentifier.cpp b/oi/type_graph/TypeIdentifier.cpp
index 33c382c..e51f1fb 100644
--- a/oi/type_graph/TypeIdentifier.cpp
+++ b/oi/type_graph/TypeIdentifier.cpp
@@ -57,47 +57,48 @@ bool isAllocator(Type* t) {
 }  // namespace
 
 void TypeIdentifier::visit(Container& c) {
-  // TODO will containers exist at this point?
-  // maybe don't need this function
-
-  //  auto *c = make_type<Container>(containerInfo);
-
   const auto& stubParams = c.containerInfo_.stubTemplateParams;
   // TODO these two arrays could be looped over in sync for better performance
   for (size_t i = 0; i < c.templateParams.size(); i++) {
+    const auto& param = c.templateParams[i];
+    if (dynamic_cast<Dummy*>(param.type) ||
+        dynamic_cast<DummyAllocator*>(param.type)) {
+      // In case the TypeIdentifier pass is run multiple times, we don't want to
+      // replace dummies again as the context of the original replacement has
+      // been lost.
+      continue;
+    }
+
     if (std::find(stubParams.begin(), stubParams.end(), i) !=
         stubParams.end()) {
-      const auto& param = c.templateParams[i];
-      if (isAllocator(param.type)) {
-        //        auto *allocator = dynamic_cast<Class*>(param.type); // TODO
-        //        please don't do this... auto &typeToAllocate =
-        //        *allocator->templateParams.at(0).type; auto *dummy =
-        //        make_type<DummyAllocator>(typeToAllocate, param.type->size(),
-        //        param.type->align()); c.templateParams[i] = dummy;
+      size_t size = param.type->size();
+      if (size == 1) {
+        // Hack: when we get a reported size of 1 for these parameters, it
+        // turns out that a size of 0 is actually expected.
+        size = 0;
+      }
 
-        // TODO allocators are tricky... just remove them entirely for now
-        // The problem is a std::map<int, int> requires an allocator of type
-        // std::allocator<std::pair<const int, int>>, but we do not record
-        // constness of types.
-        if (i != c.templateParams.size() - 1) {
-          throw std::runtime_error("Unsupported allocator parameter");
+      if (isAllocator(param.type)) {
+        auto* allocator =
+            dynamic_cast<Class*>(param.type);  // TODO please don't do this...
+        Type* typeToAllocate;
+        if (!allocator->templateParams.empty()) {
+          typeToAllocate = allocator->templateParams.at(0).type;
+        } else {
+          // We've encountered some bad DWARF. Let's guess that the type to
+          // allocate is the first parameter of the container.
+          typeToAllocate = c.templateParams[0].type;
         }
-        c.templateParams.erase(c.templateParams.begin() + i,
-                               c.templateParams.end());
+        auto* dummy = typeGraph_.make_type<DummyAllocator>(
+            *typeToAllocate, size, param.type->align());
+        c.templateParams[i] = dummy;
       } else {
-        size_t size = param.type->size();
-        if (size == 1) {  // TODO this is a hack
-          size = 0;
-        }
         auto* dummy = typeGraph_.make_type<Dummy>(size, param.type->align());
         c.templateParams[i] = dummy;
       }
     }
   }
 
-  // TODO replace current node with "c" (if we want to transform classes into
-  // containers here)
-
   for (const auto& param : c.templateParams) {
     visit(*param.type);
   }
diff --git a/oi/type_graph/Types.h b/oi/type_graph/Types.h
index 1075d96..55fe572 100644
--- a/oi/type_graph/Types.h
+++ b/oi/type_graph/Types.h
@@ -410,10 +410,8 @@ class DummyAllocator : public Type {
   DECLARE_ACCEPT
 
   virtual std::string name() const override {
-    return "std::allocator<" + type_.name() + ">";
-    // TODO custom sized allocators:
-    //    return "DummyAllocator<" + type_.name() + ", " + std::to_string(size_)
-    //    + "," + std::to_string(align_) + ">";
+    return "DummyAllocator<" + type_.name() + ", " + std::to_string(size_) +
+           ", " + std::to_string(align_) + ">";
   }
 
   virtual size_t size() const override {
diff --git a/test/test_name_gen.cpp b/test/test_name_gen.cpp
index 44ba065..d19f826 100644
--- a/test/test_name_gen.cpp
+++ b/test/test_name_gen.cpp
@@ -242,6 +242,15 @@ TEST(NameGenTest, Pointer) {
   EXPECT_EQ(mypointer->name(), "std::vector<MyParam_0, MyParam_1>*");
 }
 
+TEST(NameGenTest, Dummy) {
+  auto dummy = std::make_unique<Dummy>(12, 34);
+
+  NameGen nameGen;
+  nameGen.generateNames({*dummy});
+
+  EXPECT_EQ(dummy->name(), "DummySizedOperator<12, 34>");
+}
+
 TEST(NameGenTest, DummyAllocator) {
   auto myparam1 = std::make_unique<Class>(Class::Kind::Struct, "MyParam", 13);
   auto myparam2 = std::make_unique<Class>(Class::Kind::Struct, "MyParam", 13);
@@ -250,7 +259,7 @@ TEST(NameGenTest, DummyAllocator) {
   mycontainer.templateParams.push_back(myparam1.get());
   mycontainer.templateParams.push_back(myparam2.get());
 
-  auto myalloc = std::make_unique<DummyAllocator>(mycontainer, 0, 0);
+  auto myalloc = std::make_unique<DummyAllocator>(mycontainer, 12, 34);
 
   NameGen nameGen;
   nameGen.generateNames({*myalloc});
@@ -259,7 +268,7 @@ TEST(NameGenTest, DummyAllocator) {
   EXPECT_EQ(myparam2->name(), "MyParam_1");
   EXPECT_EQ(mycontainer.name(), "std::vector<MyParam_0, MyParam_1>");
   EXPECT_EQ(myalloc->name(),
-            "std::allocator<std::vector<MyParam_0, MyParam_1>>");
+            "DummyAllocator<std::vector<MyParam_0, MyParam_1>, 12, 34>");
 }
 
 TEST(NameGenTest, Cycle) {
diff --git a/test/test_type_identifier.cpp b/test/test_type_identifier.cpp
index 39816f4..646391b 100644
--- a/test/test_type_identifier.cpp
+++ b/test/test_type_identifier.cpp
@@ -1,7 +1,109 @@
 #include <gtest/gtest.h>
 
+#include "oi/type_graph/Printer.h"
 #include "oi/type_graph/TypeIdentifier.h"
+#include "oi/type_graph/Types.h"
+#include "test/type_graph_utils.h"
 
 using namespace type_graph;
 
-// TODO tests!!!
+TEST(TypeIdentifierTest, StubbedParam) {
+  auto myint = Primitive{Primitive::Kind::Int32};
+
+  auto myparam = Class{Class::Kind::Struct, "MyParam", 4};
+  myparam.members.push_back(Member{&myint, "a", 0});
+
+  auto container = getVector();
+  container.templateParams.push_back(TemplateParam{&myint});
+  container.templateParams.push_back(TemplateParam{&myparam});
+  container.templateParams.push_back(TemplateParam{&myint});
+
+  test(TypeIdentifier::createPass(), {container}, R"(
+[0] Container: std::vector (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+[1]     Struct: MyParam (size: 4)
+          Member: a (offset: 0)
+            Primitive: int32_t
+      Param
+        Primitive: int32_t
+)",
+       R"(
+[0] Container: std::vector (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+        Dummy (size: 4)
+      Param
+        Primitive: int32_t
+)");
+}
+
+TEST(TypeIdentifierTest, Allocator) {
+  auto myint = Primitive{Primitive::Kind::Int32};
+
+  auto myalloc = Class{Class::Kind::Struct, "MyAlloc", 8};
+  myalloc.templateParams.push_back(TemplateParam{&myint});
+  myalloc.functions.push_back(Function{"allocate"});
+  myalloc.functions.push_back(Function{"deallocate"});
+
+  auto container = getVector();
+  container.templateParams.push_back(TemplateParam{&myint});
+  container.templateParams.push_back(TemplateParam{&myalloc});
+  container.templateParams.push_back(TemplateParam{&myint});
+
+  test(TypeIdentifier::createPass(), {container}, R"(
+[0] Container: std::vector (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+[1]     Struct: MyAlloc (size: 8)
+          Param
+            Primitive: int32_t
+          Function: allocate
+          Function: deallocate
+      Param
+        Primitive: int32_t
+)",
+       R"(
+[0] Container: std::vector (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+        DummyAllocator (size: 8)
+          Primitive: int32_t
+      Param
+        Primitive: int32_t
+)");
+}
+
+TEST(TypeIdentifierTest, AllocatorNoParam) {
+  auto myint = Primitive{Primitive::Kind::Int32};
+
+  auto myalloc = Class{Class::Kind::Struct, "MyAlloc", 8};
+  myalloc.functions.push_back(Function{"allocate"});
+  myalloc.functions.push_back(Function{"deallocate"});
+
+  auto container = getVector();
+  container.templateParams.push_back(TemplateParam{&myint});
+  container.templateParams.push_back(TemplateParam{&myalloc});
+
+  test(TypeIdentifier::createPass(), {container}, R"(
+[0] Container: std::vector (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+[1]     Struct: MyAlloc (size: 8)
+          Function: allocate
+          Function: deallocate
+)",
+       R"(
+[0] Container: std::vector (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+        DummyAllocator (size: 8)
+          Primitive: int32_t
+)");
+}
diff --git a/types/cxx11_string_type.toml b/types/cxx11_string_type.toml
index 85e3f6c..c7036da 100644
--- a/types/cxx11_string_type.toml
+++ b/types/cxx11_string_type.toml
@@ -12,13 +12,13 @@ replaceTemplateParamIndex = []
 
 [codegen]
 decl = """
-template<typename T>
-void getSizeType(const %1%<T> &container, size_t& returnArg);
+template<typename T, typename Traits, typename Allocator>
+void getSizeType(const %1%<T, Traits, Allocator> &container, size_t& returnArg);
 """
 
 func = """
-template<typename T>
-void getSizeType(const %1%<T> &container, size_t& returnArg)
+template<typename T, typename Traits, typename Allocator>
+void getSizeType(const %1%<T, Traits, Allocator> &container, size_t& returnArg)
 {
     SAVE_SIZE(sizeof(%1%<T>));
 
diff --git a/types/seq_type.toml b/types/seq_type.toml
index e0432dc..9880da6 100644
--- a/types/seq_type.toml
+++ b/types/seq_type.toml
@@ -12,13 +12,13 @@ allocatorIndex = 1
 
 [codegen]
 decl = """
-template<typename T>
-void getSizeType(const %1%<T> &container, size_t& returnArg);
+template<typename T, typename Allocator>
+void getSizeType(const %1%<T, Allocator> &container, size_t& returnArg);
 """
 
 func = """
-template<typename T>
-void getSizeType(const %1%<T> &container, size_t& returnArg)
+template<typename T, typename Allocator>
+void getSizeType(const %1%<T, Allocator> &container, size_t& returnArg)
 {
     SAVE_SIZE(sizeof(%1%<T>));
 
diff --git a/types/string_type.toml b/types/string_type.toml
index 007b73e..fa4a2ae 100644
--- a/types/string_type.toml
+++ b/types/string_type.toml
@@ -12,13 +12,13 @@ replaceTemplateParamIndex = []
 
 [codegen]
 decl = """
-template<typename T>
-void getSizeType(const %1%<T> &container, size_t& returnArg);
+template<typename T, typename Traits, typename Allocator>
+void getSizeType(const %1%<T, Traits, Allocator> &container, size_t& returnArg);
 """
 
 func = """
-template<typename T>
-void getSizeType(const %1%<T> &container, size_t& returnArg)
+template<typename T, typename Traits, typename Allocator>
+void getSizeType(const %1%<T, Traits, Allocator> &container, size_t& returnArg)
 {
     SAVE_SIZE(sizeof(%1%<T>));