From 27c6ee10c3f58268e9f240ad3329fd04f3585889 Mon Sep 17 00:00:00 2001
From: Alastair Robertson <ajor@meta.com>
Date: Wed, 24 May 2023 03:04:06 -0700
Subject: [PATCH] OID: Use new TypeGraph CodeGen when requested

This adds a new feature flag, "type-graph", which controls whether to
use legacy OICodeGen or the new TypeGraph-based CodeGen.
---
 oi/Features.h     |  1 +
 oi/OIDebugger.cpp | 19 +++++++++++++++----
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/oi/Features.h b/oi/Features.h
index 62ad7c0..1f12313 100644
--- a/oi/Features.h
+++ b/oi/Features.h
@@ -24,6 +24,7 @@
   X(PackStructs, "pack-structs")                \
   X(GenPaddingStats, "gen-padding-stats")       \
   X(CaptureThriftIsset, "capture-thrift-isset") \
+  X(TypeGraph, "type-graph")                    \
   X(PolymorphicInheritance, "polymorphic-inheritance")
 
 namespace ObjectIntrospection {
diff --git a/oi/OIDebugger.cpp b/oi/OIDebugger.cpp
index 794904c..b225653 100644
--- a/oi/OIDebugger.cpp
+++ b/oi/OIDebugger.cpp
@@ -44,6 +44,7 @@ extern "C" {
 
 #include <glog/logging.h>
 
+#include "oi/CodeGen.h"
 #include "oi/ContainerInfo.h"
 #include "oi/Headers.h"
 #include "oi/Metrics.h"
@@ -51,6 +52,7 @@ extern "C" {
 #include "oi/OIUtils.h"
 #include "oi/PaddingHunter.h"
 #include "oi/Syscall.h"
+#include "oi/type_graph/TypeGraph.h"
 
 #ifndef OSS_ENABLE
 #include "cea/object-introspection/internal/GobsService.h"
@@ -2918,15 +2920,24 @@ std::optional<std::string> OIDebugger::generateCode(const irequest& req) {
     return std::nullopt;
   }
 
-  if (auto sourcePath = cache.getPath(req, OICache::Entity::Source)) {
-    std::ofstream(*sourcePath) << code;
-  }
-
   typeInfos.emplace(
       req,
       std::make_tuple(RootInfo{rootInfo.varName, codegen->getRootType()},
                       codegen->getTypeHierarchy(), codegen->getPaddingInfo()));
 
+  if (generatorConfig.features[Feature::TypeGraph]) {
+    type_graph::TypeGraph typeGraph;
+    CodeGen codegen2(typeGraph, generatorConfig, *symbols);
+    codegen2.loadConfig(generatorConfig.containerConfigPaths);
+    if (!codegen2.generate(root->type.type, code)) {
+      return nullopt;
+    }
+  }
+
+  if (auto sourcePath = cache.getPath(req, OICache::Entity::Source)) {
+    std::ofstream(*sourcePath) << code;
+  }
+
   if (!customCodeFile.empty()) {
     auto ifs = std::ifstream(customCodeFile);
     code.assign(std::istreambuf_iterator<char>(ifs),