TypeGraph: Split CodeGen into separate functions for testing

oi/CMakeLists.txt

@@ -42,6 +42,7 @@ add_library(codegen
 )
 target_link_libraries(codegen
   container_info
+  resources
   symbol_service
   type_graph
 

oi/CodeGen.cpp

@@ -24,6 +24,7 @@
 
 #include "oi/FuncGen.h"
 #include "oi/Headers.h"
+#include "oi/SymbolService.h"
 #include "type_graph/AddChildren.h"
 #include "type_graph/AddPadding.h"
 #include "type_graph/AlignmentCalc.h"
@@ -545,21 +546,51 @@ void addTypeHandlers(const TypeGraph& typeGraph, std::string& code) {
 }
 }  // namespace
 
-bool CodeGen::generate(drgn_type* drgnType, std::string& code) {
-  type_graph::DrgnParser drgnParser{
-      typeGraph_, containerInfos_, config_.features[Feature::ChaseRawPointers]};
+bool CodeGen::codegenFromDrgn(struct drgn_type* drgnType, std::string& code) {
   try {
-    Type* parsedRoot = drgnParser.parse(drgnType);
-    typeGraph_.addRoot(*parsedRoot);
+    containerInfos_.reserve(config_.containerConfigPaths.size());
+    for (const auto& path : config_.containerConfigPaths) {
+      registerContainer(path);
+    }
+  } catch (const ContainerInfoError& err) {
+    LOG(ERROR) << "Error reading container TOML file " << err.what();
+    return false;
+  }
+
+  type_graph::TypeGraph typeGraph;
+  try {
+    addDrgnRoot(drgnType, typeGraph);
   } catch (const type_graph::DrgnParserError& err) {
     LOG(ERROR) << "Error parsing DWARF: " << err.what();
     return false;
   }
 
+  transform(typeGraph);
+  generate(typeGraph, code, drgnType);
+  return true;
+}
+
+void CodeGen::registerContainer(const fs::path& path) {
+  const auto& info = containerInfos_.emplace_back(path);
+  VLOG(1) << "Registered container: " << info.typeName;
+}
+
+void CodeGen::addDrgnRoot(struct drgn_type* drgnType,
+                          type_graph::TypeGraph& typeGraph) {
+  type_graph::DrgnParser drgnParser{
+      typeGraph, containerInfos_, config_.features[Feature::ChaseRawPointers]};
+  Type* parsedRoot = drgnParser.parse(drgnType);
+  typeGraph.addRoot(*parsedRoot);
+}
+
+void CodeGen::transform(type_graph::TypeGraph& typeGraph) {
   type_graph::PassManager pm;
   pm.addPass(type_graph::Flattener::createPass());
   pm.addPass(type_graph::TypeIdentifier::createPass());
   if (config_.features[Feature::PolymorphicInheritance]) {
+    type_graph::DrgnParser drgnParser{
+        typeGraph, containerInfos_,
+        config_.features[Feature::ChaseRawPointers]};
     pm.addPass(type_graph::AddChildren::createPass(drgnParser, symbols_));
     // Re-run passes over newly added children
     pm.addPass(type_graph::Flattener::createPass());
@@ -571,18 +602,24 @@ bool CodeGen::generate(drgn_type* drgnType, std::string& code) {
   pm.addPass(type_graph::AlignmentCalc::createPass());
   pm.addPass(type_graph::RemoveTopLevelPointer::createPass());
   pm.addPass(type_graph::TopoSorter::createPass());
-  pm.run(typeGraph_);
+  pm.run(typeGraph);
 
   LOG(INFO) << "Sorted types:\n";
-  for (Type& t : typeGraph_.finalTypes) {
+  for (Type& t : typeGraph.finalTypes) {
     LOG(INFO) << "  " << t.name() << std::endl;
   };
+}
 
+void CodeGen::generate(
+    type_graph::TypeGraph& typeGraph,
+    std::string& code,
+    struct drgn_type* drgnType /* TODO: this argument should not be required */
+) {
   code = headers::OITraceCode_cpp;
   if (!config_.features[Feature::TypedDataSegment]) {
     defineMacros(code);
   }
-  addIncludes(typeGraph_, config_.features, code);
+  addIncludes(typeGraph, config_.features, code);
   defineArray(code);
   defineJitLog(code);  // TODO: feature gate this
 
@@ -614,24 +651,24 @@ bool CodeGen::generate(drgn_type* drgnType, std::string& code) {
   }
   FuncGen::DeclareGetContainer(code);
 
-  genDecls(typeGraph_, code);
-  genDefs(typeGraph_, code);
-  genStaticAsserts(typeGraph_, code);
+  genDecls(typeGraph, code);
+  genDefs(typeGraph, code);
+  genStaticAsserts(typeGraph, code);
 
   if (config_.features[Feature::TypedDataSegment]) {
     addStandardTypeHandlers(code);
-    addTypeHandlers(typeGraph_, code);
+    addTypeHandlers(typeGraph, code);
   } else {
     addStandardGetSizeFuncDecls(code);
-    addGetSizeFuncDecls(typeGraph_, code);
+    addGetSizeFuncDecls(typeGraph, code);
 
     addStandardGetSizeFuncDefs(code);
-    addGetSizeFuncDefs(typeGraph_, symbols_,
+    addGetSizeFuncDefs(typeGraph, symbols_,
                        config_.features[Feature::PolymorphicInheritance], code);
   }
 
-  assert(typeGraph_.rootTypes().size() == 1);
-  Type& rootType = typeGraph_.rootTypes()[0];
+  assert(typeGraph.rootTypes().size() == 1);
+  Type& rootType = typeGraph.rootTypes()[0];
   code += "\nusing __ROOT_TYPE__ = " + rootType.name() + ";\n";
   code += "} // namespace\n} // namespace OIInternal\n";
 
@@ -648,22 +685,4 @@ bool CodeGen::generate(drgn_type* drgnType, std::string& code) {
     // VLOG truncates output, so use std::cout
     std::cout << code;
   }
-  return true;
-}
-
-void CodeGen::loadConfig(const std::set<fs::path>& containerConfigPaths) {
-  containerInfos_.reserve(containerConfigPaths.size());
-  for (const auto& path : containerConfigPaths) {
-    registerContainer(path);
-  }
-}
-
-void CodeGen::registerContainer(const fs::path& path) {
-  try {
-    const auto& info = containerInfos_.emplace_back(path);
-    VLOG(1) << "Registered container: " << info.typeName;
-  } catch (const std::runtime_error& err) {
-    LOG(ERROR) << "Error reading container TOML file " << path << ": "
-               << err.what();
-  }
 }

oi/CodeGen.h

@@ -21,33 +21,39 @@
 
 #include "ContainerInfo.h"
 #include "OICodeGen.h"
-#include "SymbolService.h"
 
 struct drgn_type;
 
+class SymbolService;
+
 namespace type_graph {
-class Class;
-class Container;
-class Type;
 class TypeGraph;
 }  // namespace type_graph
 
 class CodeGen {
  public:
-  CodeGen(type_graph::TypeGraph& typeGraph,
-          OICodeGen::Config& config,
-          SymbolService& symbols)
-      : typeGraph_(typeGraph), config_(config), symbols_(symbols) {
+  CodeGen(OICodeGen::Config& config, SymbolService& symbols)
+      : config_(config), symbols_(symbols) {
   }
 
-  bool generate(drgn_type* drgnType, std::string& code);
-  void loadConfig(const std::set<std::filesystem::path>& containerConfigPaths);
+  /*
+   * Helper function to perform all the steps required for code generation for a
+   * single drgn_type.
+   */
+  bool codegenFromDrgn(struct drgn_type* drgnType, std::string& code);
+
+  void registerContainer(const std::filesystem::path& path);
+  void addDrgnRoot(struct drgn_type* drgnType,
+                   type_graph::TypeGraph& typeGraph);
+  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 */
+  );
 
  private:
-  void registerContainer(const std::filesystem::path& path);
-
-  type_graph::TypeGraph& typeGraph_;
   OICodeGen::Config config_;
-  std::vector<ContainerInfo> containerInfos_;
   SymbolService& symbols_;
+  std::vector<ContainerInfo> containerInfos_;
 };

oi/ContainerInfo.cpp

@@ -189,16 +189,17 @@ ContainerInfo::ContainerInfo(const fs::path& path) {
   try {
     container = toml::parse_file(std::string(path));
   } catch (const toml::parse_error& err) {
-    // Convert into a std::runtime_error, just to avoid having to include
+    // Convert into a ContainerInfoError, just to avoid having to include
     // the huge TOML++ header in the caller's file. Use toml::parse_error's
     // operator<< to generate a pretty message with error location.
     std::stringstream ss;
     ss << err;
-    throw std::runtime_error(ss.str());
+    throw ContainerInfoError(path, ss.str());
   }
 
   if (!container["info"].is_table()) {
-    throw std::runtime_error("a container info file requires an `info` table");
+    throw ContainerInfoError(path,
+                             "a container info file requires an `info` table");
   }
 
   const auto& info = container["info"];
@@ -206,7 +207,7 @@ ContainerInfo::ContainerInfo(const fs::path& path) {
   if (std::optional<std::string> str = info["type_name"].value<std::string>()) {
     typeName = std::move(*str);
   } else {
-    throw std::runtime_error("`info.type_name` is a required field");
+    throw ContainerInfoError(path, "`info.type_name` is a required field");
   }
 
   matcher = getMatcher(typeName);
@@ -214,16 +215,17 @@ ContainerInfo::ContainerInfo(const fs::path& path) {
   if (std::optional<std::string> str = info["ctype"].value<std::string>()) {
     ctype = containerTypeEnumFromStr(*str);
     if (ctype == UNKNOWN_TYPE) {
-      throw std::runtime_error("`" + *str + "` is not a valid container type");
+      throw ContainerInfoError(path,
+                               "`" + *str + "` is not a valid container type");
     }
   } else {
-    throw std::runtime_error("`info.ctype` is a required field");
+    throw ContainerInfoError(path, "`info.ctype` is a required field");
   }
 
   if (std::optional<std::string> str = info["header"].value<std::string>()) {
     header = std::move(*str);
   } else {
-    throw std::runtime_error("`info.header` is a required field");
+    throw ContainerInfoError(path, "`info.header` is a required field");
   }
 
   if (toml::array* arr = info["stub_template_params"].as_array()) {
@@ -232,8 +234,8 @@ ContainerInfo::ContainerInfo(const fs::path& path) {
       if constexpr (toml::is_integer<decltype(el)>) {
         stubTemplateParams.push_back(*el);
       } else {
-        throw std::runtime_error(
-            "stub_template_params should only contain integers");
+        throw ContainerInfoError(
+            path, "stub_template_params should only contain integers");
       }
     });
   }
@@ -241,8 +243,8 @@ ContainerInfo::ContainerInfo(const fs::path& path) {
   underlyingContainerIndex = info["underlying_container_index"].value<size_t>();
 
   if (!container["codegen"].is_table()) {
-    throw std::runtime_error(
-        "a container info file requires a `codegen` table");
+    throw ContainerInfoError(
+        path, "a container info file requires a `codegen` table");
   }
 
   const auto& codegenToml = container["codegen"];
@@ -251,13 +253,13 @@ ContainerInfo::ContainerInfo(const fs::path& path) {
           codegenToml["func"].value<std::string>()) {
     codegen.func = std::move(*str);
   } else {
-    throw std::runtime_error("`codegen.func` is a required field");
+    throw ContainerInfoError(path, "`codegen.func` is a required field");
   }
   if (std::optional<std::string> str =
           codegenToml["decl"].value<std::string>()) {
     codegen.decl = std::move(*str);
   } else {
-    throw std::runtime_error("`codegen.decl` is a required field");
+    throw ContainerInfoError(path, "`codegen.decl` is a required field");
   }
   if (std::optional<std::string> str =
           codegenToml["handler"].value<std::string>()) {

oi/ContainerInfo.h

@@ -93,6 +93,13 @@ struct ContainerInfo {
   }
 };
 
+class ContainerInfoError : public std::runtime_error {
+ public:
+  ContainerInfoError(const std::filesystem::path& path, const std::string& msg)
+      : std::runtime_error{std::string{path} + ": " + msg} {
+  }
+};
+
 using ContainerInfoRefSet =
     std::set<std::reference_wrapper<const ContainerInfo>,
              std::less<ContainerInfo>>;

oi/OIDebugger.cpp

@@ -52,6 +52,7 @@ extern "C" {
 #include "oi/OIUtils.h"
 #include "oi/PaddingHunter.h"
 #include "oi/Syscall.h"
+#include "oi/type_graph/DrgnParser.h"
 #include "oi/type_graph/TypeGraph.h"
 
 #ifndef OSS_ENABLE
@@ -2928,12 +2929,8 @@ std::optional<std::string> OIDebugger::generateCode(const irequest& req) {
                       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;
-    }
+    CodeGen codegen2{generatorConfig, *symbols};
+    codegen2.codegenFromDrgn(root->type.type, code);
   }
 
   if (auto sourcePath = cache.getPath(req, OICache::Entity::Source)) {