TypeGraph: Add CaptureKeys node

Author: ajor

oi/type_graph/NameGen.cpp

@@ -199,4 +199,9 @@ void NameGen::visit(DummyAllocator& d) {
   d.regenerateName();
 }
 
+void NameGen::visit(CaptureKeys& c) {
+  RecursiveVisitor::visit(c);
+  c.regenerateName();
+}
+
 }  // namespace oi::detail::type_graph

oi/type_graph/NameGen.h

@@ -47,6 +47,7 @@ class NameGen final : public RecursiveVisitor {
   void visit(Typedef& td) override;
   void visit(Pointer& p) override;
   void visit(DummyAllocator& d) override;
+  void visit(CaptureKeys& d) override;
 
   static const inline std::string AnonPrefix = "__oi_anon";
 

oi/type_graph/Printer.cpp

@@ -156,6 +156,12 @@ void Printer::visit(const DummyAllocator& d) {
   print(d.allocType());
 }
 
+void Printer::visit(const CaptureKeys& d) {
+  prefix();
+  out_ << "CaptureKeys" << std::endl;
+  print(d.container());
+}
+
 void Printer::prefix() {
   int indent = baseIndent_ + depth_ * 2;
   out_ << std::string(indent, ' ');

oi/type_graph/Printer.h

@@ -41,6 +41,7 @@ class Printer : public ConstVisitor {
   void visit(const Pointer& p) override;
   void visit(const Dummy& d) override;
   void visit(const DummyAllocator& d) override;
+  void visit(const CaptureKeys& d) override;
 
  private:
   void prefix();

oi/type_graph/TopoSorter.cpp

@@ -136,6 +136,11 @@ void TopoSorter::visit(Pointer& p) {
   acceptAfter(p.pointeeType());
 }
 
+void TopoSorter::visit(CaptureKeys& c) {
+  accept(c.container());
+  sortedTypes_.push_back(c);
+}
+
 /*
  * A type graph may contain cycles, so we need to slightly tweak the standard
  * topological sorting algorithm. Cycles can only be introduced by certain

oi/type_graph/TopoSorter.h

@@ -47,6 +47,7 @@ class TopoSorter : public RecursiveVisitor {
   void visit(Typedef& td) override;
   void visit(Pointer& p) override;
   void visit(Primitive& p) override;
+  void visit(CaptureKeys& p) override;
 
  private:
   std::unordered_set<Type*> visited_;

oi/type_graph/Types.h

@@ -48,7 +48,8 @@
   X(Typedef)         \
   X(Pointer)         \
   X(Dummy)           \
-  X(DummyAllocator)
+  X(DummyAllocator)  \
+  X(CaptureKeys)
 
 struct ContainerInfo;
 
@@ -730,6 +731,60 @@ class DummyAllocator : public Type {
   std::string inputName_;
 };
 
+/*
+ * CaptureKeys
+ *
+ * The held Container will have its keys captured.
+ */
+class CaptureKeys : public Type {
+ public:
+  explicit CaptureKeys(Container& c, const ContainerInfo& info)
+      : container_(c), containerInfo_(info) {
+    regenerateName();
+  }
+
+  static inline constexpr bool has_node_id = false;
+
+  DECLARE_ACCEPT
+
+  virtual const std::string& name() const override {
+    return name_;
+  }
+
+  virtual std::string_view inputName() const override {
+    return container_.inputName();
+  }
+
+  void regenerateName() {
+    name_ = "OICaptureKeys<" + container_.name() + ">";
+  }
+
+  virtual size_t size() const override {
+    return container_.size();
+  }
+
+  virtual uint64_t align() const override {
+    return container_.align();
+  }
+
+  virtual NodeId id() const override {
+    return -1;
+  }
+
+  const ContainerInfo& containerInfo() const {
+    return containerInfo_;
+  }
+
+  Container& container() const {
+    return container_;
+  }
+
+ private:
+  Container& container_;
+  const ContainerInfo& containerInfo_;
+  std::string name_;
+};
+
 Type& stripTypedefs(Type& type);
 
 }  // namespace oi::detail::type_graph

oi/type_graph/Visitor.h

@@ -108,6 +108,9 @@ class RecursiveVisitor : public Visitor {
   virtual void visit(DummyAllocator& d) {
     accept(d.allocType());
   }
+  virtual void visit(CaptureKeys& c) {
+    accept(c.container());
+  }
 };
 
 /*

test/test_name_gen.cpp

@@ -399,6 +399,31 @@ TEST(NameGenTest, DummyAllocator) {
   EXPECT_EQ(myalloc.inputName(), "BigAllocator");
 }
 
+TEST(NameGenTest, CaptureKeys) {
+  auto myparam1 = Class{0, Class::Kind::Struct, "MyParam", 13};
+  auto myparam2 = Class{1, Class::Kind::Struct, "MyParam", 13};
+
+  auto mycontainer = getVector();
+  mycontainer.templateParams.push_back(myparam1);
+  mycontainer.templateParams.push_back(myparam2);
+
+  auto captureKeys = CaptureKeys{mycontainer, mycontainer.containerInfo_};
+
+  NameGen nameGen;
+  nameGen.generateNames({captureKeys});
+
+  EXPECT_EQ(myparam1.name(), "MyParam_0");
+  EXPECT_EQ(myparam2.name(), "MyParam_1");
+  EXPECT_EQ(mycontainer.name(), "std::vector<MyParam_0, MyParam_1>");
+  EXPECT_EQ(captureKeys.name(),
+            "OICaptureKeys<std::vector<MyParam_0, MyParam_1>>");
+
+  EXPECT_EQ(myparam1.inputName(), "MyParam");
+  EXPECT_EQ(myparam2.inputName(), "MyParam");
+  EXPECT_EQ(mycontainer.inputName(), "std::vector<MyParam, MyParam>");
+  EXPECT_EQ(captureKeys.inputName(), "std::vector<MyParam, MyParam>");
+}
+
 TEST(NameGenTest, Cycle) {
   auto classA = Class{0, Class::Kind::Class, "ClassA", 69};
   auto classB = Class{1, Class::Kind::Class, "ClassB", 69};

test/test_topo_sorter.cpp

@@ -299,3 +299,20 @@ MyStruct
 MyClass
 )");
 }
+
+TEST(TopoSorterTest, CaptureKeys) {
+  auto myparam1 = Class{1, Class::Kind::Struct, "MyParam1", 13};
+  auto myparam2 = Class{2, Class::Kind::Struct, "MyParam2", 13};
+  auto mycontainer = getMap();
+  mycontainer.templateParams.push_back((myparam1));
+  mycontainer.templateParams.push_back((myparam2));
+
+  auto captureKeys = CaptureKeys{mycontainer, mycontainer.containerInfo_};
+
+  test({captureKeys}, R"(
+MyParam1
+MyParam2
+std::map
+OICaptureKeys<std::map>
+)");
+}