Instrument edges instead of basic blocks

We are currently deriving edge coverage instrumentation from basic block
instrumentation via the AFL XOR-technique. This has several downsides:

* Different edges can be assigned the same position in the coverage map,
  which leads to underreported coverage.
* The coverage map needs to be large enough for collisions to be
  unlikely (on the order of num_edges^2). In addition to being wasteful,
  it is also hard to determine the correct size given that we don't know
  the number of edges.

In addition to the design limitations, the current implementation
additionally does not take into account that most Java method
invocations can throw exceptions and thus need to be instrumented.

These issues are resolved by switching to true LLVM-style edge coverage
instrumentation. The new coverage instrumentation is based on a lightly
patched version of the JaCoCo internals.

Note:
//agent/src/test/java/com/code_intelligence/jazzer/instrumentor:coverage_instrumentation_test
is not passing for this commit. It will be fixed with the next commit.

3 files changed, 38 insertions, 57 deletions

diff --git a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/CoverageInstrumentationTarget.java b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/CoverageInstrumentationTarget.java
index b738ab19..5725ba23 100644
--- a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/CoverageInstrumentationTarget.java
+++ b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/CoverageInstrumentationTarget.java
@@ -18,14 +18,11 @@ import java.util.HashMap;
 import java.util.Map;
 
 public class CoverageInstrumentationTarget implements DynamicTestContract {
-  // Constructor loc: constructorStart
-
   volatile int int1 = 3;
   volatile int int2 = 213234;
 
   @Override
   public Map<String, Boolean> selfCheck() {
-    // loc: selfCheckStart
     Map<String, Boolean> results = new HashMap<>();
 
     results.put("for0", false);
@@ -37,44 +34,32 @@ public class CoverageInstrumentationTarget implements DynamicTestContract {
     results.put("baz", true);
 
     if (int1 < int2) {
-      // loc: ifFirstBranch
       results.put("block1", true);
     } else {
-      // loc: not reached
       results.put("block2", false);
     }
-    // loc: ifEnd
 
-    for (int i = 0; /* loc: outerForCondition */ i < 2; /* loc: outerForIncrementCounter */ i++) {
-      /* loc: outerForBody */
-      for (int j = 0; /* loc: innerForCondition */ j < 5; /* loc: innerForIncrementCounter */ j++) {
-        // loc: innerForBody
-        results.put("for" + j,
-            i != 0); // != 0 loc: innerForBodyIfSecondRun, == 0 loc: innerForBodyIfFirstRun
+    for (int i = 0; i < 2; i++) {
+      for (int j = 0; j < 5; j++) {
+        results.put("for" + j, i != 0);
       }
     }
-    // loc: outerForAfter
 
     foo(results);
-    // baz(results);
 
     return results;
   }
 
   private void foo(Map<String, Boolean> results) {
-    // loc: fooStart
     bar(results);
   }
 
   private void bar(Map<String, Boolean> results) {
-    // loc: barStart
     results.put("foobar", true);
   }
 
-  // Not called.
   @SuppressWarnings("unused")
   private void baz(Map<String, Boolean> results) {
-    // loc: not reached
     results.put("baz", false);
   }
 }
diff --git a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/CoverageInstrumentationTest.kt b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/CoverageInstrumentationTest.kt
index 5de70f8a..c689e1ac 100644
--- a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/CoverageInstrumentationTest.kt
+++ b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/CoverageInstrumentationTest.kt
@@ -19,7 +19,9 @@ import java.io.File
 import kotlin.test.assertEquals
 
 private fun applyInstrumentation(bytecode: ByteArray): ByteArray {
-    return AFLCoverageMapInstrumentor(MockCoverageMap::class.java).instrument(bytecode)
+    EdgeCoverageInstrumentor.resetNextGlobalEdgeIdForTestingOnly()
+    EdgeCoverageInstrumentor.setCoverageMapClassForTestingOnly(MockCoverageMap::class.java)
+    return EdgeCoverageInstrumentor.instrument(bytecode)
 }
 
 private fun getOriginalInstrumentationTargetInstance(): DynamicTestContract {
@@ -43,21 +45,21 @@ private fun assertControlFlow(expectedLocations: List<Int>) {
 
 class CoverageInstrumentationTest {
 
-    private val constructorStart = 54445
-    private val selfCheckStart = 8397
-    private val ifFirstBranch = 1555
-    private val ifEnd = 26354
-    private val outerForCondition = 37842
-    private val outerForBody = 53325
-    private val innerForCondition = 38432
-    private val innerForBody = 5673
-    private val innerForBodyIfFirstRun = 2378
-    private val innerForBodyIfSecondRun = 57606
-    private val innerForIncrementCounter = 7617
-    private val outerForIncrementCounter = 14668
-    private val outerForAfter = 9328
-    private val fooStart = 32182
-    private val barStart = 1381
+    private val constructorReturn = 0
+    private val ifFirstBranch = 1
+    @Suppress("unused")
+    private val ifSecondBranch = 2
+    private val ifEnd = 3
+    private val outerForCondition = 4
+    private val innerForBodyIfFirstRun = 6
+    private val innerForBodyIfSecondRun = 5
+    private val innerForIncrementCounter = 7
+    private val outerForIncrementCounter = 8
+    private val selfCheckReturn = 9
+    private val fooReturn = 10
+    private val barReturn = 11
+    @Suppress("unused")
+    private val bazReturn = 12
 
     @Test
     fun testOriginal() {
@@ -71,35 +73,28 @@ class CoverageInstrumentationTest {
 
         val innerForFirstRunControlFlow = mutableListOf<Int>().apply {
             repeat(5) {
-                addAll(listOf(innerForCondition, innerForBody, innerForBodyIfFirstRun, innerForIncrementCounter))
+                addAll(listOf(innerForBodyIfFirstRun, innerForIncrementCounter))
             }
-            add(innerForCondition)
         }.toList()
         val innerForSecondRunControlFlow = mutableListOf<Int>().apply {
             repeat(5) {
-                addAll(listOf(innerForCondition, innerForBody, innerForBodyIfSecondRun, innerForIncrementCounter))
+                addAll(listOf(innerForBodyIfSecondRun, innerForIncrementCounter))
             }
-            add(innerForCondition)
         }.toList()
-        val outerForControlFlow = listOf(outerForCondition, outerForBody) +
-            innerForFirstRunControlFlow +
-            listOf(outerForIncrementCounter, outerForCondition, outerForBody) +
-            innerForSecondRunControlFlow +
-            listOf(outerForIncrementCounter, outerForCondition)
+        val outerForControlFlow =
+            listOf(outerForCondition) +
+                innerForFirstRunControlFlow +
+                listOf(outerForIncrementCounter, outerForCondition) +
+                innerForSecondRunControlFlow +
+                listOf(outerForIncrementCounter)
 
         assertControlFlow(
-            listOf(constructorStart, selfCheckStart, ifFirstBranch, ifEnd) +
+            listOf(constructorReturn, ifFirstBranch, ifEnd) +
                 outerForControlFlow +
-                listOf(outerForAfter, fooStart, barStart)
+                listOf(barReturn, fooReturn, selfCheckReturn)
         )
     }
 
-    /**
-     * Computes the position of the counter in the coverage map to be incremented when control flows
-     * from the first member of [blocks] to the second.
-     */
-    fun edge(blocks: Pair<Int, Int>) = (blocks.first shr 1) xor blocks.second
-
     @OptIn(ExperimentalUnsignedTypes::class)
     @Test
     fun testCounters() {
@@ -107,9 +102,9 @@ class CoverageInstrumentationTest {
 
         val target = getInstrumentedInstrumentationTargetInstance()
         // The constructor of the target is run only once.
-        val takenOnceEdge = edge(constructorStart to selfCheckStart)
-        // Control flows from the start of selfCheck to the first if branch once per run.
-        val takenOnEveryRunEdge = edge(selfCheckStart to ifFirstBranch)
+        val takenOnceEdge = constructorReturn
+        // Control flows through the first if branch once per run.
+        val takenOnEveryRunEdge = ifFirstBranch
 
         for (i in 1..300) {
             assertSelfCheck(target)
@@ -128,7 +123,9 @@ class CoverageInstrumentationTest {
         // Make the patched class available in bazel-testlogs/.../test.outputs for manual inspection.
         val outDir = System.getenv("TEST_UNDECLARED_OUTPUTS_DIR")
         File("$outDir/${CoverageInstrumentationSpecialCasesTarget::class.simpleName}.class").writeBytes(originalBytecode)
-        File("$outDir/${CoverageInstrumentationSpecialCasesTarget::class.simpleName}.patched.class").writeBytes(patchedBytecode)
+        File("$outDir/${CoverageInstrumentationSpecialCasesTarget::class.simpleName}.patched.class").writeBytes(
+            patchedBytecode
+        )
         val patchedClass = bytecodeToClass(CoverageInstrumentationSpecialCasesTarget::class.java.name, patchedBytecode)
         // Trigger a class load
         patchedClass.declaredMethods
diff --git a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/MockCoverageMap.java b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/MockCoverageMap.java
index b289d20c..787ea493 100644
--- a/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/MockCoverageMap.java
+++ b/agent/src/test/java/com/code_intelligence/jazzer/instrumentor/MockCoverageMap.java
@@ -33,8 +33,7 @@ public class MockCoverageMap {
         updated_pos = i;
       }
     }
-    int cur_location = updated_pos ^ prev_location;
-    locations.add(cur_location);
+    locations.add(updated_pos);
     System.arraycopy(mem.array(), 0, previous_mem.array(), 0, SIZE);
   }