move afl_loop check

1 files changed, 27 insertions, 25 deletions

diff --git a/llvm_mode/afl-llvm-lto-instrumentation.so.cc b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
index 300951fb..10482705 100644
--- a/llvm_mode/afl-llvm-lto-instrumentation.so.cc
+++ b/llvm_mode/afl-llvm-lto-instrumentation.so.cc
@@ -602,17 +602,41 @@ bool AFLLTOPass::runOnModule(Module &M) {
 
     for (auto &BB : F) {
 
-      uint32_t succ = 0;
+      if (F.size() == 1) {
+
+        InsBlocks.push_back(&BB);
+        continue;
 
-      if (F.size() == 1) InsBlocks.push_back(&BB);
+      }
 
+      uint32_t succ = 0;
       for (succ_iterator SI = succ_begin(&BB), SE = succ_end(&BB); SI != SE;
            ++SI)
         if ((*SI)->size() > 0) succ++;
-
       if (succ < 2)  // no need to instrument
         continue;
 
+      if (BlockList.size()) {
+
+        int skip = 0;
+        for (uint32_t k = 0; k < BlockList.size(); k++) {
+
+          if (origBB == BlockList[k]) {
+
+            if (debug)
+              fprintf(stderr,
+                      "DEBUG: Function %s skipping BB with/after __afl_loop\n",
+                      F.getName().str().c_str());
+            skip = 1;
+
+          }
+
+        }
+
+        if (skip) continue;
+
+      }
+
       InsBlocks.push_back(&BB);
 
     }
@@ -631,28 +655,6 @@ bool AFLLTOPass::runOnModule(Module &M) {
         uint32_t                  fs = origBB->getParent()->size();
         uint32_t                  countto;
 
-        if (BlockList.size()) {
-
-          int skip = 0;
-          for (uint32_t k = 0; k < BlockList.size(); k++) {
-
-            if (origBB == BlockList[k]) {
-
-              if (debug)
-                fprintf(
-                    stderr,
-                    "DEBUG: Function %s skipping BB with/after __afl_loop\n",
-                    F.getName().str().c_str());
-              skip = 1;
-
-            }
-
-          }
-
-          if (skip) continue;
-
-        }
-
         for (succ_iterator SI = succ_begin(origBB), SE = succ_end(origBB);
              SI != SE; ++SI) {