Move more logic into the function layout override

include/rellic/AST/FunctionLayoutOverride.h

@@ -33,6 +33,7 @@ class FunctionLayoutOverride {
   virtual bool VisitInstruction(llvm::Instruction& insn,
                                 clang::FunctionDecl* fdecl,
                                 clang::ValueDecl*& vdecl);
+  virtual bool NeedsDereference(llvm::Function& func, llvm::Value& val);
 };
 
 class FunctionLayoutOverrideCombiner final : public FunctionLayoutOverride {
@@ -56,5 +57,6 @@ class FunctionLayoutOverrideCombiner final : public FunctionLayoutOverride {
                           clang::FunctionDecl* fdecl) final;
   bool VisitInstruction(llvm::Instruction& insn, clang::FunctionDecl* fdecl,
                         clang::ValueDecl*& vdecl) final;
+  bool NeedsDereference(llvm::Function& func, llvm::Value& val) final;
 };
 }  // namespace rellic

include/rellic/AST/IRToASTVisitor.h

@@ -30,8 +30,6 @@ class IRToASTVisitor {
   DecompilationContext &dec_ctx;
   ASTBuilder *
 
-  void VisitArgument(llvm::Argument &arg);
-
  public:
   IRToASTVisitor(DecompilationContext &dec_ctx);
 

lib/AST/FunctionLayoutOverride.cpp

@@ -11,6 +11,7 @@
 #include <clang/AST/Type.h>
 #include <llvm/IR/Argument.h>
 
+#include "rellic/AST/DecompilationContext.h"
 #include "rellic/AST/Util.h"
 
 namespace rellic {
@@ -34,9 +35,65 @@ bool FunctionLayoutOverride::VisitInstruction(llvm::Instruction &insn,
   return false;
 }
 
+bool FunctionLayoutOverride::NeedsDereference(llvm::Function &func,
+                                              llvm::Value &val) {
+  return false;
+}
+
+class FallbackFunctionLayoutOverride : public FunctionLayoutOverride {
+ public:
+  FallbackFunctionLayoutOverride(DecompilationContext &dec_ctx)
+      : FunctionLayoutOverride(dec_ctx) {}
+
+  bool HasOverride(llvm::Function &func) final { return true; }
+
+  std::vector<clang::QualType> GetArguments(llvm::Function &func) final {
+    std::vector<clang::QualType> arg_types;
+    for (auto &arg : func.args()) {
+      arg_types.push_back(dec_ctx.type_provider->GetArgumentType(arg));
+    }
+    return arg_types;
+  }
+
+  void BeginFunctionVisit(llvm::Function &func,
+                          clang::FunctionDecl *fdecl) final {
+    std::vector<clang::ParmVarDecl *> params;
+    for (auto &arg : func.args()) {
+      auto &parm{dec_ctx.value_decls[&arg]};
+      if (parm) {
+        return;
+      }
+      // Create a name
+      auto name{arg.hasName() ? arg.getName().str()
+                              : "arg" + std::to_string(arg.getArgNo())};
+      // Get parent function declaration
+      auto func{arg.getParent()};
+      auto fdecl{clang::cast<clang::FunctionDecl>(dec_ctx.value_decls[func])};
+      auto argtype = dec_ctx.type_provider->GetArgumentType(arg);
+      // Create a declaration
+      parm = dec_ctx.ast.CreateParamDecl(fdecl, argtype, name);
+      params.push_back(
+          clang::dyn_cast<clang::ParmVarDecl>(dec_ctx.value_decls[&arg]));
+    }
+
+    fdecl->setParams(params);
+  }
+
+  bool VisitInstruction(llvm::Instruction &inst, clang::FunctionDecl *fdecl,
+                        clang::ValueDecl *&vdecl) final {
+    return false;
+  }
+
+  bool NeedsDereference(llvm::Function &func, llvm::Value &val) final {
+    return llvm::isa<llvm::AllocaInst>(val);
+  }
+};
+
 FunctionLayoutOverrideCombiner::FunctionLayoutOverrideCombiner(
     DecompilationContext &dec_ctx)
-    : FunctionLayoutOverride(dec_ctx) {}
+    : FunctionLayoutOverride(dec_ctx) {
+  AddOverride<FallbackFunctionLayoutOverride>();
+}
 
 void FunctionLayoutOverrideCombiner::AddOverride(
     std::unique_ptr<FunctionLayoutOverride> provider) {
@@ -87,4 +144,15 @@ bool FunctionLayoutOverrideCombiner::VisitInstruction(
   return false;
 }
 
+bool FunctionLayoutOverrideCombiner::NeedsDereference(llvm::Function &func,
+                                                      llvm::Value &val) {
+  for (auto it{overrides.rbegin()}; it != overrides.rend(); ++it) {
+    auto &override{*it};
+    if (override->HasOverride(func)) {
+      return override->NeedsDereference(func, val);
+    }
+  }
+  return false;
+}
+
 }  // namespace rellic

lib/AST/IRToASTVisitor.cpp

@@ -346,10 +346,15 @@ clang::Expr *ExprGen::CreateOperandExpr(llvm::Use &val) {
   }};
 
   clang::Expr *res{nullptr};
+  llvm::Instruction *maybe_inst{llvm::dyn_cast<llvm::Instruction>(&val)};
+  llvm::Argument *maybe_arg{llvm::dyn_cast<llvm::Argument>(&val)};
   if (auto constant = llvm::dyn_cast<llvm::Constant>(val)) {
     // Operand is a constant value
     res = CreateConstantExpr(constant);
-  } else if (llvm::isa<llvm::AllocaInst>(val)) {
+  } else if ((maybe_inst && dec_ctx.function_layout_override->NeedsDereference(
+                                *maybe_inst->getFunction(), *maybe_inst)) ||
+             (maybe_arg && dec_ctx.function_layout_override->NeedsDereference(
+                               *maybe_arg->getParent(), *maybe_arg))) {
     // Operand is an l-value (variable, function, ...)
     // Add a `&` operator
     res = ast.CreateAddrOf(CreateRef());
@@ -1077,23 +1082,6 @@ void IRToASTVisitor::VisitGlobalVar(llvm::GlobalVariable &gvar) {
   expr_gen.VisitGlobalVar(gvar);
 }
 
-void IRToASTVisitor::VisitArgument(llvm::Argument &arg) {
-  DLOG(INFO) << "VisitArgument: " << LLVMThingToString(&arg);
-  auto &parm{dec_ctx.value_decls[&arg]};
-  if (parm) {
-    return;
-  }
-  // Create a name
-  auto name{arg.hasName() ? arg.getName().str()
-                          : "arg" + std::to_string(arg.getArgNo())};
-  // Get parent function declaration
-  auto func{arg.getParent()};
-  auto fdecl{clang::cast<clang::FunctionDecl>(dec_ctx.value_decls[func])};
-  auto argtype = dec_ctx.type_provider->GetArgumentType(arg);
-  // Create a declaration
-  parm = ast.CreateParamDecl(fdecl, argtype, name);
-}
-
 void IRToASTVisitor::VisitBasicBlock(llvm::BasicBlock &block,
                                      std::vector<clang::Stmt *> &stmts) {
   ExprGen expr_gen{dec_ctx};
@@ -1132,17 +1120,7 @@ void IRToASTVisitor::VisitFunctionDecl(llvm::Function &func) {
   DLOG(INFO) << "Creating FunctionDecl for " << name;
   auto tudecl{dec_ctx.ast_ctx.getTranslationUnitDecl()};
 
-  bool override_function_layout =
-      dec_ctx.function_layout_override->HasOverride(func);
-
-  std::vector<clang::QualType> arg_types;
-  if (override_function_layout) {
-    arg_types = dec_ctx.function_layout_override->GetArguments(func);
-  } else {
-    for (auto &arg : func.args()) {
-      arg_types.push_back(dec_ctx.type_provider->GetArgumentType(arg));
-    }
-  }
+  auto arg_types{dec_ctx.function_layout_override->GetArguments(func)};
   auto ret_type{dec_ctx.type_provider->GetFunctionReturnType(func)};
   clang::FunctionProtoType::ExtProtoInfo epi;
   epi.Variadic = func.isVarArg();
@@ -1151,18 +1129,7 @@ void IRToASTVisitor::VisitFunctionDecl(llvm::Function &func) {
 
   tudecl->addDecl(decl);
   auto fdecl{decl->getAsFunction()};
-  if (override_function_layout) {
-    dec_ctx.function_layout_override->BeginFunctionVisit(func, fdecl);
-  } else {
-    std::vector<clang::ParmVarDecl *> params;
-    for (auto &arg : func.args()) {
-      VisitArgument(arg);
-      params.push_back(
-          clang::dyn_cast<clang::ParmVarDecl>(dec_ctx.value_decls[&arg]));
-    }
-
-    fdecl->setParams(params);
-  }
+  dec_ctx.function_layout_override->BeginFunctionVisit(func, fdecl);
 
   for (auto &inst : llvm::instructions(func)) {
     auto &var{dec_ctx.value_decls[&inst]};