add exporters::TypeCheckingWalker

CMakeLists.txt

@@ -309,7 +309,10 @@ target_link_libraries(oicore
 )
 
 ### TreeBuilder
-add_library(treebuilder oi/TreeBuilder.cpp)
+add_library(treebuilder
+  oi/TreeBuilder.cpp
+  oi/exporters/TypeCheckingWalker.cpp
+)
 add_dependencies(treebuilder librocksdb)
 target_link_libraries(treebuilder
   ${rocksdb_BINARY_DIR}/librocksdb.a

include/oi/types/dy.h

@@ -93,6 +93,6 @@ class List {
   Dynamic element;
 };
 
-}  // namespace ObjectIntrospection
+}  // namespace ObjectIntrospection::types::dy
 
 #endif

oi/CodeGen.cpp

@@ -109,7 +109,7 @@ void addIncludes(const TypeGraph& typeGraph,
   if (features[Feature::TreeBuilderTypeChecking]) {
     includes.emplace("oi/types/dy.h");
 
-    code += "#define DEFINE_DESCRIBE 1\n"; // added before all includes
+    code += "#define DEFINE_DESCRIBE 1\n";  // added before all includes
   }
   if (features[Feature::JitTiming]) {
     includes.emplace("chrono");

oi/exporters/TypeCheckingWalker.cpp

@@ -0,0 +1,80 @@
+/*
+ * 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 "TypeCheckingWalker.h"
+
+#include <type_traits>
+
+template <class>
+inline constexpr bool always_false_v = false;
+
+namespace ObjectIntrospection {
+namespace exporters {
+
+std::optional<TypeCheckingWalker::Element> TypeCheckingWalker::advance() {
+  if (stack.empty()) {
+    return std::nullopt;
+  }
+
+  auto el = stack.top();
+  stack.pop();
+
+  return std::visit(
+      [this](auto&& r) -> std::optional<TypeCheckingWalker::Element> {
+        const auto& ty = r.get();
+        using T = std::decay_t<decltype(ty)>;
+
+        if constexpr (std::is_same_v<T, types::dy::Unit>) {
+          // Unit type - noop.
+          return advance();
+        } else if constexpr (std::is_same_v<T, types::dy::VarInt>) {
+          // VarInt type - pop one element and return as a `VarInt`.
+          return TypeCheckingWalker::VarInt{popFront()};
+        } else if constexpr (std::is_same_v<T, types::dy::Pair>) {
+          // Pair type - read all of left then all of right. Recurse to get the
+          // values.
+          stack.push(ty.second);
+          stack.push(ty.first);
+          return advance();
+        } else if constexpr (std::is_same_v<T, types::dy::List>) {
+          // List type - pop one element as the length of the list, and place
+          // that many of the element type on the stack. Return the value as a
+          // `ListLength`.
+          auto el = popFront();
+          for (size_t i = 0; i < el; i++) {
+            stack.push(ty.element);
+          }
+          return TypeCheckingWalker::ListLength{el};
+        } else if constexpr (std::is_same_v<T, types::dy::Sum>) {
+          // Sum type - pop one element as the index of the tagged union, and
+          // place the matching element type on the stack. Return the value as a
+          // `SumIndex`. Throw if the index is invalid.
+          auto el = popFront();
+          if (el >= ty.variants.size()) {
+            throw std::runtime_error(std::string("invalid sum index: got ") +
+                                     std::to_string(el) + ", max " +
+                                     std::to_string(ty.variants.size()));
+          }
+          stack.push(ty.variants[el]);
+          return TypeCheckingWalker::SumIndex{el};
+        } else {
+          static_assert(always_false_v<T>, "non-exhaustive visitor!");
+        }
+      },
+      el);
+}
+
+}  // namespace exporters
+}  // namespace ObjectIntrospection

oi/exporters/TypeCheckingWalker.h

@@ -0,0 +1,78 @@
+/*
+ * 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
+
+/*
+ * TypeCheckingWalker
+ *
+ * Walks a dynamic data segment type simultaneously with the contents of the
+ * data segment, providing context to each extracted element.
+ *
+ * The dynamic types are implemented as a stack machine. Handling a type
+ * pops the top element, handles it which may involve pushing more types onto
+ * the stack, and returns an element from the data segment if appropriate or
+ * recurses. See the implementation for details.
+ */
+
+#include <cstdint>
+#include <optional>
+#include <span>
+#include <stack>
+#include <stdexcept>
+#include <variant>
+
+#include "oi/types/dy.h"
+
+namespace ObjectIntrospection {
+namespace exporters {
+
+class TypeCheckingWalker {
+ public:
+  struct VarInt {
+    uint64_t value;
+  };
+  struct SumIndex {
+    uint64_t index;
+  };
+  struct ListLength {
+    uint64_t length;
+  };
+  using Element = std::variant<VarInt, SumIndex, ListLength>;
+
+  TypeCheckingWalker(types::dy::Dynamic rootType,
+                     std::span<const uint64_t> buffer)
+      : stack({rootType}), buf(buffer) {
+  }
+
+  std::optional<Element> advance();
+
+ private:
+  std::stack<types::dy::Dynamic> stack;
+  std::span<const uint64_t> buf;
+
+ private:
+  uint64_t popFront() {
+    if (buf.empty()) {
+      throw std::runtime_error("unexpected end of data segment");
+    }
+    auto el = buf.front();
+    buf = buf.last(buf.size() - 1);
+    return el;
+  }
+};
+
+}  // namespace exporters
+}  // namespace ObjectIntrospection

oi/exporters/test/TypeCheckingWalkerTest.cpp

@@ -0,0 +1,185 @@
+#include <gtest/gtest.h>
+
+#include "oi/exporters/TypeCheckingWalker.h"
+
+using namespace ObjectIntrospection;
+using exporters::TypeCheckingWalker;
+
+TEST(TypeCheckingWalker, TestUnit) {
+  // ASSIGN
+  std::vector<uint64_t> data;
+  types::dy::Unit rootType;
+
+  TypeCheckingWalker walker(rootType, data);
+
+  // ACT
+  auto first = walker.advance();
+
+  // ASSERT
+  ASSERT_FALSE(first.has_value());
+}
+
+TEST(TypeCheckingWalker, TestVarInt) {
+  // ASSIGN
+  uint64_t val = 51566;
+  std::vector<uint64_t> data{val};
+  types::dy::VarInt rootType;
+
+  TypeCheckingWalker walker(rootType, data);
+
+  // ACT
+  auto first = walker.advance();
+
+  // ASSERT
+  ASSERT_TRUE(first.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::VarInt>(*first));
+  EXPECT_EQ(std::get<TypeCheckingWalker::VarInt>(*first).value, val);
+}
+
+TEST(TypeCheckingWalker, TestPair) {
+  // ASSIGN
+  uint64_t firstVal = 37894;
+  uint64_t secondVal = 6667;
+  std::vector<uint64_t> data{firstVal, secondVal};
+
+  types::dy::VarInt left;
+  types::dy::VarInt right;
+  types::dy::Pair rootType{left, right};
+
+  TypeCheckingWalker walker(rootType, data);
+
+  // ACT
+  auto first = walker.advance();
+  auto second = walker.advance();
+  auto third = walker.advance();
+
+  // ASSERT
+  ASSERT_TRUE(first.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::VarInt>(*first));
+  EXPECT_EQ(std::get<TypeCheckingWalker::VarInt>(*first).value, firstVal);
+
+  ASSERT_TRUE(second.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::VarInt>(*second));
+  EXPECT_EQ(std::get<TypeCheckingWalker::VarInt>(*second).value, secondVal);
+
+  ASSERT_FALSE(third.has_value());
+}
+
+TEST(TypeCheckingWalker, TestSumUnit) {
+  // ASSIGN
+  uint64_t sumIndex = 0;
+  std::vector<uint64_t> data{sumIndex};
+
+  types::dy::Unit unit;
+  types::dy::VarInt varint;
+  std::array<types::dy::Dynamic, 2> elements{unit, varint};
+  types::dy::Sum rootType{elements};
+
+  TypeCheckingWalker walker(rootType, data);
+
+  // ACT
+  auto first = walker.advance();
+  auto second = walker.advance();
+
+  // ASSERT
+  ASSERT_TRUE(first.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::SumIndex>(*first));
+  EXPECT_EQ(std::get<TypeCheckingWalker::SumIndex>(*first).index, sumIndex);
+
+  ASSERT_FALSE(second.has_value());
+}
+
+TEST(TypeCheckingWalker, TestSumVarInt) {
+  // ASSIGN
+  uint64_t sumIndex = 1;
+  uint64_t val = 63557;
+  std::vector<uint64_t> data{sumIndex, val};
+
+  types::dy::Unit unit;
+  types::dy::VarInt varint;
+  std::array<types::dy::Dynamic, 2> elements{unit, varint};
+  types::dy::Sum rootType{elements};
+
+  TypeCheckingWalker walker(rootType, data);
+
+  // ACT
+  auto first = walker.advance();
+  auto second = walker.advance();
+  auto third = walker.advance();
+
+  // ASSERT
+  ASSERT_TRUE(first.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::SumIndex>(*first));
+  EXPECT_EQ(std::get<TypeCheckingWalker::SumIndex>(*first).index, sumIndex);
+
+  ASSERT_TRUE(second.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::VarInt>(*second));
+  EXPECT_EQ(std::get<TypeCheckingWalker::VarInt>(*second).value, val);
+
+  ASSERT_FALSE(third.has_value());
+}
+
+TEST(TypeCheckingWalker, TestListEmpty) {
+  // ASSIGN
+  uint64_t listLength = 0;
+  std::vector<uint64_t> data{listLength};
+
+  types::dy::VarInt varint;
+  types::dy::List rootType{varint};
+
+  TypeCheckingWalker walker(rootType, data);
+
+  // ACT
+  auto first = walker.advance();
+  auto second = walker.advance();
+
+  // ASSERT
+  ASSERT_TRUE(first.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::ListLength>(*first));
+  EXPECT_EQ(std::get<TypeCheckingWalker::ListLength>(*first).length,
+            listLength);
+
+  ASSERT_FALSE(second.has_value());
+}
+
+TEST(TypeCheckingWalker, TestListSome) {
+  // ASSIGN
+  std::array<uint64_t, 3> listElements{59942, 44126, 64525};
+  std::vector<uint64_t> data{listElements.size(), listElements[0],
+                             listElements[1], listElements[2]};
+
+  types::dy::VarInt varint;
+  types::dy::List rootType{varint};
+
+  TypeCheckingWalker walker(rootType, data);
+
+  // ACT
+  auto first = walker.advance();
+  auto second = walker.advance();
+  auto third = walker.advance();
+  auto fourth = walker.advance();
+  auto fifth = walker.advance();
+
+  // ASSERT
+  ASSERT_TRUE(first.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::ListLength>(*first));
+  EXPECT_EQ(std::get<TypeCheckingWalker::ListLength>(*first).length,
+            listElements.size());
+
+  ASSERT_TRUE(second.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::VarInt>(*second));
+  EXPECT_EQ(std::get<TypeCheckingWalker::VarInt>(*second).value,
+            listElements[0]);
+
+  ASSERT_TRUE(third.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::VarInt>(*third));
+  EXPECT_EQ(std::get<TypeCheckingWalker::VarInt>(*third).value,
+            listElements[1]);
+
+  ASSERT_TRUE(fourth.has_value());
+  ASSERT_TRUE(std::holds_alternative<TypeCheckingWalker::VarInt>(*fourth));
+  EXPECT_EQ(std::get<TypeCheckingWalker::VarInt>(*fourth).value,
+            listElements[2]);
+
+  ASSERT_FALSE(fifth.has_value());
+}

test/CMakeLists.txt

@@ -101,6 +101,12 @@ cpp_unittest(
   DEPS oicore
 )
 
+cpp_unittest(
+  NAME type_checking_walker_test
+  SRCS ../oi/exporters/test/TypeCheckingWalkerTest.cpp
+  DEPS treebuilder
+)
+
 # Integration tests
 if (WITH_FLAKY_TESTS)
   add_test(