Compile Time OIL for OIL v2

Summary:
Update `OIGenerator` and out BUCK stuff for compile time OIL with OIL v2. Main things:
- Switch `OIGenerator` from the `getObjectSize` call to the new `introspect` call.
- Switch from looking at template parameters to looking at function parameters, as this was exposing a bug in our elfutils/drgn and this way it's the same as OID.
- Migrate `OIGenerator` to CodeGen v2 and update CodeGen v2 to accept a linkage name.
- Update the compile time example to be the same as the JIT example, using the new interface and the JSON exporter.
- Clean up the `ObjectIntrospection.h` header.

Differential Revision: D48687728

fbshipit-source-id: 2c3c041fd1b6499c5e02eb5e2082a977bfa529d7

include/ObjectIntrospection.h

@@ -1,268 +0,0 @@
-/*
- * 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 <atomic>
-#include <cstddef>
-#include <string>
-
-/*
- * Library interface for Object Introspection
- *
- * On the first call for each type these library functions perform significant
- * setup. In a single-threaded context, the calling thread blocks on the first
- * call. In a multi-threaded context, the first caller blocks, and other callers
- * see Response::OIL_INITIALISING until initialisation completes.
- *
- * The options passed to library functions MUST NOT change after the first call
- * for each type. By default, this will result in Response::OIL_CHANGED_OPTIONS.
- * This check can be disabled for decreased latency by passing checkOptions as
- * false.
- *
- * Generally the only required option is configFilePath. See sample.oid.toml for
- * an example configuration file.
- *
- * -- SINGLE-THREADED
- * ObjectIntrospection::options opts = { .configFilePath = "sample.oid.toml" };
- * size_t size;
- * int responseCode = ObjectIntrospection::getObjectSize(obj, size, opts);
- * if (responseCode != ObjectIntrospection::Response::OIL_SUCCESS) {
- *   // handle error
- * }
- *
- * -- MULTI-THREADED (NO SETUP)
- * ObjectIntrospection::options opts = { .configFilePath = "sample.oid.toml" };
- * size_t size;
- * int responseCode = ObjectIntrospection::getObjectSize(obj, size, opts);
- * if (responseCode > ObjectIntrospection::Response::OIL_INITIALISING) {
- *   // handle error
- * } else if (responseCode == ObjectIntrospection::Response::OIL_SUCCESS) {
- *   // do something
- * } // do nothing if still initialising
- *
- * -- MULTI-THREADED (WITH SETUP)
- * ObjectIntrospection::options opts = { .configFilePath = "sample.oid.toml" };
- * int responseCode = ObjectIntrospection::CodegenHandler<T>::init(opts);
- * if (responseCode != ObjectIntrospection::Response::OIL_SUCCESS) {
- *   // handle error
- * }
- * size_t size;
- * int responseCode = ObjectIntrospection::getObjectSize(obj, size);
- * if (responseCode == ObjectIntrospection::Response::OIL_UNINITIALISED) {
- *   // handle error - impossible if successfully inited
- * }
- *
- */
-namespace ObjectIntrospection {
-
-enum Response : int {
-  OIL_SUCCESS = 0,
-  OIL_INITIALISING = 1,
-  OIL_CHANGED_OPTIONS = 2,
-  OIL_BAD_CONFIG_FILE = 3,
-  OIL_SEGMENT_INIT_FAIL = 4,
-  OIL_COMPILATION_FAILURE = 5,
-  OIL_RELOCATION_FAILURE = 6,
-  OIL_UNINITIALISED = 7,
-};
-
-#ifndef OIL_AOT_COMPILATION
-
-struct options {
-  std::string configFilePath{};
-  std::string debugFilePath{};
-
-  int debugLevel = 0;
-  std::string sourceFileDumpPath{};
-
-  bool layout = false;
-  bool chaseRawPointers = false;
-  bool generateJitDebugInfo = false;
-
-  friend bool operator==(const options& lhs, const options& rhs);
-  friend bool operator!=(const options& lhs, const options& rhs);
-};
-
-constexpr std::string_view OI_SECTION_PREFIX = ".oi.";
-
-class OILibrary {
-  friend class OILibraryImpl;
-
- public:
-  OILibrary(void* TemplateFunc, options opt);
-  ~OILibrary();
-  int init();
-  int getObjectSize(void* objectAddr, size_t& size);
-
-  options opts;
-
- private:
-  class OILibraryImpl* pimpl_;
-
-  size_t (*fp)(const void*) = nullptr;
-};
-
-template <class T>
-class CodegenHandler {
- public:
-  static int init(const options& opts = {}, bool checkOptions = true) {
-    OILibrary* lib;
-    return getLibrary(lib, opts, checkOptions);
-  }
-
-  static void teardown() {
-    OILibrary* lib;
-    if (int responseCode = getLibrary(lib);
-        responseCode != Response::OIL_SUCCESS) {
-      return;
-    }
-
-    getLib()->store(nullptr);
-    getBoxedLib()->store(nullptr);
-    delete lib;
-  }
-
-  static int getObjectSize(const T& objectAddr, size_t& objectSize) {
-    OILibrary* lib;
-    if (int responseCode = getLibrary(lib);
-        responseCode != Response::OIL_SUCCESS) {
-      return responseCode;
-    }
-
-    return lib->getObjectSize((void*)&objectAddr, objectSize);
-  }
-
-  static int getObjectSize(const T& objectAddr,
-                           size_t& objectSize,
-                           const options& opts,
-                           bool checkOptions = true) {
-    OILibrary* lib;
-    if (int responseCode = getLibrary(lib, opts, checkOptions);
-        responseCode != Response::OIL_SUCCESS) {
-      return responseCode;
-    }
-
-    return lib->getObjectSize((void*)&objectAddr, objectSize);
-  }
-
- private:
-  static std::atomic<OILibrary*>* getLib() {
-    static std::atomic<OILibrary*> lib = nullptr;
-    return &lib;
-  }
-
-  static std::atomic<std::atomic<OILibrary*>*>* getBoxedLib() {
-    static std::atomic<std::atomic<OILibrary*>*> boxedLib = nullptr;
-    return &boxedLib;
-  }
-
-  static int getLibrary(OILibrary*& result) {
-    std::atomic<OILibrary*>* curBoxedLib = getBoxedLib()->load();
-    if (curBoxedLib == nullptr)
-      return Response::OIL_UNINITIALISED;
-
-    OILibrary* curLib = curBoxedLib->load();
-    if (curLib == nullptr)
-      return Response::OIL_UNINITIALISED;
-
-    result = curLib;
-    return Response::OIL_SUCCESS;
-  }
-
-  static int getLibrary(OILibrary*& result,
-                        const options& opts,
-                        bool checkOptions) {
-    std::atomic<OILibrary*>* curBoxedLib = getBoxedLib()->load();
-
-    if (curBoxedLib == nullptr) {
-      if (!getBoxedLib()->compare_exchange_strong(curBoxedLib, getLib())) {
-        return Response::OIL_INITIALISING;
-      }
-      curBoxedLib = getLib();
-
-      int (*sizeFp)(const T&, size_t&) = &getObjectSize;
-      void* typedFp = reinterpret_cast<void*>(sizeFp);
-      OILibrary* newLib = new OILibrary(typedFp, opts);
-
-      if (int initCode = newLib->init(); initCode != Response::OIL_SUCCESS) {
-        delete newLib;
-        getBoxedLib()->store(nullptr);  // allow next attempt to initialise
-        return initCode;
-      }
-
-      getLib()->store(newLib);
-    }
-
-    OILibrary* curLib = curBoxedLib->load();
-    if (curLib == nullptr) {
-      return Response::OIL_INITIALISING;
-    }
-
-    if (checkOptions && opts != curLib->opts) {
-      return Response::OIL_CHANGED_OPTIONS;
-    }
-
-    result = curLib;
-    return Response::OIL_SUCCESS;
-  }
-};
-
-/*
- * Call this from anywhere in your program. It blocks on the first call for
- * each type when seen for the first time. Usage patterns are given at the
- * top of this file. This method should not be called when utilising
- * Ahead-Of-Time (AOT) compilation.
- */
-template <class T>
-int getObjectSize(const T& objectAddr,
-                  size_t& objectSize,
-                  const options& opts,
-                  bool checkOptions = true) {
-  return CodegenHandler<T>::getObjectSize(objectAddr, objectSize, opts,
-                                          checkOptions);
-}
-
-#else
-
-template <class T>
-int __attribute__((weak))
-getObjectSizeImpl(const T& objectAddr, size_t& objectSize);
-
-#endif
-
-/*
- * You may only call this after a call to the previous signature, or a
- * call to CodegenHandler<T>::init(...) for the used type.
- *
- * As we can choose to compile the OIL blob Ahead-Of-Time (AOT) rather
- * than Just-In-Time (JIT), this default is provided as a weak symbol.
- * When in AOT mode this will no-op, removing the burden of JIT on a
- * production system.
- */
-template <class T>
-int __attribute__((noinline))
-getObjectSize(const T& objectAddr, size_t& objectSize) {
-#ifdef OIL_AOT_COMPILATION
-  if (!getObjectSizeImpl<T>) {
-    return Response::OIL_UNINITIALISED;
-  }
-  return getObjectSizeImpl(objectAddr, objectSize);
-#else
-  return CodegenHandler<T>::getObjectSize(objectAddr, objectSize);
-#endif
-}
-
-}  // namespace ObjectIntrospection

include/oi/oi.h

@@ -20,6 +20,7 @@
 #include <oi/types/dy.h>
 
 #include <cstdint>
+#include <stdexcept>
 #include <vector>
 
 namespace oi {
@@ -33,15 +34,26 @@ enum class Feature {
 template <typename T, Feature... Fs>
 IntrospectionResult __attribute__((noinline)) introspect(const T& objectAddr);
 
+#ifdef OIL_AOT_COMPILATION
+
+template <class T, Feature... Fs>
+IntrospectionResult __attribute__((weak)) introspectImpl(const T& objectAddr);
+
+template <typename T, Feature... Fs>
+IntrospectionResult introspect(const T& objectAddr) {
+  if (!introspectImpl<T, Fs...>)
+    throw std::logic_error(
+        "OIL is expecting AoT compilation but it doesn't appear to have run.");
+
+  return introspectImpl<T, Fs...>(objectAddr);
+}
+
+#endif
+
 }  // namespace oi
 
 #ifndef OIL_AOT_COMPILATION
 #include "oi/oi-jit.h"
-#else
-
-template <typename T, Feature... Fs>
-IntrospectionResult __attribute__((noinline)) introspect(const T& objectAddr);
-{ static_assert(false, "OIL v2 does not yet support AoT compilation."); }
+#endif
 
 #endif
-#endif

oi/CodeGen.cpp

@@ -162,8 +162,10 @@ void addIncludes(const TypeGraph& typeGraph,
   }
   if (features[Feature::TreeBuilderV2])
     includes.emplace("oi/exporters/inst.h");
-  if (features[Feature::Library])
+  if (features[Feature::Library]) {
     includes.emplace("vector");
+    includes.emplace("oi/IntrospectionResult.h");
+  }
   if (features[Feature::JitTiming]) {
     includes.emplace("chrono");
   }
@@ -1026,6 +1028,13 @@ void CodeGen::addTypeHandlers(const TypeGraph& typeGraph, std::string& code) {
   }
 }
 
+bool CodeGen::codegenFromDrgn(struct drgn_type* drgnType,
+                              std::string linkageName,
+                              std::string& code) {
+  linkageName_ = std::move(linkageName);
+  return codegenFromDrgn(drgnType, code);
+}
+
 bool CodeGen::codegenFromDrgn(struct drgn_type* drgnType, std::string& code) {
   try {
     containerInfos_.reserve(config_.containerConfigPaths.size());
@@ -1206,6 +1215,9 @@ void CodeGen::generate(
     FuncGen::DefineOutputType(code, typeName);
   }
 
+  if (!linkageName_.empty())
+    FuncGen::DefineTopLevelIntrospectNamed(code, typeName, linkageName_);
+
   if (VLOG_IS_ON(3)) {
     VLOG(3) << "Generated trace code:\n";
     // VLOG truncates output, so use std::cerr

oi/CodeGen.h

@@ -48,6 +48,9 @@ class CodeGen {
    * single drgn_type.
    */
   bool codegenFromDrgn(struct drgn_type* drgnType, std::string& code);
+  bool codegenFromDrgn(struct drgn_type* drgnType,
+                       std::string linkageName,
+                       std::string& code);
 
   void registerContainer(const std::filesystem::path& path);
   void addDrgnRoot(struct drgn_type* drgnType,
@@ -66,6 +69,7 @@ class CodeGen {
   std::unordered_set<const ContainerInfo*> definedContainers_;
   std::unordered_map<const type_graph::Class*, const type_graph::Member*>
       thriftIssetMembers_;
+  std::string linkageName_;
 
   void genDefsThrift(const type_graph::TypeGraph& typeGraph, std::string& code);
   void addGetSizeFuncDefs(const type_graph::TypeGraph& typeGraph,

oi/FuncGen.cpp

@@ -274,6 +274,28 @@ void __attribute__((used, retain)) introspect_%2$016x(
       (boost::format(func) % type % std::hash<std::string>{}(type)).str());
 }
 
+void FuncGen::DefineTopLevelIntrospectNamed(std::string& code,
+                                            const std::string& type,
+                                            const std::string& linkageName) {
+  std::string typeHash =
+      (boost::format("%1$016x") % std::hash<std::string>{}(type)).str();
+
+  code += "/* RawType: ";
+  code += type;
+  code += " */\n";
+  code += "extern \"C\" IntrospectionResult ";
+  code += linkageName;
+  code += "(const OIInternal::__ROOT_TYPE__& t) {\n";
+  code += "  std::vector<uint8_t> v{};\n";
+  code += "  introspect_";
+  code += typeHash;
+  code += "(t, v);\n";
+  code += "  return IntrospectionResult{std::move(v), treeBuilderInstructions";
+  code += typeHash;
+  code += "};\n";
+  code += "}\n";
+}
+
 void FuncGen::DefineTopLevelGetSizeRef(std::string& testCode,
                                        const std::string& rawType,
                                        FeatureSet features) {

oi/FuncGen.h

@@ -61,6 +61,9 @@ class FuncGen {
                                           const std::string& linkageName);
   static void DefineTopLevelIntrospect(std::string& code,
                                        const std::string& type);
+  static void DefineTopLevelIntrospectNamed(std::string& code,
+                                            const std::string& type,
+                                            const std::string& linkageName);
 
   static void DefineTopLevelGetSizeRef(std::string& testCode,
                                        const std::string& rawType,

oi/Headers.h

@@ -18,6 +18,8 @@
 namespace oi::detail::headers {
 
 // These externs are provided by our build system. See resources/CMakeLists.txt
+extern const std::string_view oi_IntrospectionResult_h;
+extern const std::string_view oi_IntrospectionResult_inl_h;
 extern const std::string_view oi_OITraceCode_cpp;
 extern const std::string_view oi_exporters_ParsedData_h;
 extern const std::string_view oi_exporters_inst_h;

oi/OICompiler.cpp

@@ -518,7 +518,7 @@ bool OICompiler::compile(const std::string& code,
   }
 
   static const auto syntheticHeaders = std::array<
-      std::pair<Feature, std::pair<std::string_view, std::string>>, 5>{{
+      std::pair<Feature, std::pair<std::string_view, std::string>>, 7>{{
       {Feature::TypedDataSegment, {headers::oi_types_st_h, "oi/types/st.h"}},
       {Feature::TreeBuilderTypeChecking,
        {headers::oi_types_dy_h, "oi/types/dy.h"}},
@@ -528,6 +528,10 @@ bool OICompiler::compile(const std::string& code,
        {headers::oi_exporters_ParsedData_h, "oi/exporters/ParsedData.h"}},
       {Feature::TreeBuilderV2,
        {headers::oi_result_Element_h, "oi/result/Element.h"}},
+      {Feature::Library,
+       {headers::oi_IntrospectionResult_h, "oi/IntrospectionResult.h"}},
+      {Feature::Library,
+       {headers::oi_IntrospectionResult_inl_h, "oi/IntrospectionResult-inl.h"}},
   }};
   for (const auto& [k, v] : syntheticHeaders) {
     if (!config.features[k])

oi/OIGenerator.cpp

@@ -24,6 +24,7 @@
 #include <unordered_map>
 #include <variant>
 
+#include "oi/CodeGen.h"
 #include "oi/DrgnUtils.h"
 #include "oi/Headers.h"
 #include "oi/OIUtils.h"
@@ -33,9 +34,9 @@ namespace oi::detail {
 std::unordered_map<std::string, std::string>
 OIGenerator::oilStrongToWeakSymbolsMap(drgnplusplus::program& prog) {
   static constexpr std::string_view strongSymbolPrefix =
-      "int ObjectIntrospection::getObjectSize<";
+      "oi::IntrospectionResult oi::introspect<";
   static constexpr std::string_view weakSymbolPrefix =
-      "int ObjectIntrospection::getObjectSizeImpl<";
+      "oi::IntrospectionResult oi::introspectImpl<";
 
   std::unordered_map<std::string, std::pair<std::string, std::string>>
       templateArgsToSymbolsMap;
@@ -50,14 +51,14 @@ OIGenerator::oilStrongToWeakSymbolsMap(drgnplusplus::program& prog) {
           strongSymbolPrefix.length())];
       if (!matchedSyms.first.empty()) {
         LOG(WARNING) << "non-unique symbols found: `" << matchedSyms.first
-                     << "` and `" << symName << "`";
+                     << "` and `" << symName << '`';
       }
       matchedSyms.first = symName;
     } else if (demangled.starts_with(weakSymbolPrefix)) {
       auto& matchedSyms =
           templateArgsToSymbolsMap[demangled.substr(weakSymbolPrefix.length())];
       if (!matchedSyms.second.empty()) {
-        LOG(WARNING) << "non-unique symbols found: `" << matchedSyms.first
+        LOG(WARNING) << "non-unique symbols found: `" << matchedSyms.second
                      << "` and `" << symName << "`";
       }
       matchedSyms.second = symName;
@@ -75,15 +76,13 @@ OIGenerator::oilStrongToWeakSymbolsMap(drgnplusplus::program& prog) {
   return strongToWeakSymbols;
 }
 
-std::vector<std::tuple<drgn_qualified_type, std::string>>
+std::unordered_map<std::string, drgn_qualified_type>
 OIGenerator::findOilTypesAndNames(drgnplusplus::program& prog) {
   auto strongToWeakSymbols = oilStrongToWeakSymbolsMap(prog);
 
-  std::vector<std::tuple<drgn_qualified_type, std::string>> out;
+  std::unordered_map<std::string, drgn_qualified_type> out;
 
-  // TODO: Clean up this loop when switching to
-  // drgn_program_find_function_by_address.
-  for (auto& func : drgnplusplus::func_iterator(prog)) {
+  for (drgn_qualified_type& func : drgnplusplus::func_iterator(prog)) {
     std::string strongLinkageName;
     {
       const char* linkageNameCstr;
@@ -103,22 +102,24 @@ OIGenerator::findOilTypesAndNames(drgnplusplus::program& prog) {
       continue;  // not an oil strong symbol
     }
 
-    auto templateParameters = drgn_type_template_parameters(func.type);
-    drgn_type_template_parameter param = templateParameters[0];
+    // IntrospectionResult (*)(const T&)
+    CHECK(drgn_type_has_parameters(func.type)) << "functions have parameters";
+    CHECK(drgn_type_num_parameters(func.type) == 1)
+        << "introspection func has one parameter";
 
-    drgn_qualified_type paramType;
-    if (auto err = drgnplusplus::error(
-            drgn_template_parameter_type(&param, &paramType))) {
-      LOG(ERROR) << "error getting drgn template parameter type: " << err;
+    auto* params = drgn_type_parameters(func.type);
+    drgn_qualified_type tType;
+    if (auto err =
+            drgnplusplus::error(drgn_parameter_type(&params[0], &tType))) {
       throw err;
     }
 
-    if (drgn_type_has_name(paramType.type)) {
-      LOG(INFO) << "found OIL type: " << drgn_type_name(paramType.type);
+    if (drgn_type_has_name(tType.type)) {
+      LOG(INFO) << "found OIL type: " << drgn_type_name(tType.type);
     } else {
       LOG(INFO) << "found OIL type: (no name)";
     }
-    out.push_back({paramType, std::move(weakLinkageName)});
+    out.emplace(std::move(weakLinkageName), tType);
   }
 
   return out;
@@ -129,19 +130,11 @@ fs::path OIGenerator::generateForType(const OICodeGen::Config& generatorConfig,
                                       const drgn_qualified_type& type,
                                       const std::string& linkageName,
                                       SymbolService& symbols) {
-  auto codegen = OICodeGen::buildFromConfig(generatorConfig, symbols);
-  if (!codegen) {
-    LOG(ERROR) << "failed to initialise codegen";
-    return {};
-  }
+  CodeGen codegen{generatorConfig, symbols};
 
-  std::string code(headers::oi_OITraceCode_cpp);
-
-  codegen->setRootType(type);
-  codegen->setLinkageName(linkageName);
-
-  if (!codegen->generate(code)) {
-    LOG(ERROR) << "failed to generate code";
+  std::string code;
+  if (!codegen.codegenFromDrgn(type.type, linkageName, code)) {
+    LOG(ERROR) << "codegen failed!";
     return {};
   }
 
@@ -180,10 +173,14 @@ int OIGenerator::generate(fs::path& primaryObject, SymbolService& symbols) {
     }
   }
 
-  std::vector<std::tuple<drgn_qualified_type, std::string>> oilTypes =
-      findOilTypesAndNames(prog);
+  auto oilTypes = findOilTypesAndNames(prog);
 
   std::map<Feature, bool> featuresMap = {
+      {Feature::TypeGraph, true},
+      {Feature::TypedDataSegment, true},
+      {Feature::TreeBuilderTypeChecking, true},
+      {Feature::TreeBuilderV2, true},
+      {Feature::Library, true},
       {Feature::PackStructs, true},
       {Feature::PruneTypeGraph, true},
   };
@@ -198,10 +195,10 @@ int OIGenerator::generate(fs::path& primaryObject, SymbolService& symbols) {
     return -1;
   }
   generatorConfig.features = *features;
-  generatorConfig.useDataSegment = false;
+  compilerConfig.features = *features;
 
   size_t failures = 0;
-  for (const auto& [type, linkageName] : oilTypes) {
+  for (const auto& [linkageName, type] : oilTypes) {
     if (auto obj = generateForType(generatorConfig, compilerConfig, type,
                                    linkageName, symbols);
         !obj.empty()) {

oi/OIGenerator.h

@@ -50,8 +50,8 @@ class OIGenerator {
   std::unordered_map<std::string, std::string> oilStrongToWeakSymbolsMap(
       drgnplusplus::program& prog);
 
-  std::vector<std::tuple<drgn_qualified_type, std::string>>
-  findOilTypesAndNames(drgnplusplus::program& prog);
+  std::unordered_map<std::string, drgn_qualified_type> findOilTypesAndNames(
+      drgnplusplus::program& prog);
 
   std::filesystem::path generateForType(
       const OICodeGen::Config& generatorConfig,

resources/CMakeLists.txt

@@ -7,6 +7,8 @@ function(embed_headers output)
   file(APPEND ${output} "namespace oi::detail::headers {\n")
 
   set(HEADERS
+    ../include/oi/IntrospectionResult-inl.h
+    ../include/oi/IntrospectionResult.h
     ../include/oi/exporters/ParsedData.h
     ../include/oi/exporters/inst.h
     ../include/oi/result/Element.h