4 files changed, 82 insertions, 0 deletions
diff --git a/lib/Module/CMakeLists.txt b/lib/Module/CMakeLists.txt
index 73f003c2..341af9df 100644
--- a/lib/Module/CMakeLists.txt
+++ b/lib/Module/CMakeLists.txt
@@ -16,6 +16,7 @@ set(KLEE_MODULE_COMPONENT_SRCS
   LowerSwitch.cpp
   ModuleUtil.cpp
   Optimize.cpp
+  OptNone.cpp
   PhiCleaner.cpp
   RaiseAsm.cpp
   Scalarizer.cpp
diff --git a/lib/Module/KModule.cpp b/lib/Module/KModule.cpp
index 3ce95a9e..2a15b02f 100644
--- a/lib/Module/KModule.cpp
+++ b/lib/Module/KModule.cpp
@@ -112,6 +112,12 @@ namespace {
   DontVerify("disable-verify",
              cl::desc("Do not verify the module integrity (default=false)"),
              cl::init(false), cl::cat(klee::ModuleCat));
+
+  cl::opt<bool>
+  OptimiseKLEECall("klee-call-optimisation",
+                             cl::desc("Allow optimization of functions that "
+                                      "contain KLEE calls (default=true)"),
+                             cl::init(true), cl::cat(ModuleCat));
 }
 
 /***/
@@ -262,6 +268,14 @@ void KModule::instrument(const Interpreter::ModuleOptions &opts) {
 void KModule::optimiseAndPrepare(
     const Interpreter::ModuleOptions &opts,
     llvm::ArrayRef<const char *> preservedFunctions) {
+  // Preserve all functions containing klee-related function calls from being
+  // optimised around
+  if (!OptimiseKLEECall) {
+    LegacyLLVMPassManagerTy pm;
+    pm.add(new OptNonePass());
+    pm.run(*module);
+  }
+
   if (opts.Optimize)
     Optimize(module.get(), preservedFunctions);
 
diff --git a/lib/Module/OptNone.cpp b/lib/Module/OptNone.cpp
new file mode 100644
index 00000000..08837488
--- /dev/null
+++ b/lib/Module/OptNone.cpp
@@ -0,0 +1,59 @@
+//===-- OptNone.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 "klee/Config/Version.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Module.h"
+
+namespace klee {
+
+char OptNonePass::ID;
+
+bool OptNonePass::runOnModule(llvm::Module &M) {
+  // Find list of functions that start with `klee_`
+  // and mark all functions that contain such call or invoke as optnone
+  llvm::SmallPtrSet<llvm::Function *,16> CallingFunctions;
+  for (auto &F : M) {
+    if (!F.hasName() || !F.getName().startswith("klee_"))
+      continue;
+#if LLVM_VERSION_CODE >= LLVM_VERSION(3, 5)
+    for (auto *U : F.users()) {
+      // skip non-calls and non-invokes
+      if (!llvm::isa<llvm::CallInst>(U) && !llvm::isa<llvm::InvokeInst>(U))
+        continue;
+      auto *Inst = llvm::cast<llvm::Instruction>(U);
+      CallingFunctions.insert(Inst->getParent()->getParent());
+    }
+#else
+    for (auto i = F.use_begin(), e = F.use_end(); i != e; ++i) {
+      if (auto Inst = llvm::dyn_cast<llvm::Instruction>(*i)) {
+        CallingFunctions.insert(Inst->getParent()->getParent());
+      }
+    }
+#endif
+  }
+
+  bool changed = false;
+  for (auto F : CallingFunctions) {
+    // Skip if already annotated
+    if (F->hasFnAttribute(llvm::Attribute::OptimizeNone))
+      continue;
+    F->addFnAttr(llvm::Attribute::OptimizeNone);
+    F->addFnAttr(llvm::Attribute::NoInline);
+    changed = true;
+  }
+
+  return changed;
+}
+} // namespace klee
diff --git a/lib/Module/Passes.h b/lib/Module/Passes.h
index 417ebe7d..6a302819 100644
--- a/lib/Module/Passes.h
+++ b/lib/Module/Passes.h
@@ -185,6 +185,14 @@ public:
   bool runOnModule(llvm::Module &M) override;
 };
 #endif
+
+/// Instruments every function that contains a KLEE function call as nonopt
+class OptNonePass : public llvm::ModulePass {
+public:
+  static char ID;
+  OptNonePass() : llvm::ModulePass(ID) {}
+  bool runOnModule(llvm::Module &M) override;
+};
 } // namespace klee
 
 #endif