TypeGraph: Add KeyCapture pass

oi/type_graph/CMakeLists.txt

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

oi/type_graph/KeyCapture.cpp

@@ -0,0 +1,122 @@
+/*
+ * 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 "KeyCapture.h"
+
+#include "TypeGraph.h"
+
+namespace oi::detail::type_graph {
+
+Pass KeyCapture::createPass(
+    const std::vector<OICodeGen::Config::KeyToCapture>& keysToCapture,
+    std::vector<std::unique_ptr<ContainerInfo>>& containerInfos) {
+  auto fn = [&keysToCapture, &containerInfos](TypeGraph& typeGraph,
+                                              NodeTracker& tracker) {
+    KeyCapture pass{tracker, typeGraph, keysToCapture, containerInfos};
+    pass.insertCaptureDataNodes(typeGraph.rootTypes());
+  };
+
+  return Pass("KeyCapture", fn);
+}
+
+/*
+ * This function should be used as the main entry point to this pass, to add
+ * special handling of top-level types.
+ */
+void KeyCapture::insertCaptureDataNodes(
+    std::vector<std::reference_wrapper<Type>>& types) {
+  for (const auto& keyToCapture : keysToCapture_) {
+    if (!keyToCapture.topLevel)
+      continue;
+
+    // Capture keys from all top-level types
+    for (size_t i = 0; i < types.size(); i++) {
+      types[i] = captureKey(types[i]);
+    }
+    break;
+  }
+
+  for (const auto& type : types) {
+    accept(type);
+  }
+}
+
+void KeyCapture::accept(Type& type) {
+  if (tracker_.visit(type))
+    return;
+
+  type.accept(*this);
+}
+
+void KeyCapture::visit(Class& c) {
+  for (const auto& keyToCapture : keysToCapture_) {
+    if (!keyToCapture.type.has_value() || c.name() != *keyToCapture.type)
+      continue;
+    if (!keyToCapture.member.has_value())
+      continue;
+    for (size_t i = 0; i < c.members.size(); i++) {
+      auto& member = c.members[i];
+      if (member.name != *keyToCapture.member)
+        continue;
+
+      member = Member{captureKey(member.type()), member};
+    }
+  }
+
+  RecursiveVisitor::visit(c);
+}
+
+/*
+ * captureKey
+ *
+ * If the given type is a container, insert a CaptureKey node above it.
+ * Otherwise, just return the container node unchanged.
+ *
+ * Before:
+ *   Container: std::map
+ *     Param
+ *       [KEY]
+ *     Param
+ *       [VAL]
+ *
+ * After:
+ *   CaptureKeys
+ *     Container: std::map
+ *       Param
+ *         [KEY]
+ *       Param
+ *         [VAL]
+ */
+Type& KeyCapture::captureKey(Type& type) {
+  auto* container = dynamic_cast<Container*>(&type);
+  if (!container)  // We only want to capture keys from containers
+    return type;
+
+  /*
+   * Create a copy of the container info for capturing keys.
+   * CodeGen and other places may deduplicate containers based on the container
+   * info object, so it is necessary to create a new one when we want different
+   * behaviour.
+   */
+  auto newContainerInfo = container->containerInfo_.clone();
+  newContainerInfo.captureKeys = true;
+  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);
+  return captureKeysNode;
+}
+
+}  // namespace oi::detail::type_graph

oi/type_graph/KeyCapture.h

@@ -0,0 +1,67 @@
+/*
+ * 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 <memory>
+#include <vector>
+
+#include "NodeTracker.h"
+#include "PassManager.h"
+#include "Types.h"
+#include "Visitor.h"
+#include "oi/OICodeGen.h"
+
+namespace oi::detail::type_graph {
+
+/*
+ * KeyCapture
+ *
+ * Marks containers for which the user has requested key-capture.
+ */
+class KeyCapture : public RecursiveVisitor {
+ public:
+  static Pass createPass(
+      const std::vector<OICodeGen::Config::KeyToCapture>& keysToCapture,
+      std::vector<std::unique_ptr<ContainerInfo>>& containerInfos);
+
+  KeyCapture(NodeTracker& tracker,
+             TypeGraph& typeGraph,
+             const std::vector<OICodeGen::Config::KeyToCapture>& keysToCapture,
+             std::vector<std::unique_ptr<ContainerInfo>>& containerInfos)
+      : tracker_(tracker),
+        typeGraph_(typeGraph),
+        keysToCapture_(keysToCapture),
+        containerInfos_(containerInfos) {
+  }
+
+  using RecursiveVisitor::accept;
+  using RecursiveVisitor::visit;
+
+  void insertCaptureDataNodes(std::vector<std::reference_wrapper<Type>>& types);
+  void visit(Class& c) override;
+
+ private:
+  NodeTracker& tracker_;
+  TypeGraph& typeGraph_;
+  const std::vector<OICodeGen::Config::KeyToCapture>& keysToCapture_;
+  std::vector<std::unique_ptr<ContainerInfo>>& containerInfos_;
+
+  void accept(Type& type) override;
+  Type& captureKey(Type& type);
+};
+
+}  // namespace oi::detail::type_graph

oi/type_graph/Types.h

@@ -109,6 +109,14 @@ class Member {
         bitsize(bitsize) {
   }
 
+  Member(Type& type, const Member& other)
+      : type_(type),
+        name(other.name),
+        inputName(other.inputName),
+        bitOffset(other.bitOffset),
+        bitsize(other.bitsize) {
+  }
+
   Type& type() const {
     return type_;
   }
@@ -299,6 +307,17 @@ class Container : public Type {
         id_(id) {
   }
 
+  Container(NodeId id,
+            const Container& other,
+            const ContainerInfo& containerInfo)
+      : templateParams(other.templateParams),
+        containerInfo_(containerInfo),
+        name_(other.name_),
+        inputName_(other.inputName_),
+        size_(other.size_),
+        id_(id) {
+  }
+
   static inline constexpr bool has_node_id = true;
 
   DECLARE_ACCEPT

test/CMakeLists.txt

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

test/test_key_capture.cpp

@@ -0,0 +1,111 @@
+#include <gtest/gtest.h>
+
+#include "oi/type_graph/KeyCapture.h"
+#include "oi/type_graph/Types.h"
+#include "test/type_graph_utils.h"
+
+using namespace type_graph;
+
+TEST(KeyCaptureTest, InClass) {
+  std::vector<OICodeGen::Config::KeyToCapture> keysToCapture = {
+      {"MyClass", "b"},
+  };
+  std::vector<std::unique_ptr<ContainerInfo>> containerInfos;
+  test(KeyCapture::createPass(keysToCapture, containerInfos), R"(
+[0] Class: MyClass (size: 12)
+      Member: a (offset: 0)
+        Primitive: int32_t
+      Member: b (offset: 8)
+[1]     Container: std::map (size: 24)
+          Param
+            Primitive: int32_t
+          Param
+            Primitive: int32_t
+      Member: c (offset: 8)
+        [1]
+)",
+       R"(
+[0] Class: MyClass (size: 12)
+      Member: a (offset: 0)
+        Primitive: int32_t
+      Member: b (offset: 8)
+        CaptureKeys
+[1]       Container: std::map (size: 24)
+            Param
+              Primitive: int32_t
+            Param
+              Primitive: int32_t
+      Member: c (offset: 8)
+        [1]
+)");
+}
+
+TEST(KeyCaptureTest, MapInMap) {
+  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]     Container: std::map (size: 24)
+          Param
+            Primitive: int32_t
+          Param
+[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]       Container: std::map (size: 24)
+            Param
+              Primitive: int32_t
+            Param
+[2]           Container: std::map (size: 24)
+                Param
+                  Primitive: int32_t
+                Param
+                  Primitive: int32_t
+)");
+}
+
+TEST(KeyCaptureTest, TopLevel) {
+  std::vector<OICodeGen::Config::KeyToCapture> keysToCapture = {
+      {{}, {}, true},
+  };
+  std::vector<std::unique_ptr<ContainerInfo>> containerInfos;
+  test(KeyCapture::createPass(keysToCapture, containerInfos), R"(
+[0] Container: std::map (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+        Primitive: int32_t
+)",
+       R"(
+    CaptureKeys
+[0]   Container: std::map (size: 24)
+        Param
+          Primitive: int32_t
+        Param
+          Primitive: int32_t
+)");
+}
+
+TEST(KeyCaptureTest, TopLevelNotCaptured) {
+  std::vector<OICodeGen::Config::KeyToCapture> keysToCapture = {
+      {"MyClass", "a"},
+  };
+  std::vector<std::unique_ptr<ContainerInfo>> containerInfos;
+  testNoChange(KeyCapture::createPass(keysToCapture, containerInfos), R"(
+[0] Container: std::map (size: 24)
+      Param
+        Primitive: int32_t
+      Param
+        Primitive: int32_t
+)");
+}