From 65691196c5fddcf22f26d8f24a40d40fda79a3f5 Mon Sep 17 00:00:00 2001
From: Yurii Kostyukov <kostyukov.yurii@gmail.com>
Date: Wed, 25 Oct 2023 11:10:49 +0300
Subject: [PATCH] [feat] Removed unwanted calls pass

---
 lib/Module/CMakeLists.txt   |  1 +
 lib/Module/CallRemover.cpp  | 36 ++++++++++++++++++++++++++++++++++++
 lib/Module/CallSplitter.cpp |  8 ++++----
 lib/Module/KModule.cpp      |  7 +++++++
 lib/Module/Optimize.cpp     |  7 -------
 lib/Module/Passes.h         |  8 ++++++++
 6 files changed, 56 insertions(+), 11 deletions(-)
 create mode 100644 lib/Module/CallRemover.cpp

diff --git a/lib/Module/CMakeLists.txt b/lib/Module/CMakeLists.txt
index 42ef35556b..bca27a10ea 100644
--- a/lib/Module/CMakeLists.txt
+++ b/lib/Module/CMakeLists.txt
@@ -8,6 +8,7 @@
 #===------------------------------------------------------------------------===#
 set(KLEE_MODULE_COMPONENT_SRCS
   CallSplitter.cpp
+  CallRemover.cpp
   Checks.cpp
   CodeGraphInfo.cpp
   FunctionAlias.cpp
diff --git a/lib/Module/CallRemover.cpp b/lib/Module/CallRemover.cpp
new file mode 100644
index 0000000000..d2854ebd47
--- /dev/null
+++ b/lib/Module/CallRemover.cpp
@@ -0,0 +1,36 @@
+//===-- CallRemover.cpp----------------------------------------------------===//
+//
+//                     The KLEE Symbolic Virtual Machine
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Passes.h"
+#include <unordered_set>
+
+namespace klee {
+
+using namespace llvm;
+
+char CallRemover::ID;
+
+bool CallRemover::runOnModule(llvm::Module &M) {
+  std::vector<std::string> badFuncs = {"llvm.dbg.declare", "llvm.dbg.label"};
+
+  for (const auto &f : badFuncs) {
+    auto Declare = M.getFunction(f);
+    if (!Declare)
+      continue;
+    while (!Declare->use_empty()) {
+      auto CI = cast<CallInst>(Declare->user_back());
+      assert(CI->use_empty() && "deleted function must have void result");
+      CI->eraseFromParent();
+    }
+    Declare->eraseFromParent();
+  }
+
+  return true;
+}
+} // namespace klee
diff --git a/lib/Module/CallSplitter.cpp b/lib/Module/CallSplitter.cpp
index 4c4a71c4f3..a90e17516d 100644
--- a/lib/Module/CallSplitter.cpp
+++ b/lib/Module/CallSplitter.cpp
@@ -1,5 +1,4 @@
-//===-- CallSplitter.cpp
-//-------------------------------------------------------===//
+//===-- CallSplitter.cpp --------------------------------------------------===//
 //
 //                     The KLEE Symbolic Virtual Machine
 //
@@ -46,8 +45,9 @@ bool CallSplitter::runOnFunction(Function &F) {
         if (callInst != firstInst) {
           fbb = fbb->splitBasicBlock(callInst);
         }
-        if (isa<BranchInst>(afterCallInst) &&
-            cast<BranchInst>(afterCallInst)->isUnconditional()) {
+        if ((isa<BranchInst>(afterCallInst) &&
+             cast<BranchInst>(afterCallInst)->isUnconditional()) ||
+            isa<UnreachableInst>(afterCallInst)) {
           break;
         }
         fbb = fbb->splitBasicBlock(afterCallInst);
diff --git a/lib/Module/KModule.cpp b/lib/Module/KModule.cpp
index 343974b472..ce09edfcc2 100644
--- a/lib/Module/KModule.cpp
+++ b/lib/Module/KModule.cpp
@@ -117,6 +117,11 @@ cl::opt<bool>
                cl::desc("Split each call in own basic block (default=true)"),
                cl::init(true), cl::cat(klee::ModuleCat));
 
+static cl::opt<bool>
+    StripUnwantedCalls("strip-unwanted-calls",
+                       cl::desc("Strip all unwanted calls (llvm.dbg.* stuff)"),
+                       cl::init(false), cl::cat(klee::ModuleCat));
+
 cl::opt<bool> SplitReturns(
     "split-returns",
     cl::desc("Split each return in own basic block (default=true)"),
@@ -332,6 +337,8 @@ void KModule::optimiseAndPrepare(
   pm3.add(createScalarizerPass());
   pm3.add(new PhiCleanerPass());
   pm3.add(new FunctionAliasPass());
+  if (StripUnwantedCalls)
+    pm3.add(new CallRemover());
   if (SplitCalls) {
     pm3.add(new CallSplitter());
   }
diff --git a/lib/Module/Optimize.cpp b/lib/Module/Optimize.cpp
index 77e11b2e2d..cda430bd96 100644
--- a/lib/Module/Optimize.cpp
+++ b/lib/Module/Optimize.cpp
@@ -72,11 +72,6 @@ static cl::opt<bool>
                cl::desc("Strip debugger symbol info from executable"),
                cl::init(false), cl::cat(klee::ModuleCat));
 
-static cl::opt<bool>
-    StripDebugDeclare("strip-debug-declare",
-                      cl::desc("Strip all llvm.dbg.declare intrinsics"),
-                      cl::init(true), cl::cat(klee::ModuleCat));
-
 static cl::alias A1("S", cl::desc("Alias for --strip-debug"),
                     cl::aliasopt(StripDebug));
 
@@ -103,8 +98,6 @@ static void AddStandardCompilePasses(legacy::PassManager &PM) {
   // If the -strip-debug command line option was specified, do it.
   if (StripDebug)
     addPass(PM, createStripSymbolsPass(true));
-  if (StripDebugDeclare)
-    addPass(PM, createStripDebugDeclarePass());
 
   addPass(PM, createCFGSimplificationPass());       // Clean up disgusting code
   addPass(PM, createPromoteMemoryToRegisterPass()); // Kill useless allocas
diff --git a/lib/Module/Passes.h b/lib/Module/Passes.h
index 76499428df..01983e0b85 100644
--- a/lib/Module/Passes.h
+++ b/lib/Module/Passes.h
@@ -206,6 +206,14 @@ class CallSplitter : public llvm::FunctionPass {
   bool runOnFunction(llvm::Function &F) override;
 };
 
+/// Remove unwanted calls
+class CallRemover : public llvm::ModulePass {
+public:
+  static char ID;
+  CallRemover() : llvm::ModulePass(ID) {}
+  bool runOnModule(llvm::Module &M) override;
+};
+
 class ReturnSplitter : public llvm::FunctionPass {
 public:
   static char ID;