1 files changed, 38 insertions, 42 deletions

diff --git a/lib/Module/Checks.cpp b/lib/Module/Checks.cpp
index aac63e1d..a9ef0dd4 100644
--- a/lib/Module/Checks.cpp
+++ b/lib/Module/Checks.cpp
@@ -12,16 +12,17 @@
 #include "klee/Config/Version.h"
 
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
-#include "llvm/IR/Module.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
-#include "llvm/IR/DataLayout.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
@@ -31,50 +32,45 @@ using namespace klee;
 
 char DivCheckPass::ID;
 
-bool DivCheckPass::runOnModule(Module &M) { 
-  Function *divZeroCheckFunction = 0;
-  LLVMContext &ctx = M.getContext();
-
-  bool moduleChanged = false;
-  
-  for (Module::iterator f = M.begin(), fe = M.end(); f != fe; ++f) {
-    for (Function::iterator b = f->begin(), be = f->end(); b != be; ++b) {
-      for (BasicBlock::iterator i = b->begin(), ie = b->end(); i != ie; ++i) {     
-          if (BinaryOperator* binOp = dyn_cast<BinaryOperator>(i)) {
-          // find all [s|u][div|mod] instructions
-          Instruction::BinaryOps opcode = binOp->getOpcode();
-          if (opcode == Instruction::SDiv || opcode == Instruction::UDiv ||
-              opcode == Instruction::SRem || opcode == Instruction::URem) {
-            
-            CastInst *denominator =
-              CastInst::CreateIntegerCast(i->getOperand(1),
-                                          Type::getInt64Ty(ctx),
-                                          false,  /* sign doesn't matter */
-                                          "int_cast_to_i64",
-                                          &*i);
-            
-            // Lazily bind the function to avoid always importing it.
-            if (!divZeroCheckFunction) {
-              Constant *fc = M.getOrInsertFunction("klee_div_zero_check", 
-                                                   Type::getVoidTy(ctx),
-                                                   Type::getInt64Ty(ctx),
-                                                   NULL);
-              divZeroCheckFunction = cast<Function>(fc);
-            }
-
-            CallInst * ci = CallInst::Create(divZeroCheckFunction, denominator, "", &*i);
-
-            // Set debug location of checking call to that of the div/rem
-            // operation so error locations are reported in the correct
-            // location.
-            ci->setDebugLoc(binOp->getDebugLoc());
-            moduleChanged = true;
-          }
+bool DivCheckPass::runOnModule(Module &M) {
+  std::vector<llvm::BinaryOperator *> divInstruction;
+
+  for (auto &F : M) {
+    for (auto &BB : F) {
+      for (auto &I : BB) {
+        auto binOp = dyn_cast<BinaryOperator>(&I);
+        if (!binOp)
+          continue;
+
+        // find all [s|u][div|rem] instructions
+        auto opcode = binOp->getOpcode();
+        if (opcode != Instruction::SDiv && opcode != Instruction::UDiv &&
+            opcode != Instruction::SRem && opcode != Instruction::URem)
+          continue;
         }
       }
     }
   }
-  return moduleChanged;
+
+  // If nothing to do, return
+  if (divInstruction.empty())
+    return false;
+
+  LLVMContext &ctx = M.getContext();
+  auto divZeroCheckFunction = cast<Function>(
+      M.getOrInsertFunction("klee_div_zero_check", Type::getVoidTy(ctx),
+                            Type::getInt64Ty(ctx), NULL));
+
+  for (auto &divInst : divInstruction) {
+    llvm::IRBuilder<> Builder(divInst /* Inserts before divInst*/);
+    auto denominator =
+        Builder.CreateIntCast(divInst->getOperand(1), Type::getInt64Ty(ctx),
+                              false, /* sign doesn't matter */
+                              "int_cast_to_i64");
+    Builder.CreateCall(divZeroCheckFunction, denominator);
+  }
+
+  return true;
 }
 
 char OvershiftCheckPass::ID;