From e125f1d28339791ebae6ea380841e465a0a20661 Mon Sep 17 00:00:00 2001
From: Aleksandr Misonizhnik <misonijnik@gmail.com>
Date: Fri, 11 Nov 2022 22:12:39 +0300
Subject: [PATCH] [wip] Add kmoduleOrig

---
 include/klee/Core/Interpreter.h |  3 ++-
 lib/Core/Executor.cpp           |  9 ++++++++-
 lib/Core/Executor.h             |  4 +++-
 lib/Runner/run_klee.cpp         | 15 ++++++++++++++-
 4 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/include/klee/Core/Interpreter.h b/include/klee/Core/Interpreter.h
index 6f8bce6fbf..7ef30fd491 100644
--- a/include/klee/Core/Interpreter.h
+++ b/include/klee/Core/Interpreter.h
@@ -121,7 +121,8 @@ class Interpreter {
   /// \return The final module after it has been optimized, checks
   /// inserted, and modified for interpretation.
   virtual llvm::Module *
-  setModule(std::vector<std::unique_ptr<llvm::Module>> &modules,
+  setModule(std::unique_ptr<llvm::Module> mainModule,
+            std::vector<std::unique_ptr<llvm::Module>> &modules,
             const ModuleOptions &opts,
             const std::vector<std::string> &mainFunctions) = 0;
 
diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
index 5170ad1f6b..6c653a4136 100644
--- a/lib/Core/Executor.cpp
+++ b/lib/Core/Executor.cpp
@@ -593,13 +593,15 @@ Executor::Executor(LLVMContext &ctx, const InterpreterOptions &opts,
 }
 
 llvm::Module *
-Executor::setModule(std::vector<std::unique_ptr<llvm::Module>> &modules,
+Executor::setModule(std::unique_ptr<llvm::Module> mainModule,
+                    std::vector<std::unique_ptr<llvm::Module>> &modules,
                     const ModuleOptions &opts,
                     const std::vector<std::string> &mainFunctions) {
   assert(!kmodule && !modules.empty() &&
          "can only register one module"); // XXX gross
 
   kmodule = std::unique_ptr<KModule>(new KModule());
+  kmoduleOrig = std::unique_ptr<KModule>(new KModule());
 
   // Preparing the final module happens in multiple stages
 
@@ -613,6 +615,11 @@ Executor::setModule(std::vector<std::unique_ptr<llvm::Module>> &modules,
     klee_error("Could not load KLEE intrinsic file %s", LibPath.c_str());
   }
 
+  std::vector<std::unique_ptr<llvm::Module>> mainModuleVector;
+  mainModuleVector.push_back(std::move(mainModule));
+  kmoduleOrig->link(mainModuleVector, opts.EntryPoint);
+  kmoduleOrig->checkModule();
+
   // 1.) Link the modules together
   while (kmodule->link(modules, opts.EntryPoint)) {
     // 2.) Apply different instrumentation
diff --git a/lib/Core/Executor.h b/lib/Core/Executor.h
index 185a03dbd9..d3096d1fc5 100644
--- a/lib/Core/Executor.h
+++ b/lib/Core/Executor.h
@@ -123,6 +123,7 @@ class Executor : public Interpreter {
   static const std::unordered_set <llvm::Intrinsic::ID> modelledFPIntrinsics;
 
   std::unique_ptr<KModule> kmodule;
+  std::unique_ptr<KModule> kmoduleOrig;
   InterpreterHandler *interpreterHandler;
   Searcher *searcher;
 
@@ -569,7 +570,8 @@ class Executor : public Interpreter {
   }
 
   llvm::Module *
-  setModule(std::vector<std::unique_ptr<llvm::Module>> &modules,
+  setModule(std::unique_ptr<llvm::Module> mainModule,
+            std::vector<std::unique_ptr<llvm::Module>> &modules,
             const ModuleOptions &opts,
             const std::vector<std::string> &mainFunctions) override;
 
diff --git a/lib/Runner/run_klee.cpp b/lib/Runner/run_klee.cpp
index f8040a193f..b546d40d38 100644
--- a/lib/Runner/run_klee.cpp
+++ b/lib/Runner/run_klee.cpp
@@ -1846,11 +1846,17 @@ int run_klee(int argc, char **argv, char **envp) {
     // Load the bytecode...
     std::string errorMsg;
     LLVMContext ctx;
+    LLVMContext originalCtx;
     std::vector<std::unique_ptr<llvm::Module>> loadedModules;
+    std::vector<std::unique_ptr<llvm::Module>> originalLoadedModules;
     if (!klee::loadFile(InputFile, ctx, loadedModules, errorMsg)) {
         klee_error("error loading program '%s': %s", InputFile.c_str(),
                    errorMsg.c_str());
     }
+    if (!klee::loadFile(InputFile, originalCtx, originalLoadedModules, errorMsg)) {
+        klee_error("error loading program '%s': %s", InputFile.c_str(),
+                   errorMsg.c_str());
+    }
     // Load and link the whole files content. The assumption is that this is the
     // application under test.
     // Nothing gets removed in the first place.
@@ -1861,6 +1867,13 @@ int run_klee(int argc, char **argv, char **envp) {
                    errorMsg.c_str());
     }
 
+    std::unique_ptr<llvm::Module> OM(klee::linkModules(
+            originalLoadedModules, "" /* link all modules together */, errorMsg));
+    if (!OM) {
+        klee_error("error loading program '%s': %s", InputFile.c_str(),
+                   errorMsg.c_str());
+    }
+
     llvm::Module *mainModule = M.get();
 
     if (ExecutionMode == Interpreter::GuidanceKind::ErrorGuidance) {
@@ -2080,7 +2093,7 @@ int run_klee(int argc, char **argv, char **envp) {
     // Get the desired main function.  klee_main initializes uClibc
     // locale and other data and then calls main.
 
-    auto finalModule = interpreter->setModule(loadedModules, Opts, mainFunctions);
+    auto finalModule = interpreter->setModule(std::move(OM), loadedModules, Opts, mainFunctions);
 
     if (InteractiveMode) {
       klee_message("KLEE finish preprocessing.");