2 files changed, 40 insertions, 42 deletions

diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
index bef34059..caf0de36 100644
--- a/lib/Core/Executor.cpp
+++ b/lib/Core/Executor.cpp
@@ -572,20 +572,31 @@ void Executor::branch(ExecutionState &state,
   unsigned N = conditions.size();
   assert(N);
 
-  stats::forks += N-1;
-
-  // XXX do proper balance or keep random?
-  result.push_back(&state);
-  for (unsigned i=1; i<N; ++i) {
-    ExecutionState *es = result[theRNG.getInt32() % i];
-    ExecutionState *ns = es->branch();
-    addedStates.insert(ns);
-    result.push_back(ns);
-    es->ptreeNode->data = 0;
-    std::pair<PTree::Node*,PTree::Node*> res = 
-      processTree->split(es->ptreeNode, ns, es);
-    ns->ptreeNode = res.first;
-    es->ptreeNode = res.second;
+  if (MaxForks!=~0u && stats::forks >= MaxForks) {
+    unsigned next = theRNG.getInt32() % N;
+    for (unsigned i=0; i<N; ++i) {
+      if (i == next) {
+        result.push_back(&state);
+      } else {
+        result.push_back(NULL);
+      }
+    }
+  } else {
+    stats::forks += N-1;
+
+    // XXX do proper balance or keep random?
+    result.push_back(&state);
+    for (unsigned i=1; i<N; ++i) {
+      ExecutionState *es = result[theRNG.getInt32() % i];
+      ExecutionState *ns = es->branch();
+      addedStates.insert(ns);
+      result.push_back(ns);
+      es->ptreeNode->data = 0;
+      std::pair<PTree::Node*,PTree::Node*> res = 
+        processTree->split(es->ptreeNode, ns, es);
+      ns->ptreeNode = res.first;
+      es->ptreeNode = res.second;
+    }
   }
 
   // If necessary redistribute seeds to match conditions, killing
@@ -620,12 +631,14 @@ void Executor::branch(ExecutionState &state,
       if (i==N)
         i = theRNG.getInt32() % N;
 
-      seedMap[result[i]].push_back(*siit);
+      // Extra check in case we're replaying seeds with a max-fork
+      if (result[i])
+        seedMap[result[i]].push_back(*siit);
     }
 
     if (OnlyReplaySeeds) {
       for (unsigned i=0; i<N; ++i) {
-        if (!seedMap.count(result[i])) {
+        if (result[i] && !seedMap.count(result[i])) {
           terminateState(*result[i]);
           result[i] = NULL;
         }
diff --git a/lib/Module/ModuleUtil.cpp b/lib/Module/ModuleUtil.cpp
index 617a91ca..43eef0b7 100644
--- a/lib/Module/ModuleUtil.cpp
+++ b/lib/Module/ModuleUtil.cpp
@@ -15,6 +15,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IRReader/IRReader.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/DataStream.h"
@@ -54,34 +55,18 @@ using namespace klee;
 Module *klee::linkWithLibrary(Module *module, 
                               const std::string &libraryName) {
 #if LLVM_VERSION_CODE >= LLVM_VERSION(3, 3)
-  Linker linker(module);
-  std::string errorMessage;
-
-  DataStreamer * streamer = getDataFileStreamer(libraryName, &errorMessage);
-
-  if (!streamer)
-    fprintf(stderr, "Error Loading file: %s\n", errorMessage.c_str());
-  assert(streamer);
-  
-  OwningPtr<Module> library_module;
-  library_module.reset(getStreamedBitcodeModule(libraryName, streamer, getGlobalContext(), &errorMessage));
-  if (library_module.get() != 0
-	  && library_module->MaterializeAllPermanently(&errorMessage)) {
-	  library_module.reset();
+  SMDiagnostic err;
+  std::string err_str;
+  sys::Path libraryPath(libraryName);
+  Module *new_mod = ParseIRFile(libraryPath.str(), err, 
+module->getContext());
+
+  if (Linker::LinkModules(module, new_mod, Linker::DestroySource, 
+&err_str)) {
+    assert(0 && "linked in library failed!");
   }
 
-  if (library_module.get() == 0) {
-	  errs() << errorMessage << " for " << libraryName << "\n";
-	  assert(library_module.get());
-  }
-  if (linker.linkInModule(library_module.get(), &errorMessage)){
-	  fprintf(stderr, "Error in Linking %s; Existing module: %s, library to be linked in %s\n", errorMessage.c_str(),
-	      module->getModuleIdentifier().c_str(), libraryName.c_str());
-	  assert(0 && "linking in library failed!");
-  }
-
-  return linker.getModule();
-
+  return module;
 #else
   Linker linker("klee", module, false);