Bring ast_unit, ast_ctx, ast into dec_ctx (#303)

* Bring `ast_unit`, `ast_ctx`, `ast` into `dec_ctx` * Remove amount of `dec_ctx` indirection * More removals * Refactor `GenerateAST`
lifting-bits · Sep 26, 2022 · 05c2835 · 05c2835
1 parent 2fffd6b
commit 05c2835
Show file tree

Hide file tree

Showing 35 changed files with 231 additions and 315 deletions.
diff --git a/include/rellic/AST/ASTPass.h b/include/rellic/AST/ASTPass.h
@@ -22,21 +22,15 @@ class ASTPass {
 
  protected:
   DecompilationContext& dec_ctx;
-  clang::ASTUnit& ast_unit;
-  clang::ASTContext& ast_ctx;
-  ASTBuilder ast;
 
   bool changed{false};
 
   virtual void RunImpl() = 0;
   virtual void StopImpl() {}
 
  public:
-  ASTPass(DecompilationContext& dec_ctx, clang::ASTUnit& ast_unit)
-      : dec_ctx(dec_ctx),
-        ast_unit(ast_unit),
-        ast_ctx(ast_unit.getASTContext()),
-        ast(ast_unit) {}
+  ASTPass(DecompilationContext& dec_ctx)
+      : dec_ctx(dec_ctx) {}
   virtual ~ASTPass() = default;
   void Stop() {
     stop = true;
@@ -89,8 +83,8 @@ class CompositeASTPass : public ASTPass {
   }
 
  public:
-  CompositeASTPass(DecompilationContext& dec_ctx, clang::ASTUnit& ast_unit)
-      : ASTPass(dec_ctx, ast_unit) {}
+  CompositeASTPass(DecompilationContext& dec_ctx)
+      : ASTPass(dec_ctx) {}
   std::vector<std::unique_ptr<ASTPass>>& GetPasses() { return passes; }
 };
 }  // namespace rellic
diff --git a/include/rellic/AST/CondBasedRefine.h b/include/rellic/AST/CondBasedRefine.h
@@ -38,7 +38,7 @@ class CondBasedRefine : public TransformVisitor<CondBasedRefine> {
   void RunImpl() override;
 
  public:
-  CondBasedRefine(DecompilationContext &dec_ctx, clang::ASTUnit &unit);
+  CondBasedRefine(DecompilationContext &dec_ctx);
 
   bool VisitCompoundStmt(clang::CompoundStmt *compound);
 };

diff --git a/include/rellic/AST/DeadStmtElim.h b/include/rellic/AST/DeadStmtElim.h
@@ -22,7 +22,7 @@ class DeadStmtElim : public TransformVisitor<DeadStmtElim> {
   void RunImpl() override;
 
  public:
-  DeadStmtElim(DecompilationContext &dec_ctx, clang::ASTUnit &unit);
+  DeadStmtElim(DecompilationContext &dec_ctx);
 
   bool VisitIfStmt(clang::IfStmt *ifstmt);
   bool VisitCompoundStmt(clang::CompoundStmt *compound);

diff --git a/include/rellic/AST/ExprCombine.h b/include/rellic/AST/ExprCombine.h
@@ -21,7 +21,7 @@ class ExprCombine : public TransformVisitor<ExprCombine> {
   void RunImpl() override;
 
  public:
-  ExprCombine(DecompilationContext &dec_ctx, clang::ASTUnit &unit);
+  ExprCombine(DecompilationContext &dec_ctx);
 
   bool VisitCStyleCastExpr(clang::CStyleCastExpr *cast);
   bool VisitUnaryOperator(clang::UnaryOperator *op);

diff --git a/include/rellic/AST/GenerateAST.h b/include/rellic/AST/GenerateAST.h
@@ -17,6 +17,7 @@
 #include <map>
 #include <unordered_set>
 
+#include "rellic/AST/ASTBuilder.h"
 #include "rellic/AST/IRToASTVisitor.h"
 
 namespace rellic {
@@ -29,11 +30,9 @@ class GenerateAST : public llvm::AnalysisInfoMixin<GenerateAST> {
   constexpr static unsigned poison_idx = std::numeric_limits<unsigned>::max();
   z3::expr ToExpr(unsigned idx);
 
-  clang::ASTUnit &unit;
-  clang::ASTContext *ast_ctx;
   rellic::IRToASTVisitor ast_gen;
-  rellic::ASTBuilder ast;
   DecompilationContext &dec_ctx;
+  ASTBuilder &ast;
   bool reaching_conds_changed{true};
   std::unordered_map<llvm::BasicBlock *, clang::IfStmt *> block_stmts;
   std::unordered_map<llvm::Region *, clang::CompoundStmt *> region_stmts;
@@ -78,13 +77,12 @@ class GenerateAST : public llvm::AnalysisInfoMixin<GenerateAST> {
 
  public:
   using Result = llvm::PreservedAnalyses;
-  GenerateAST(DecompilationContext &dec_ctx, clang::ASTUnit &unit);
+  GenerateAST(DecompilationContext &dec_ctx);
 
   Result run(llvm::Function &F, llvm::FunctionAnalysisManager &FAM);
   Result run(llvm::Module &M, llvm::ModuleAnalysisManager &MAM);
 
-  static void run(llvm::Module &M, DecompilationContext &dec_ctx,
-                  clang::ASTUnit &unit);
+  static void run(llvm::Module &M, DecompilationContext &dec_ctx);
 };
 
 }  // namespace rellic
diff --git a/include/rellic/AST/IRToASTVisitor.h b/include/rellic/AST/IRToASTVisitor.h
@@ -27,17 +27,13 @@
 namespace rellic {
 class IRToASTVisitor {
  private:
-  clang::ASTUnit &ast_unit;
-  clang::ASTContext &ast_ctx;
-
-  ASTBuilder ast;
-
   DecompilationContext &dec_ctx;
+  ASTBuilder &ast;
 
   void VisitArgument(llvm::Argument &arg);
 
  public:
-  IRToASTVisitor(clang::ASTUnit &unit, DecompilationContext &dec_ctx);
+  IRToASTVisitor(DecompilationContext &dec_ctx);
 
   clang::Expr *CreateOperandExpr(llvm::Use &val);
   clang::Expr *CreateConstantExpr(llvm::Constant *constant);

diff --git a/include/rellic/AST/InferenceRule.h b/include/rellic/AST/InferenceRule.h
@@ -35,12 +35,11 @@ class InferenceRule : public clang::ast_matchers::MatchFinder::MatchCallback {
   }
 
   virtual clang::Stmt *GetOrCreateSubstitution(DecompilationContext &dec_ctx,
-                                               clang::ASTUnit &unit,
                                                clang::Stmt *stmt) = 0;
 };
 
 clang::Stmt *ApplyFirstMatchingRule(
-    DecompilationContext &dec_ctx, clang::ASTUnit &unit, clang::Stmt *stmt,
+    DecompilationContext &dec_ctx, clang::Stmt *stmt,
     std::vector<std::unique_ptr<InferenceRule>> &rules);
 
 }  // namespace rellic
diff --git a/include/rellic/AST/LocalDeclRenamer.h b/include/rellic/AST/LocalDeclRenamer.h
@@ -35,8 +35,7 @@ class LocalDeclRenamer : public TransformVisitor<LocalDeclRenamer> {
   void RunImpl() override;
 
  public:
-  LocalDeclRenamer(DecompilationContext &dec_ctx, clang::ASTUnit &unit,
-                   IRToNameMap &names);
+  LocalDeclRenamer(DecompilationContext &dec_ctx, IRToNameMap &names);
 
   bool shouldTraversePostOrder() override;
   bool VisitVarDecl(clang::VarDecl *decl);

diff --git a/include/rellic/AST/LoopRefine.h b/include/rellic/AST/LoopRefine.h
@@ -34,7 +34,7 @@ class LoopRefine : public TransformVisitor<LoopRefine> {
   void RunImpl() override;
 
  public:
-  LoopRefine(DecompilationContext &dec_ctx, clang::ASTUnit &unit);
+  LoopRefine(DecompilationContext &dec_ctx);
 
   bool VisitWhileStmt(clang::WhileStmt *loop);
 };

diff --git a/include/rellic/AST/MaterializeConds.h b/include/rellic/AST/MaterializeConds.h
@@ -27,7 +27,7 @@ class MaterializeConds : public TransformVisitor<MaterializeConds> {
   void RunImpl() override;
 
  public:
-  MaterializeConds(DecompilationContext &dec_ctx, clang::ASTUnit &unit);
+  MaterializeConds(DecompilationContext &dec_ctx);
 
   bool VisitIfStmt(clang::IfStmt *stmt);
   bool VisitWhileStmt(clang::WhileStmt *loop);

diff --git a/include/rellic/AST/NestedCondProp.h b/include/rellic/AST/NestedCondProp.h
@@ -54,7 +54,7 @@ class NestedCondProp : public ASTPass {
   void RunImpl() override;
 
  public:
-  NestedCondProp(DecompilationContext& dec_ctx, clang::ASTUnit& unit);
+  NestedCondProp(DecompilationContext& dec_ctx);
 };
 
 }  // namespace rellic
diff --git a/include/rellic/AST/NestedScopeCombine.h b/include/rellic/AST/NestedScopeCombine.h
@@ -35,7 +35,7 @@ class NestedScopeCombine : public TransformVisitor<NestedScopeCombine> {
   void RunImpl() override;
 
  public:
-  NestedScopeCombine(DecompilationContext &dec_ctx, clang::ASTUnit &unit);
+  NestedScopeCombine(DecompilationContext &dec_ctx);
 
   bool VisitIfStmt(clang::IfStmt *ifstmt);
   bool VisitWhileStmt(clang::WhileStmt *stmt);

diff --git a/include/rellic/AST/ReachBasedRefine.h b/include/rellic/AST/ReachBasedRefine.h
@@ -41,7 +41,7 @@ class ReachBasedRefine : public TransformVisitor<ReachBasedRefine> {
   void RunImpl() override;
 
  public:
-  ReachBasedRefine(DecompilationContext &dec_ctx, clang::ASTUnit &unit);
+  ReachBasedRefine(DecompilationContext &dec_ctx);
 
   bool VisitCompoundStmt(clang::CompoundStmt *compound);
 };

diff --git a/include/rellic/AST/StructFieldRenamer.h b/include/rellic/AST/StructFieldRenamer.h
@@ -29,8 +29,7 @@ class StructFieldRenamer
   void RunImpl() override;
 
  public:
-  StructFieldRenamer(DecompilationContext &dec_ctx, clang::ASTUnit &unit,
-                     IRTypeToDITypeMap &types);
+  StructFieldRenamer(DecompilationContext &dec_ctx, IRTypeToDITypeMap &types);
 
   bool VisitRecordDecl(clang::RecordDecl *decl);
 };

diff --git a/include/rellic/AST/TransformVisitor.h b/include/rellic/AST/TransformVisitor.h
@@ -55,8 +55,7 @@ class TransformVisitor : public ASTPass,
   void RunImpl() override { substitutions.clear(); }
 
  public:
-  TransformVisitor(DecompilationContext &dec_ctx, clang::ASTUnit &unit)
-      : ASTPass(dec_ctx, unit) {}
+  TransformVisitor(DecompilationContext &dec_ctx) : ASTPass(dec_ctx) {}
 
   virtual bool shouldTraversePostOrder() { return true; }
 

diff --git a/include/rellic/AST/Util.h b/include/rellic/AST/Util.h
@@ -8,6 +8,7 @@
 
 #pragma once
 
+#include <clang/AST/ASTContext.h>
 #include <clang/AST/DeclBase.h>
 #include <clang/AST/Stmt.h>
 #include <clang/Frontend/ASTUnit.h>
@@ -66,6 +67,12 @@ struct DecompilationContext {
   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;

diff --git a/include/rellic/AST/Z3CondSimplify.h b/include/rellic/AST/Z3CondSimplify.h
@@ -24,7 +24,7 @@ class Z3CondSimplify : public ASTPass {
   void RunImpl() override;
 
  public:
-  Z3CondSimplify(DecompilationContext& dec_ctx, clang::ASTUnit& unit);
+  Z3CondSimplify(DecompilationContext& dec_ctx);
 };
 
 }  // namespace rellic
diff --git a/lib/AST/CondBasedRefine.cpp b/lib/AST/CondBasedRefine.cpp
@@ -15,9 +15,8 @@
 
 namespace rellic {
 
-CondBasedRefine::CondBasedRefine(DecompilationContext &dec_ctx,
-                                 clang::ASTUnit &unit)
-    : TransformVisitor<CondBasedRefine>(dec_ctx, unit) {}
+CondBasedRefine::CondBasedRefine(DecompilationContext &dec_ctx)
+    : TransformVisitor<CondBasedRefine>(dec_ctx) {}
 
 bool CondBasedRefine::VisitCompoundStmt(clang::CompoundStmt *compound) {
   std::vector<clang::Stmt *> body{compound->body_begin(), compound->body_end()};
@@ -52,9 +51,9 @@ bool CondBasedRefine::VisitCompoundStmt(clang::CompoundStmt *compound) {
       // becomes
       // if(a) { X1; X2; } else { Y1; Y2; }
       new_then_body.push_back(then_b);
-      auto new_then{ast.CreateCompoundStmt(new_then_body)};
+      auto new_then{dec_ctx.ast.CreateCompoundStmt(new_then_body)};
 
-      new_if = ast.CreateIf(dec_ctx.marker_expr, new_then);
+      new_if = dec_ctx.ast.CreateIf(dec_ctx.marker_expr, new_then);
 
       if (else_a || else_b) {
         // At least one of the two `if` statements has an `else` branch
@@ -68,7 +67,7 @@ bool CondBasedRefine::VisitCompoundStmt(clang::CompoundStmt *compound) {
           new_else_body.push_back(else_b);
         }
 
-        auto new_else{ast.CreateCompoundStmt(new_else_body)};
+        auto new_else{dec_ctx.ast.CreateCompoundStmt(new_else_body)};
         new_if->setElse(new_else);
       }
 
@@ -86,17 +85,17 @@ bool CondBasedRefine::VisitCompoundStmt(clang::CompoundStmt *compound) {
         new_then_body.push_back(else_b);
       }
 
-      auto new_then{ast.CreateCompoundStmt(new_then_body)};
+      auto new_then{dec_ctx.ast.CreateCompoundStmt(new_then_body)};
 
       std::vector<clang::Stmt *> new_else_body{};
       if (else_a) {
         new_else_body.push_back(else_a);
       }
       new_else_body.push_back(then_b);
 
-      new_if = ast.CreateIf(dec_ctx.marker_expr, new_then);
+      new_if = dec_ctx.ast.CreateIf(dec_ctx.marker_expr, new_then);
 
-      auto new_else{ast.CreateCompoundStmt(new_else_body)};
+      auto new_else{dec_ctx.ast.CreateCompoundStmt(new_else_body)};
       new_if->setElse(new_else);
 
       did_something = true;
@@ -109,15 +108,15 @@ bool CondBasedRefine::VisitCompoundStmt(clang::CompoundStmt *compound) {
     }
   }
   if (did_something) {
-    substitutions[compound] = ast.CreateCompoundStmt(body);
+    substitutions[compound] = dec_ctx.ast.CreateCompoundStmt(body);
   }
   return !Stopped();
 }
 
 void CondBasedRefine::RunImpl() {
   LOG(INFO) << "Condition-based refinement";
   TransformVisitor<CondBasedRefine>::RunImpl();
-  TraverseDecl(ast_ctx.getTranslationUnitDecl());
+  TraverseDecl(dec_ctx.ast_ctx.getTranslationUnitDecl());
 }
 
 }  // namespace rellic
diff --git a/lib/AST/DeadStmtElim.cpp b/lib/AST/DeadStmtElim.cpp
@@ -12,8 +12,8 @@
 
 namespace rellic {
 
-DeadStmtElim::DeadStmtElim(DecompilationContext &dec_ctx, clang::ASTUnit &unit)
-    : TransformVisitor<DeadStmtElim>(dec_ctx, unit) {}
+DeadStmtElim::DeadStmtElim(DecompilationContext &dec_ctx)
+    : TransformVisitor<DeadStmtElim>(dec_ctx) {}
 
 bool DeadStmtElim::VisitIfStmt(clang::IfStmt *ifstmt) {
   // DLOG(INFO) << "VisitIfStmt";
@@ -40,7 +40,7 @@ bool DeadStmtElim::VisitCompoundStmt(clang::CompoundStmt *compound) {
     }
     // Add only necessary statements
     if (auto expr = clang::dyn_cast<clang::Expr>(stmt)) {
-      if (expr->HasSideEffects(ast_ctx)) {
+      if (expr->HasSideEffects(dec_ctx.ast_ctx)) {
         new_body.push_back(stmt);
       }
     } else if (!clang::dyn_cast<clang::NullStmt>(stmt)) {
@@ -49,15 +49,15 @@ bool DeadStmtElim::VisitCompoundStmt(clang::CompoundStmt *compound) {
   }
   // Create the a new compound
   if (changed || new_body.size() < compound->size()) {
-    substitutions[compound] = ast.CreateCompoundStmt(new_body);
+    substitutions[compound] = dec_ctx.ast.CreateCompoundStmt(new_body);
   }
   return !Stopped();
 }
 
 void DeadStmtElim::RunImpl() {
   LOG(INFO) << "Eliminating dead statements";
   TransformVisitor<DeadStmtElim>::RunImpl();
-  TraverseDecl(ast_ctx.getTranslationUnitDecl());
+  TraverseDecl(dec_ctx.ast_ctx.getTranslationUnitDecl());
 }
 
 }  // namespace rellic