From 37b89d789dcb6b09ad0d458f1d9bdf15791df579 Mon Sep 17 00:00:00 2001
From: Jake Hillion <jakehillion@meta.com>
Date: Tue, 19 Dec 2023 15:00:35 +0000
Subject: [PATCH] codegen: remove reliance on drgn type for top level name

Currently we rely on `SymbolService::getTypeName` for getting the hash that's
included in the generated function's name. The value of this must stay the same
to match with the value expected by OIDebugger - changing it causes failure to
relocate when attaching with OID and JIT OIL.

Calculate this name in the `codegenFromDrgn` method and pass it through where
appropriate rather than passing the `drgn_type` itself through.

We don't need to name the type like that when using AoT OIL. Let's
hash the linkage name instead as that is more unique.

Test Plan:
- CI
---
 oi/CodeGen.cpp | 46 ++++++++++++++++++++++++++++++++--------------
 oi/CodeGen.h   | 17 +++++++++++++----
 2 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/oi/CodeGen.cpp b/oi/CodeGen.cpp
index d866df3..74c6f39 100644
--- a/oi/CodeGen.cpp
+++ b/oi/CodeGen.cpp
@@ -43,6 +43,9 @@
 #include "type_graph/TypeIdentifier.h"
 #include "type_graph/Types.h"
 
+template <typename T>
+inline constexpr bool always_false_v = false;
+
 namespace oi::detail {
 
 using type_graph::AddChildren;
@@ -1101,11 +1104,17 @@ 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);
+  return codegenFromDrgn(drgnType, code, ExactName{std::move(linkageName)});
 }
 
 bool CodeGen::codegenFromDrgn(struct drgn_type* drgnType, std::string& code) {
+  return codegenFromDrgn(
+      drgnType, code, HashedComponent{SymbolService::getTypeName(drgnType)});
+}
+
+bool CodeGen::codegenFromDrgn(struct drgn_type* drgnType,
+                              std::string& code,
+                              RootFunctionName name) {
   try {
     containerInfos_.reserve(config_.containerConfigPaths.size());
     for (const auto& path : config_.containerConfigPaths) {
@@ -1124,7 +1133,7 @@ bool CodeGen::codegenFromDrgn(struct drgn_type* drgnType, std::string& code) {
   }
 
   transform(typeGraph_);
-  generate(typeGraph_, code, drgnType);
+  generate(typeGraph_, code, std::move(name));
   return true;
 }
 
@@ -1213,11 +1222,9 @@ void CodeGen::transform(TypeGraph& typeGraph) {
   };
 }
 
-void CodeGen::generate(
-    TypeGraph& typeGraph,
-    std::string& code,
-    struct drgn_type* drgnType /* TODO: this argument should not be required */
-) {
+void CodeGen::generate(TypeGraph& typeGraph,
+                       std::string& code,
+                       RootFunctionName rootName) {
   code = headers::oi_OITraceCode_cpp;
   if (!config_.features[Feature::Library]) {
     FuncGen::DeclareExterns(code);
@@ -1296,22 +1303,33 @@ void CodeGen::generate(
   code += "\nusing __ROOT_TYPE__ = " + rootType.name() + ";\n";
   code += "} // namespace\n} // namespace OIInternal\n";
 
-  const auto typeName = SymbolService::getTypeName(drgnType);
+  const auto& typeToHash = std::visit(
+      [](const auto& v) -> const std::string& {
+        using T = std::decay_t<decltype(v)>;
+        if constexpr (std::is_same_v<ExactName, T> ||
+                      std::is_same_v<HashedComponent, T>) {
+          return v.name;
+        } else {
+          static_assert(always_false_v<T>, "missing visit");
+        }
+      },
+      rootName);
+
   if (config_.features[Feature::TreeBuilderV2]) {
-    FuncGen::DefineTopLevelIntrospect(code, typeName);
+    FuncGen::DefineTopLevelIntrospect(code, typeToHash);
   } else {
-    FuncGen::DefineTopLevelGetSizeRef(code, typeName, config_.features);
+    FuncGen::DefineTopLevelGetSizeRef(code, typeToHash, config_.features);
   }
 
   if (config_.features[Feature::TreeBuilderV2]) {
     FuncGen::DefineTreeBuilderInstructions(code,
-                                           typeName,
+                                           typeToHash,
                                            calculateExclusiveSize(rootType),
                                            enumerateTypeNames(rootType));
   }
 
-  if (!linkageName_.empty())
-    FuncGen::DefineTopLevelIntrospectNamed(code, typeName, linkageName_);
+  if (auto* n = std::get_if<ExactName>(&rootName))
+    FuncGen::DefineTopLevelIntrospectNamed(code, typeToHash, n->name);
 
   if (VLOG_IS_ON(3)) {
     VLOG(3) << "Generated trace code:\n";
diff --git a/oi/CodeGen.h b/oi/CodeGen.h
index 2c9b5e1..9c76238 100644
--- a/oi/CodeGen.h
+++ b/oi/CodeGen.h
@@ -45,6 +45,14 @@ class CodeGen {
       : config_(config), symbols_(symbols) {
   }
 
+  struct ExactName {
+    std::string name;
+  };
+  struct HashedComponent {
+    std::string name;
+  };
+  using RootFunctionName = std::variant<ExactName, HashedComponent>;
+
   /*
    * Helper function to perform all the steps required for code generation for a
    * single drgn_type.
@@ -64,9 +72,7 @@ class CodeGen {
   void transform(type_graph::TypeGraph& typeGraph);
   void generate(type_graph::TypeGraph& typeGraph,
                 std::string& code,
-                struct drgn_type*
-                    drgnType /* TODO: this argument should not be required */
-  );
+                RootFunctionName rootName);
 
  private:
   type_graph::TypeGraph typeGraph_;
@@ -76,7 +82,10 @@ class CodeGen {
   std::unordered_set<const ContainerInfo*> definedContainers_;
   std::unordered_map<const type_graph::Class*, const type_graph::Member*>
       thriftIssetMembers_;
-  std::string linkageName_;
+
+  bool codegenFromDrgn(struct drgn_type* drgnType,
+                       std::string& code,
+                       RootFunctionName name);
 
   void genDefsThrift(const type_graph::TypeGraph& typeGraph, std::string& code);
   void addGetSizeFuncDefs(const type_graph::TypeGraph& typeGraph,