1 files changed, 44 insertions, 43 deletions

diff --git a/lib/Module/KModule.cpp b/lib/Module/KModule.cpp
index d889b51f..697b6ea9 100644
--- a/lib/Module/KModule.cpp
+++ b/lib/Module/KModule.cpp
@@ -182,7 +182,8 @@ static void injectStaticConstructorsAndDestructors(Module *m) {
   
   if (ctors || dtors) {
     Function *mainFn = m->getFunction("main");
-    assert(mainFn && "unable to find main function");
+    if (!mainFn)
+      klee_error("Could not find main() function.");
 
     if (ctors)
     CallInst::Create(getStubFunctionForCtorList(m, ctors, "klee.ctor_stub"),
@@ -229,7 +230,7 @@ static void forceImport(Module *m, const char *name, LLVM_TYPE_Q Type *retType,
 /// It is intended that this function be used for inling calls to
 /// check functions like <tt>klee_div_zero_check()</tt>
 static void inlineChecks(Module *module, const char * functionName) {
-  std::vector<CallInst*> checkCalls;
+  std::vector<CallSite> checkCalls;
     Function* runtimeCheckCall = module->getFunction(functionName);
     if (runtimeCheckCall == 0)
     {
@@ -237,22 +238,19 @@ static void inlineChecks(Module *module, const char * functionName) {
       return;
     }
 
-    // Iterate through instructions in module and collect all
-    // call instructions to "functionName" that we care about.
-    for (Module::iterator f = module->begin(), fe = module->end(); f != fe; ++f) {
-      for (inst_iterator i=inst_begin(f), ie = inst_end(f); i != ie; ++i) {
-        if ( CallInst* ci = dyn_cast<CallInst>(&*i) )
-        {
-          if ( ci->getCalledFunction() == runtimeCheckCall)
-            checkCalls.push_back(ci);
-        }
+    for (Value::use_iterator i = runtimeCheckCall->use_begin(),
+        e = runtimeCheckCall->use_end(); i != e; ++i)
+      if (isa<InvokeInst>(*i) || isa<CallInst>(*i)) {
+        CallSite cs(*i);
+        if (!cs.getCalledFunction())
+          continue;
+        checkCalls.push_back(cs);
       }
-    }
 
     unsigned int successCount=0;
     unsigned int failCount=0;
     InlineFunctionInfo IFI(0,0);
-    for ( std::vector<CallInst*>::iterator ci = checkCalls.begin(),
+    for ( std::vector<CallSite>::iterator ci = checkCalls.begin(),
           cie = checkCalls.end();
           ci != cie; ++ci)
     {
@@ -269,6 +267,17 @@ static void inlineChecks(Module *module, const char * functionName) {
     DEBUG( klee_message("Tried to inline calls to %s. %u successes, %u failures",functionName, successCount, failCount) );
 }
 
+void KModule::addInternalFunction(const char* functionName){
+  Function* internalFunction = module->getFunction(functionName);
+  if (!internalFunction) {
+    DEBUG_WITH_TYPE("KModule", klee_warning(
+        "Failed to add internal function %s. Not found.", functionName));
+    return ;
+  }
+  DEBUG( klee_message("Added function %s.",functionName));
+  internalFunctions.insert(internalFunction);
+}
+
 void KModule::prepare(const Interpreter::ModuleOptions &opts,
                       InterpreterHandler *ih) {
   if (!MergeAtExit.empty()) {
@@ -366,25 +375,23 @@ void KModule::prepare(const Interpreter::ModuleOptions &opts,
   // FIXME: Find a way that we can test programs without requiring
   // this to be linked in, it makes low level debugging much more
   // annoying.
-  llvm::sys::Path path(opts.LibraryDir);
-#if LLVM_VERSION_CODE >= LLVM_VERSION(3, 3)
-  path.appendComponent("kleeRuntimeIntrinsic.bc");
+
+  SmallString<128> LibPath(opts.LibraryDir);
+  llvm::sys::path::append(LibPath,
+#if LLVM_VERSION_CODE >= LLVM_VERSION(3,3)
+      "kleeRuntimeIntrinsic.bc"
 #else
-  path.appendComponent("libkleeRuntimeIntrinsic.bca");
+      "libkleeRuntimeIntrinsic.bca"
 #endif
-  module = linkWithLibrary(module, path.c_str());
-
-  /* In order for KLEE to report ALL errors at instrumented
-   * locations the instrumentation call (e.g. "klee_div_zero_check")
-   * must be inlined. Otherwise one of the instructions in the
-   * instrumentation function will be used as the the location of
-   * the error which means that the error cannot be recorded again
-   * ( unless -emit-all-errors is used).
-   */
+    );
+  module = linkWithLibrary(module, LibPath.str());
+
+  // Add internal functions which are not used to check if instructions
+  // have been already visited
   if (opts.CheckDivZero)
-    inlineChecks(module, "klee_div_zero_check");
+    addInternalFunction("klee_div_zero_check");
   if (opts.CheckOvershift)
-    inlineChecks(module, "klee_overshift_check");
+    addInternalFunction("klee_overshift_check");
 
 
   // Needs to happen after linking (since ctors/dtors can be modified)
@@ -423,15 +430,13 @@ void KModule::prepare(const Interpreter::ModuleOptions &opts,
   // around a kcachegrind parsing bug (it puts them on new lines), so
   // that source browsing works.
   if (OutputSource) {
-    std::ostream *os = ih->openOutputFile("assembly.ll");
-    assert(os && os->good() && "unable to open source output");
-
-    llvm::raw_os_ostream *ros = new llvm::raw_os_ostream(*os);
+    llvm::raw_fd_ostream *os = ih->openOutputFile("assembly.ll");
+    assert(os && !os->has_error() && "unable to open source output");
 
     // We have an option for this in case the user wants a .ll they
     // can compile.
     if (NoTruncateSourceLines) {
-      *ros << *module;
+      *os << *module;
     } else {
       std::string string;
       llvm::raw_string_ostream rss(string);
@@ -442,30 +447,26 @@ void KModule::prepare(const Interpreter::ModuleOptions &opts,
       for (;;) {
         const char *end = index(position, '\n');
         if (!end) {
-          *ros << position;
+          *os << position;
           break;
         } else {
           unsigned count = (end - position) + 1;
           if (count<255) {
-            ros->write(position, count);
+            os->write(position, count);
           } else {
-            ros->write(position, 254);
-            *ros << "\n";
+            os->write(position, 254);
+            *os << "\n";
           }
           position = end+1;
         }
       }
     }
-    delete ros;
-
     delete os;
   }
 
   if (OutputModule) {
-    std::ostream *f = ih->openOutputFile("final.bc");
-    llvm::raw_os_ostream* rfs = new llvm::raw_os_ostream(*f);
-    WriteBitcodeToFile(module, *rfs);
-    delete rfs;
+    llvm::raw_fd_ostream *f = ih->openOutputFile("final.bc");
+    WriteBitcodeToFile(module, *f);
     delete f;
   }