Type providers (#304)

* Bring `ast_unit`, `ast_ctx`, `ast` into `dec_ctx`

* Move `GetQualType` to `DecompilationContext`

* Add type provider

* Add type provider for `main`

* Update copyright

* Revert formatting change

* Emit non-signed `char`s

* Allow adding type providers from decomp API

* Fix merge error

* More merge errors

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

@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2022-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 <llvm/IR/DerivedTypes.h>
+#include <llvm/IR/Function.h>
+#include <llvm/IR/GlobalVariable.h>
+
+#include "rellic/AST/ASTBuilder.h"
+
+namespace rellic {
+struct DecompilationContext;
+
+class TypeProvider {
+ protected:
+  DecompilationContext& dec_ctx;
+
+ public:
+  TypeProvider(DecompilationContext& dec_ctx);
+  virtual ~TypeProvider();
+
+  // 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.
+  // 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.
+  // A null return value is assumed to mean that no info is available.
+  virtual clang::QualType GetGlobalVarType(llvm::GlobalVariable& gvar);
+};
+
+class TypeProviderCombiner : public TypeProvider {
+ private:
+  std::vector<std::unique_ptr<TypeProvider>> providers;
+
+ public:
+  TypeProviderCombiner(DecompilationContext& 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);
+
+  clang::QualType GetFunctionReturnType(llvm::Function& func) override;
+  clang::QualType GetArgumentType(llvm::Argument& arg) override;
+  clang::QualType GetGlobalVarType(llvm::GlobalVariable& gvar) override;
+};
+}  // namespace rellic

include/rellic/AST/Util.h

@@ -8,17 +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/DecompilationContext.h"
 
 namespace rellic {
 
@@ -52,60 +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;
-
-  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);
-};
-
 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 {