Allow adding type providers from decomp API

include/rellic/AST/DecompilationContext.h

@@ -0,0 +1,82 @@
+/*
+ * Copyright (c) 2021-present, Trail of Bits, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed in accordance with the terms specified in
+ * the LICENSE file found in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <clang/AST/ASTContext.h>
+#include <clang/Frontend/ASTUnit.h>
+#include <llvm/IR/Instructions.h>
+#include <llvm/IR/Value.h>
+#include <z3++.h>
+
+#include <unordered_map>
+
+#include "rellic/AST/ASTBuilder.h"
+#include "rellic/AST/TypeProvider.h"
+
+namespace rellic {
+
+struct DecompilationContext {
+  using StmtToIRMap = std::unordered_map<clang::Stmt *, llvm::Value *>;
+  using ExprToUseMap = std::unordered_map<clang::Expr *, llvm::Use *>;
+  using IRToTypeDeclMap = std::unordered_map<llvm::Type *, clang::TypeDecl *>;
+  using IRToValDeclMap = std::unordered_map<llvm::Value *, clang::ValueDecl *>;
+  using IRToStmtMap = std::unordered_map<llvm::Value *, clang::Stmt *>;
+  using ArgToTempMap = std::unordered_map<llvm::Argument *, clang::VarDecl *>;
+  using BlockToUsesMap =
+      std::unordered_map<llvm::BasicBlock *, std::vector<llvm::Use *>>;
+  using Z3CondMap = std::unordered_map<clang::Stmt *, unsigned>;
+
+  using BBEdge = std::pair<llvm::BasicBlock *, llvm::BasicBlock *>;
+  using BrEdge = std::pair<llvm::BranchInst *, bool>;
+  using SwEdge = std::pair<llvm::SwitchInst *, llvm::ConstantInt *>;
+
+  DecompilationContext(clang::ASTUnit &ast_unit);
+
+  clang::ASTUnit &ast_unit;
+  clang::ASTContext &ast_ctx;
+  ASTBuilder ast;
+
+  std::unique_ptr<TypeProviderCombiner> type_provider;
+
+  StmtToIRMap stmt_provenance;
+  ExprToUseMap use_provenance;
+  IRToTypeDeclMap type_decls;
+  IRToValDeclMap value_decls;
+  ArgToTempMap temp_decls;
+  BlockToUsesMap outgoing_uses;
+  z3::context z3_ctx;
+  z3::expr_vector z3_exprs{z3_ctx};
+  Z3CondMap conds;
+
+  clang::Expr *marker_expr;
+
+  std::unordered_map<unsigned, BrEdge> z3_br_edges_inv;
+
+  // Pairs do not have a std::hash specialization so we can't use unordered maps
+  // here. If this turns out to be a performance issue, investigate adding hash
+  // specializations for these specifically
+  std::map<BrEdge, unsigned> z3_br_edges;
+
+  std::unordered_map<llvm::SwitchInst *, unsigned> z3_sw_vars;
+  std::unordered_map<unsigned, llvm::SwitchInst *> z3_sw_vars_inv;
+  std::map<SwEdge, unsigned> z3_sw_edges;
+
+  std::map<BBEdge, unsigned> z3_edges;
+  std::unordered_map<llvm::BasicBlock *, unsigned> reaching_conds;
+
+  size_t num_literal_structs = 0;
+  size_t num_declared_structs = 0;
+
+  // Inserts an expression into z3_exprs and returns its index
+  unsigned InsertZExpr(const z3::expr &e);
+
+  clang::QualType GetQualType(llvm::Type *type);
+};
+
+}  // namespace rellic

include/rellic/AST/TypeProvider.h

@@ -24,13 +24,16 @@ class TypeProvider {
   TypeProvider(DecompilationContext& dec_ctx);
   virtual ~TypeProvider();
 
-  // Returns the return type of a function if available
+  // Returns the return type of a function if available.
+  // A null return value is assumed to mean that no info is available.
   virtual clang::QualType GetFunctionReturnType(llvm::Function& func);
 
-  // Returns the type of the argument if available
+  // Returns the type of the argument if available.
+  // A null return value is assumed to mean that no info is available.
   virtual clang::QualType GetArgumentType(llvm::Argument& arg);
 
-  // Returns the type of a global variable if available
+  // Returns the type of a global variable if available.
+  // A null return value is assumed to mean that no info is available.
   virtual clang::QualType GetGlobalVarType(llvm::GlobalVariable& gvar);
 };
 
@@ -40,9 +43,10 @@ class TypeProviderCombiner : public TypeProvider {
 
  public:
   TypeProviderCombiner(DecompilationContext& dec_ctx);
-  template <typename T>
-  void AddProvider() {
-    providers.push_back(std::make_unique<T>(dec_ctx));
+  template <typename T, typename... TArgs>
+  void AddProvider(TArgs&&... args) {
+    providers.push_back(
+        std::make_unique<T>(dec_ctx, std::forward<TArgs>(args)...));
   }
 
   void AddProvider(std::unique_ptr<TypeProvider> provider);

include/rellic/AST/Util.h

@@ -8,18 +8,7 @@
 
 #pragma once
 
-#include <clang/AST/ASTContext.h>
-#include <clang/AST/DeclBase.h>
-#include <clang/AST/Stmt.h>
-#include <clang/Frontend/ASTUnit.h>
-#include <llvm/IR/Instructions.h>
-#include <llvm/IR/Value.h>
-#include <z3++.h>
-
-#include <unordered_map>
-
-#include "rellic/AST/ASTBuilder.h"
-#include "rellic/AST/TypeProvider.h"
+#include "rellic/AST/DecompilationContext.h"
 
 namespace rellic {
 
@@ -53,64 +42,6 @@ size_t GetNumDecls(clang::DeclContext *decl_ctx) {
   return result;
 }
 
-struct DecompilationContext {
-  using StmtToIRMap = std::unordered_map<clang::Stmt *, llvm::Value *>;
-  using ExprToUseMap = std::unordered_map<clang::Expr *, llvm::Use *>;
-  using IRToTypeDeclMap = std::unordered_map<llvm::Type *, clang::TypeDecl *>;
-  using IRToValDeclMap = std::unordered_map<llvm::Value *, clang::ValueDecl *>;
-  using IRToStmtMap = std::unordered_map<llvm::Value *, clang::Stmt *>;
-  using ArgToTempMap = std::unordered_map<llvm::Argument *, clang::VarDecl *>;
-  using BlockToUsesMap =
-      std::unordered_map<llvm::BasicBlock *, std::vector<llvm::Use *>>;
-  using Z3CondMap = std::unordered_map<clang::Stmt *, unsigned>;
-
-  using BBEdge = std::pair<llvm::BasicBlock *, llvm::BasicBlock *>;
-  using BrEdge = std::pair<llvm::BranchInst *, bool>;
-  using SwEdge = std::pair<llvm::SwitchInst *, llvm::ConstantInt *>;
-
-  DecompilationContext(clang::ASTUnit &ast_unit);
-
-  clang::ASTUnit &ast_unit;
-  clang::ASTContext &ast_ctx;
-  ASTBuilder ast;
-
-  std::unique_ptr<TypeProviderCombiner> type_provider;
-
-  StmtToIRMap stmt_provenance;
-  ExprToUseMap use_provenance;
-  IRToTypeDeclMap type_decls;
-  IRToValDeclMap value_decls;
-  ArgToTempMap temp_decls;
-  BlockToUsesMap outgoing_uses;
-  z3::context z3_ctx;
-  z3::expr_vector z3_exprs{z3_ctx};
-  Z3CondMap conds;
-
-  clang::Expr *marker_expr;
-
-  std::unordered_map<unsigned, BrEdge> z3_br_edges_inv;
-
-  // Pairs do not have a std::hash specialization so we can't use unordered maps
-  // here. If this turns out to be a performance issue, investigate adding hash
-  // specializations for these specifically
-  std::map<BrEdge, unsigned> z3_br_edges;
-
-  std::unordered_map<llvm::SwitchInst *, unsigned> z3_sw_vars;
-  std::unordered_map<unsigned, llvm::SwitchInst *> z3_sw_vars_inv;
-  std::map<SwEdge, unsigned> z3_sw_edges;
-
-  std::map<BBEdge, unsigned> z3_edges;
-  std::unordered_map<llvm::BasicBlock *, unsigned> reaching_conds;
-
-  size_t num_literal_structs = 0;
-  size_t num_declared_structs = 0;
-
-  // Inserts an expression into z3_exprs and returns its index
-  unsigned InsertZExpr(const z3::expr &e);
-
-  clang::QualType GetQualType(llvm::Type *type);
-};
-
 template <typename TKey1, typename TKey2, typename TKey3, typename TValue>
 void CopyProvenance(TKey1 *from, TKey2 *to,
                     std::unordered_map<TKey3 *, TValue *> &map) {

include/rellic/Decompiler.h

@@ -13,13 +13,39 @@
 
 #include <memory>
 #include <unordered_map>
+#include <vector>
 
 #include "Result.h"
+#include "rellic/AST/TypeProvider.h"
 
 namespace rellic {
+
+/* This additional level of indirection is needed to alleviate the users from
+ * the burden of having to instantiate custom TypeProviders before the actual
+ * DecompilationContext has been created */
+class TypeProviderFactory {
+ public:
+  virtual ~TypeProviderFactory() = default;
+  virtual std::unique_ptr<TypeProvider> create(DecompilationContext& ctx) = 0;
+};
+
+template <typename T>
+class SimpleTypeProviderFactory final : public TypeProviderFactory {
+ public:
+  std::unique_ptr<TypeProvider> create(DecompilationContext& ctx) override {
+    return std::make_unique<T>(ctx);
+  }
+};
+
 struct DecompilationOptions {
+  using TypeProviderFactoryPtr = std::unique_ptr<TypeProviderFactory>;
+
   bool lower_switches = false;
   bool remove_phi_nodes = false;
+
+  // Additional type providers to be used during code generation.
+  // Providers added later will have higher priority.
+  std::vector<TypeProviderFactoryPtr> additional_providers;
 };
 
 struct DecompilationResult {

lib/Decompiler.cpp

@@ -83,6 +83,11 @@ Result<DecompilationResult, DecompilationError> Decompile(
                                   module->getTargetTriple()};
     auto ast_unit{clang::tooling::buildASTFromCodeWithArgs("", args, "out.c")};
     rellic::DecompilationContext dec_ctx(*ast_unit);
+
+    for (auto& provider : options.additional_providers) {
+      dec_ctx.type_provider->AddProvider(provider->create(dec_ctx));
+    }
+
     rellic::GenerateAST::run(*module, dec_ctx);
     // TODO(surovic): Add llvm::Value* -> clang::Decl* map
     // Especially for llvm::Argument* and llvm::Function*.

tools/decomp/Decomp.cpp

@@ -124,7 +124,7 @@ int main(int argc, char* argv[]) {
   opts.lower_switches = FLAGS_lower_switch;
   opts.remove_phi_nodes = FLAGS_remove_phi_nodes;
 
-  auto result{rellic::Decompile(std::move(module), opts)};
+  auto result{rellic::Decompile(std::move(module), std::move(opts))};
   if (result.Succeeded()) {
     auto value{result.TakeValue()};
     value.ast->getASTContext().getTranslationUnitDecl()->print(output);