type checking: add description of data segment type

CMakeLists.txt

@@ -272,6 +272,7 @@ add_library(oicore
 add_dependencies(oicore libdrgn)
 target_include_directories(oicore SYSTEM PRIVATE ${LLVM_INCLUDE_DIRS} ${CLANG_INCLUDE_DIRS})
 target_compile_definitions(oicore PRIVATE ${LLVM_DEFINITIONS})
+target_include_directories(oicore PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
 
 llvm_map_components_to_libnames(llvm_libs core native mcjit x86disassembler)
 target_link_libraries(oicore

include/oi/types/dy.h

@@ -0,0 +1,98 @@
+/*
+ * 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.
+ */
+
+#ifndef OI_TYPES_DY_H
+#define OI_TYPES_DY_H 1
+
+/*
+ * Dynamic Types
+ *
+ * These types are a dynamic (runtime) description of the types in
+ * <oi/include/types/st.h>. We use static types to ensure that what is written
+ * in the data segment is a known type. Dynamic types extend this to runtime,
+ * allowing TreeBuilder to check that what it is reading out of the data segment
+ * is what went in. Static types are written with heavy usage of templating so
+ * they can be inlined and compiled to nothing in the JIT output. When
+ * translating these types to OID/OITB they cannot be compiled in, so we
+ * represent the template arguments with member fields instead.
+ *
+ * Each type in this namespace corresponds 1-1 with a type in
+ * ObjectIntrospection::types::st, except Dynamic which references them all. See
+ * the types in st.h for the description of what each type contains.
+ *
+ * All types in this file include a constexpr constructor. This allows a single
+ * extern const variable in the JIT code to include pointer references to other
+ * dynamic elements. To map from a Static Type to a Dynamic Type, we store each
+ * template parameter as a field or std::span in a field.
+ */
+
+#include <functional>
+#include <span>
+#include <variant>
+
+namespace ObjectIntrospection::types::dy {
+
+class Unit;
+class VarInt;
+class Pair;
+class Sum;
+class List;
+
+/*
+ * Dynamic
+ *
+ * The cornerstone of the dynamic types is the Dynamic type, a variant which
+ * holds a pointer of each different type. This is essentially a tagged pointer.
+ */
+using Dynamic = std::variant<std::reference_wrapper<const Unit>,
+                             std::reference_wrapper<const VarInt>,
+                             std::reference_wrapper<const Pair>,
+                             std::reference_wrapper<const Sum>,
+                             std::reference_wrapper<const List> >;
+
+class Unit {};
+class VarInt {};
+
+class Pair {
+ public:
+  constexpr Pair(Dynamic first_, Dynamic second_)
+      : first(first_), second(second_) {
+  }
+
+  Dynamic first;
+  Dynamic second;
+};
+
+class Sum {
+ public:
+  template <size_t N>
+  constexpr Sum(const std::array<Dynamic, N>& variants_) : variants(variants_) {
+  }
+
+  std::span<const Dynamic> variants;
+};
+
+class List {
+ public:
+  constexpr List(Dynamic element_) : element(element_) {
+  }
+
+  Dynamic element;
+};
+
+}  // namespace ObjectIntrospection
+
+#endif

include/oi/types/st.h

@@ -49,9 +49,18 @@
  *              In this case, `ret` is of type `ComplexType`. After the two
  *              writes, the inner function returns `Unit`. Delegate then
  *              internally converts this unit to a `VarInt`.
+ *
+ * DEFINE_DESCRIBE controls the additional feature of dynamic descriptions of
+ * types. If defined when this header is included, static types provide a
+ * dynamic description of their type as the constexpr field `describe`. Compound
+ * types compose appropriately.
  */
 namespace ObjectIntrospection::types::st {
 
+#ifdef DEFINE_DESCRIBE
+#include "oi/types/dy.h"
+#endif
+
 /*
  * Unit
  *
@@ -75,6 +84,10 @@ class Unit {
     return cb(*this);
   }
 
+#ifdef DEFINE_DESCRIBE
+  static constexpr types::dy::Unit describe{};
+#endif
+
  private:
   /*
    * Allows you to cast the Unit type to another Static Type. Think very
@@ -118,6 +131,10 @@ class VarInt {
     return Unit<DataBuffer>(_buf);
   }
 
+#ifdef DEFINE_DESCRIBE
+  static constexpr types::dy::VarInt describe{};
+#endif
+
  private:
   DataBuffer _buf;
 };
@@ -148,6 +165,10 @@ class Pair {
     return second.template cast<T2>();
   }
 
+#ifdef DEFINE_DESCRIBE
+  static constexpr types::dy::Pair describe{T1::describe, T2::describe};
+#endif
+
  private:
   DataBuffer _buf;
 };
@@ -196,6 +217,15 @@ class Sum {
     return cb(tail);
   }
 
+#ifdef DEFINE_DESCRIBE
+ private:
+  static constexpr std::array<types::dy::Dynamic, sizeof...(Types)> members{
+      Types::describe...};
+
+ public:
+  static constexpr types::dy::Sum describe{members};
+#endif
+
  private:
   DataBuffer _buf;
 };
@@ -237,7 +267,18 @@ class ListContents {
  * elements promised.
  */
 template <typename DataBuffer, typename T>
-using List = Pair<DataBuffer, VarInt<DataBuffer>, ListContents<DataBuffer, T>>;
+class List
+    : public Pair<DataBuffer, VarInt<DataBuffer>, ListContents<DataBuffer, T>> {
+ public:
+  List(DataBuffer db)
+      : Pair<DataBuffer, VarInt<DataBuffer>, ListContents<DataBuffer, T>>(db) {
+  }
+
+#ifdef DEFINE_DESCRIBE
+ public:
+  static constexpr types::dy::List describe{T::describe};
+#endif
+};
 
 }  // namespace ObjectIntrospection::types::st
 

oi/CodeGen.cpp

@@ -106,6 +106,11 @@ void addIncludes(const TypeGraph& typeGraph,
     includes.emplace("functional");
     includes.emplace("oi/types/st.h");
   }
+  if (features[Feature::TreeBuilderTypeChecking]) {
+    includes.emplace("oi/types/dy.h");
+
+    code += "#define DEFINE_DESCRIBE 1\n"; // added before all includes
+  }
   if (features[Feature::JitTiming]) {
     includes.emplace("chrono");
   }
@@ -867,12 +872,15 @@ void CodeGen::generate(
   code += "\nusing __ROOT_TYPE__ = " + rootType.name() + ";\n";
   code += "} // namespace\n} // namespace OIInternal\n";
 
+  const auto typeName = SymbolService::getTypeName(drgnType);
   if (config_.features[Feature::TypedDataSegment]) {
-    FuncGen::DefineTopLevelGetSizeRefTyped(
-        code, SymbolService::getTypeName(drgnType), config_.features);
+    FuncGen::DefineTopLevelGetSizeRefTyped(code, typeName, config_.features);
   } else {
-    FuncGen::DefineTopLevelGetSizeRef(
-        code, SymbolService::getTypeName(drgnType), config_.features);
+    FuncGen::DefineTopLevelGetSizeRef(code, typeName, config_.features);
+  }
+
+  if (config_.features[Feature::TreeBuilderTypeChecking]) {
+    FuncGen::DefineOutputType(code, typeName);
   }
 
   if (VLOG_IS_ON(3)) {

oi/Features.cpp

@@ -36,6 +36,9 @@ std::string_view featureHelp(Feature f) {
       return "Use Type Graph for code generation (CodeGen v2).";
     case Feature::TypedDataSegment:
       return "Use Typed Data Segment in generated code.";
+    case Feature::TreeBuilderTypeChecking:
+      return "Use Typed Data Segment to perform runtime Type Checking in "
+             "TreeBuilder.";
     case Feature::TreeBuilderV2:
       return "Use Tree Builder v2 for reading the data segment";
     case Feature::GenJitDebug:

oi/Features.h

@@ -21,17 +21,18 @@
 
 #include "oi/EnumBitset.h"
 
-#define OI_FEATURE_LIST                         \
-  X(ChaseRawPointers, "chase-raw-pointers")     \
-  X(PackStructs, "pack-structs")                \
-  X(GenPaddingStats, "gen-padding-stats")       \
-  X(CaptureThriftIsset, "capture-thrift-isset") \
-  X(TypeGraph, "type-graph")                    \
-  X(TypedDataSegment, "typed-data-segment")     \
-  X(TreeBuilderV2, "tree-builder-v2")           \
-  X(GenJitDebug, "gen-jit-debug")               \
-  X(JitLogging, "jit-logging")                  \
-  X(JitTiming, "jit-timing")                    \
+#define OI_FEATURE_LIST                                    \
+  X(ChaseRawPointers, "chase-raw-pointers")                \
+  X(PackStructs, "pack-structs")                           \
+  X(GenPaddingStats, "gen-padding-stats")                  \
+  X(CaptureThriftIsset, "capture-thrift-isset")            \
+  X(TypeGraph, "type-graph")                               \
+  X(TypedDataSegment, "typed-data-segment")                \
+  X(TreeBuilderTypeChecking, "tree-builder-type-checking") \
+  X(TreeBuilderV2, "tree-builder-v2")                      \
+  X(GenJitDebug, "gen-jit-debug")                          \
+  X(JitLogging, "jit-logging")                             \
+  X(JitTiming, "jit-timing")                               \
   X(PolymorphicInheritance, "polymorphic-inheritance")
 
 namespace ObjectIntrospection {

oi/FuncGen.cpp

@@ -364,6 +364,26 @@ void FuncGen::DefineTopLevelGetSizeRefTyped(std::string& testCode,
   testCode.append(fmt.str());
 }
 
+/*
+ * DefineOutputType
+ *
+ * Present the dynamic type of an object for OID/OIL/OITB to link against.
+ */
+void FuncGen::DefineOutputType(std::string& code, const std::string& rawType) {
+  std::string func = R"(
+    #pragma GCC diagnostic push
+    #pragma GCC diagnostic ignored "-Wunknown-attributes"
+    /* RawType: %1% */
+    extern const types::dy::Dynamic __attribute__((used, retain)) outputType%2$016x =
+      OIInternal::TypeHandler<DataBuffer::DataSegment, OIInternal::__ROOT_TYPE__>::type::describe;
+    #pragma GCC diagnostic pop
+  )";
+
+  boost::format fmt =
+      boost::format(func) % rawType % std::hash<std::string>{}(rawType);
+  code.append(fmt.str());
+}
+
 void FuncGen::DefineTopLevelGetSizeRefRet(std::string& testCode,
                                           const std::string& rawType) {
   std::string func = R"(

oi/FuncGen.h

@@ -63,6 +63,8 @@ class FuncGen {
       std::string& testCode,
       const std::string& rawType,
       ObjectIntrospection::FeatureSet features);
+  static void DefineOutputType(std::string& testCode,
+                               const std::string& rawType);
 
   static void DefineTopLevelGetSizeRefRet(std::string& testCode,
                                           const std::string& type);

oi/Headers.h

@@ -21,6 +21,7 @@ namespace headers {
 // These externs are provided by our build system. See resources/CMakeLists.txt
 extern const std::string_view oi_OITraceCode_cpp;
 extern const std::string_view oi_types_st_h;
+extern const std::string_view oi_types_dy_h;
 
 }  // namespace headers
 }  // namespace ObjectIntrospection

oi/OICompiler.cpp

@@ -523,6 +523,14 @@ bool OICompiler::compile(const std::string& code,
         "/synthetic/headers", clang::frontend::IncludeDirGroup::IndexHeaderMap,
         false, false);
   }
+  if (config.features[Feature::TreeBuilderTypeChecking]) {
+    compInv->getPreprocessorOpts().addRemappedFile(
+        "/synthetic/headers/oi/types/dy.h",
+        MemoryBuffer::getMemBuffer(headers::oi_types_dy_h).release());
+    headerSearchOptions.AddPath(
+        "/synthetic/headers", clang::frontend::IncludeDirGroup::IndexHeaderMap,
+        false, false);
+  }
 
   compInv->getFrontendOpts().OutputFile = objectPath;
   compInv->getTargetOpts().Triple =

oi/OIUtils.cpp

@@ -180,6 +180,18 @@ std::optional<ObjectIntrospection::FeatureSet> processConfigFile(
     }
   }
 
+  if (featuresSet[Feature::TreeBuilderTypeChecking] &&
+      !featuresSet[Feature::TypedDataSegment]) {
+    if (auto search = featureMap.find(Feature::TypedDataSegment);
+        search != featureMap.end() && !search->second) {
+      LOG(ERROR) << "TreeBuilderTypeChecking feature requires TypedDataSegment "
+                    "feature to be enabled but it was explicitly disabled!";
+      return {};
+    }
+    featuresSet[Feature::TypedDataSegment] = true;
+    LOG(WARNING) << "TreeBuilderTypeChecking feature requires TypedDataSegment "
+                    "feature to be enabled, enabling now.";
+  }
   if (featuresSet[Feature::TypedDataSegment] &&
       !featuresSet[Feature::TypeGraph]) {
     if (auto search = featureMap.find(Feature::TypeGraph);

oi/types/test/StaticTest.cpp

@@ -0,0 +1,104 @@
+#include <gtest/gtest.h>
+
+#define DEFINE_DESCRIBE 1
+#include "oi/types/dy.h"
+#include "oi/types/st.h"
+
+using namespace ObjectIntrospection;
+
+class DummyDataBuffer {};
+
+TEST(StaticTypes, TestUnitToDynamic) {
+  // ASSIGN
+  using ty = types::st::Unit<DummyDataBuffer>;
+
+  // ACT
+  types::dy::Dynamic dynamicType = ty::describe;
+
+  // ASSERT
+  ASSERT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::Unit>>(
+          dynamicType));
+}
+
+TEST(StaticTypes, TestVarIntToDynamic) {
+  // ASSIGN
+  using ty = types::st::VarInt<DummyDataBuffer>;
+
+  // ACT
+  types::dy::Dynamic dynamicType = ty::describe;
+
+  // ASSERT
+  ASSERT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::VarInt>>(
+          dynamicType));
+}
+
+TEST(StaticTypes, TestPairToDynamic) {
+  // ASSIGN
+  using left = types::st::VarInt<DummyDataBuffer>;
+  using right = types::st::VarInt<DummyDataBuffer>;
+  using ty = types::st::Pair<DummyDataBuffer, left, right>;
+
+  // ACT
+  types::dy::Dynamic dynamicType = ty::describe;
+
+  // ASSERT
+  ASSERT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::Pair>>(
+          dynamicType));
+  const types::dy::Pair& pairType =
+      std::get<std::reference_wrapper<const types::dy::Pair>>(dynamicType);
+
+  EXPECT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::VarInt>>(
+          pairType.first));
+  EXPECT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::VarInt>>(
+          pairType.second));
+}
+
+TEST(StaticTypes, TestSumToDynamic) {
+  // ASSIGN
+  using left = types::st::Unit<DummyDataBuffer>;
+  using right = types::st::VarInt<DummyDataBuffer>;
+  using ty = types::st::Sum<DummyDataBuffer, left, right>;
+
+  // ACT
+  types::dy::Dynamic dynamicType = ty::describe;
+
+  // ASSERT
+  ASSERT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::Sum>>(
+          dynamicType));
+  const types::dy::Sum& sumType =
+      std::get<std::reference_wrapper<const types::dy::Sum>>(dynamicType);
+
+  ASSERT_EQ(sumType.variants.size(), 2);
+  EXPECT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::Unit>>(
+          sumType.variants[0]));
+  EXPECT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::VarInt>>(
+          sumType.variants[1]));
+}
+
+TEST(StaticTypes, TestListToDynamic) {
+  // ASSIGN
+  using el = types::st::VarInt<DummyDataBuffer>;
+  using ty = types::st::List<DummyDataBuffer, el>;
+
+  // ACT
+  types::dy::Dynamic dynamicType = ty::describe;
+
+  // ASSERT
+  ASSERT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::List>>(
+          dynamicType));
+  const types::dy::List& listType =
+      std::get<std::reference_wrapper<const types::dy::List>>(dynamicType);
+
+  EXPECT_TRUE(
+      std::holds_alternative<std::reference_wrapper<const types::dy::VarInt>>(
+          listType.element));
+}

resources/CMakeLists.txt

@@ -8,6 +8,7 @@ function(embed_headers output)
   file(APPEND ${output} "namespace headers {\n\n")
 
   set(HEADERS
+    ../include/oi/types/dy.h
     ../include/oi/types/st.h
     ../oi/OITraceCode.cpp
   )

test/CMakeLists.txt

@@ -95,6 +95,12 @@ cpp_unittest(
   DEPS oicore
 )
 
+cpp_unittest(
+  NAME types_static_test
+  SRCS ../oi/types/test/StaticTest.cpp
+  DEPS oicore
+)
+
 # Integration tests
 if (WITH_FLAKY_TESTS)
   add_test(
@@ -110,4 +116,3 @@ endif()
 
 include_directories("${PROJECT_SOURCE_DIR}/extern/drgn/build")
 add_subdirectory("integration")
-

test/integration/arrays.toml

@@ -44,7 +44,7 @@ definitions = '''
         "dynamicSize":0
       }]}]'''
   [cases.multidim_legacy] # Test for legacy behaviour. Remove with OICodeGen
-    cli_options = ["-Ftype-graph", "-Ftyped-data-segment"]
+    cli_options = ["-Ftype-graph", "-Ftyped-data-segment", "-Ftree-builder-type-checking"]
     param_types = ["const MultiDim&"]
     setup = "return {};"
     expect_json = '''[{