2 files changed, 4 insertions, 5 deletions
diff --git a/lib/Core/StatsTracker.cpp b/lib/Core/StatsTracker.cpp
index 08df3127..1c963bfc 100644
--- a/lib/Core/StatsTracker.cpp
+++ b/lib/Core/StatsTracker.cpp
@@ -615,7 +615,7 @@ void StatsTracker::computeReachableUncovered() {
               // (which should be correct anyhow).
               callTargets.insert(std::make_pair(it,
                                                 std::vector<Function*>()));
-            } else if (Function *target = getDirectCallTarget(it)) {
+            } else if (Function *target = getDirectCallTarget(cs)) {
               callTargets[it].push_back(target);
             } else {
               callTargets[it] = 
diff --git a/lib/Module/ModuleUtil.cpp b/lib/Module/ModuleUtil.cpp
index df495da8..17c4aa91 100644
--- a/lib/Module/ModuleUtil.cpp
+++ b/lib/Module/ModuleUtil.cpp
@@ -17,6 +17,7 @@
 #include "llvm/Module.h"
 #include "llvm/Assembly/AsmAnnotationWriter.h"
 #include "llvm/Support/CFG.h"
+#include "llvm/Support/CallSite.h"
 #include "llvm/Support/InstIterator.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Analysis/ValueTracking.h"
@@ -45,10 +46,8 @@ Module *klee::linkWithLibrary(Module *module,
   return linker.releaseModule();
 }
 
-Function *klee::getDirectCallTarget(const Instruction *i) {
-  assert(isa<CallInst>(i) || isa<InvokeInst>(i));
-
-  Value *v = i->getOperand(0);
+Function *klee::getDirectCallTarget(CallSite cs) {
+  Value *v = cs.getCalledValue();
   if (Function *f = dyn_cast<Function>(v)) {
     return f;
   } else if (llvm::ConstantExpr *ce = dyn_cast<llvm::ConstantExpr>(v)) {