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.

16 files changed, 415 insertions, 250 deletions

diff --git a/WORKSPACE.bazel b/WORKSPACE.bazel
index 03e90217..c6a0f0cb 100644
--- a/WORKSPACE.bazel
+++ b/WORKSPACE.bazel
@@ -183,3 +183,20 @@ http_archive(
     strip_prefix = "libjpeg-turbo-2.0.90",
     url = "https://github.com/libjpeg-turbo/libjpeg-turbo/archive/2.0.90.tar.gz",
 )
+
+# JaCoCo
+
+jacoco_commit = "178d49870056b8a1f8ea6915e804d28b0dda5609"
+
+jacoco_sha = "da48fb5ae4ec3ffc659d4de18232aedea99476935f4ce4b0605f2d6aa1dc2553"
+
+http_archive(
+    name = "jacoco_internal",
+    build_file = "//third_party:jacoco_internal.BUILD",
+    patches = [
+        "//third_party:jacoco-make-probe-inserter-subclassable.patch",
+    ],
+    sha256 = jacoco_sha,
+    strip_prefix = "jacoco-%s" % jacoco_commit,
+    url = "https://github.com/jacoco/jacoco/archive/178d49870056b8a1f8ea6915e804d28b0dda5609.tar.gz",
+)
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/agent/Agent.kt b/agent/src/main/java/com/code_intelligence/jazzer/agent/Agent.kt
index fabfa507..1a0cccc4 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/agent/Agent.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/agent/Agent.kt
@@ -18,7 +18,6 @@ package com.code_intelligence.jazzer.agent
 
 import com.code_intelligence.jazzer.instrumentor.InstrumentationType
 import com.code_intelligence.jazzer.instrumentor.loadHooks
-import com.code_intelligence.jazzer.runtime.CoverageMap
 import com.code_intelligence.jazzer.runtime.ManifestUtils
 import java.lang.instrument.Instrumentation
 
@@ -29,7 +28,6 @@ val KNOWN_ARGUMENTS = listOf(
     "custom_hook_excludes",
     "trace",
     "custom_hooks",
-    "cov_size",
 )
 
 fun premain(agentArgs: String?, instrumentation: Instrumentation) {
@@ -76,12 +74,6 @@ fun premain(agentArgs: String?, instrumentation: Instrumentation) {
         }
     }.toSet()
 
-    argumentMap["cov_size"]?.let {
-        require(it.size == 1) { "cov_size can only be specified once" }
-        CoverageMap.reinit(it.first().toInt())
-        println("INFO: Using custom coverage map size: ${CoverageMap.SIZE}")
-    }
-
     val runtimeInstrumentor = RuntimeInstrumentor(classNameGlobber, dependencyClassNameGlobber, instrumentationTypes)
     instrumentation.apply {
         addTransformer(runtimeInstrumentor)
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/AFLCoverageMapInstrumentor.kt b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/AFLCoverageMapInstrumentor.kt
deleted file mode 100644
index bbdf867b..00000000
--- a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/AFLCoverageMapInstrumentor.kt
+++ /dev/null
@@ -1,164 +0,0 @@
-// Copyright 2021 Code Intelligence GmbH
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package com.code_intelligence.jazzer.instrumentor
-
-import com.code_intelligence.jazzer.runtime.CoverageMap
-import org.objectweb.asm.ClassReader
-import org.objectweb.asm.ClassWriter
-import org.objectweb.asm.Opcodes
-import org.objectweb.asm.tree.AbstractInsnNode
-import org.objectweb.asm.tree.AnnotationNode
-import org.objectweb.asm.tree.ClassNode
-import org.objectweb.asm.tree.FieldInsnNode
-import org.objectweb.asm.tree.FrameNode
-import org.objectweb.asm.tree.InsnList
-import org.objectweb.asm.tree.InsnNode
-import org.objectweb.asm.tree.IntInsnNode
-import org.objectweb.asm.tree.JumpInsnNode
-import org.objectweb.asm.tree.LabelNode
-import org.objectweb.asm.tree.LdcInsnNode
-import org.objectweb.asm.tree.LineNumberNode
-import org.objectweb.asm.tree.LocalVariableNode
-import org.objectweb.asm.tree.MethodInsnNode
-import org.objectweb.asm.tree.MethodNode
-import org.objectweb.asm.tree.TryCatchBlockNode
-import org.objectweb.asm.tree.TypeInsnNode
-
-internal class AFLCoverageMapInstrumentor(coverageMapClass: Class<*> = CoverageMap::class.java) : Instrumentor {
-
-    private val coverageMapInternalClassName = coverageMapClass.name.replace('.', '/')
-    private val testing = coverageMapClass != CoverageMap::class.java
-    private lateinit var random: DeterministicRandom
-
-    override fun instrument(bytecode: ByteArray): ByteArray {
-        val node = ClassNode()
-        val reader = ClassReader(bytecode)
-        reader.accept(node, 0)
-        random = DeterministicRandom("coverage", node.name)
-        for (method in node.methods) {
-            if (shouldInstrument(method)) {
-                for (inst in instrumentationPoints(method)) {
-                    method.instructions.insertBefore(inst, coverageInstrumentation())
-                }
-                method.instructions.insert(coverageInstrumentation())
-            }
-        }
-
-        val writer = ClassWriter(ClassWriter.COMPUTE_MAXS)
-        node.accept(writer)
-        return writer.toByteArray()
-    }
-
-    /**
-     * Applies bytecode instrumentation equivalent to the branch point code used by AFL.
-     * {@link https://lcamtuf.coredump.cx/afl/technical_details.txt}
-     */
-    private fun coverageInstrumentation(): InsnList {
-        // cur_location = <COMPILE_TIME_RANDOM>;
-        val cur_location = random.nextInt(CoverageMap.SIZE)
-        return InsnList().apply {
-            // Perform the following byte increment saturating at 255:
-            // mem[cur_location ^ prev_location]++;
-            add(FieldInsnNode(Opcodes.GETSTATIC, coverageMapInternalClassName, "mem", "Ljava/nio/ByteBuffer;"))
-            // Stack: mem
-            add(LdcInsnNode(cur_location))
-            // Stack: mem | cur_location
-            add(FieldInsnNode(Opcodes.GETSTATIC, coverageMapInternalClassName, "prev_location", "I"))
-            // Stack: mem | cur_location | prev_location
-            add(InsnNode(Opcodes.IXOR))
-            // Stack: mem | cur_location ^ prev_location
-            add(InsnNode(Opcodes.DUP2))
-            // Stack: mem | cur_location ^ prev_location | mem | cur_location ^ prev_location
-            add(MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/nio/ByteBuffer", "get", "(I)B", false))
-            // Stack: mem | cur_location ^ prev_location | counter (sign-extended to int)
-            add(IntInsnNode(Opcodes.SIPUSH, 0x00ff))
-            // Stack: mem | cur_location ^ prev_location | counter (sign-extended to int) | 0x000000ff
-            add(InsnNode(Opcodes.IAND))
-            // Stack: mem | cur_location ^ prev_location | counter (zero-extended to int)
-            add(InsnNode(Opcodes.ICONST_1))
-            // Stack: mem | cur_location ^ prev_location | counter | 1
-            add(InsnNode(Opcodes.IADD))
-            // Stack: mem | cur_location ^ prev_location | counter + 1
-            add(InsnNode(Opcodes.DUP))
-            // Stack: mem | cur_location ^ prev_location | counter + 1 | counter + 1
-            add(IntInsnNode(Opcodes.BIPUSH, 8))
-            // Stack: mem | cur_location ^ prev_location | counter + 1 | counter + 1 | 8
-            add(InsnNode(Opcodes.ISHR))
-            // Stack: mem | cur_location ^ prev_location | counter + 1 | 1 if the increment overflowed, 0 otherwise
-            add(InsnNode(Opcodes.ISUB))
-            // Stack: mem | cur_location ^ prev_location | counter if the increment overflowed, counter + 1 otherwise
-            add(MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/nio/ByteBuffer", "put", "(IB)Ljava/nio/ByteBuffer;", false))
-            // Stack: mem
-            add(InsnNode(Opcodes.POP))
-            if (testing) {
-                add(MethodInsnNode(Opcodes.INVOKESTATIC, coverageMapInternalClassName, "updated", "()V", false))
-            }
-            // prev_location = cur_location >> 1;
-            add(IntInsnNode(Opcodes.SIPUSH, cur_location shr 1))
-            add(FieldInsnNode(Opcodes.PUTSTATIC, coverageMapInternalClassName, "prev_location", "I"))
-        }
-    }
-
-    private fun skipFixedPositionNodes(node: AbstractInsnNode): AbstractInsnNode? {
-        var nextFreeNode: AbstractInsnNode? = node
-        while (true) {
-            when (nextFreeNode) {
-                is LabelNode,
-                is LineNumberNode,
-                is FrameNode,
-                is AnnotationNode,
-                is TryCatchBlockNode,
-                is LocalVariableNode -> {
-                    nextFreeNode = nextFreeNode.next
-                }
-                is TypeInsnNode -> {
-                    // NEW instructions are always accompanied by label nodes that tie them to
-                    // `Uninitialized` entries in the stack map table. If a NEW instruction is
-                    // a target of a jump, then this label node is reused as the target of the
-                    // jump and we must not insert other instrumentations in its place.
-                    // Since NEW does not invoke any user-supplied code (which only happens in
-                    // <init>), we can simply insert instrumentation after the NEW without
-                    // misattributing coverage in all but one case: If the NEW throws an
-                    // OutOfMemoryError, coverage for the basic block containing it will not have
-                    // been updated.
-                    if (nextFreeNode.opcode == Opcodes.NEW) {
-                        nextFreeNode = nextFreeNode.next
-                    } else {
-                        return nextFreeNode
-                    }
-                }
-                else -> return nextFreeNode
-            }
-        }
-    }
-
-    private fun instrumentationPoints(method: MethodNode): Set<AbstractInsnNode> =
-        method.instructions.asSequence()
-            .filterIsInstance<JumpInsnNode>()
-            .flatMap {
-                // At a jump node control passes either to the target of the jump (always a
-                // label node followed by the actual next instruction) or to the instruction
-                // right after the jump (possibly skipping a label). We want to return the
-                // instruction prior to the next node that must not remain in the position
-                // right after the jump (which is the case for labels, frames and NEW
-                // instructions, among others).
-                val nextIfJumping = skipFixedPositionNodes(it.label)
-                // GOTO is an unconditional jump and execution thus never passes to the
-                // next instruction.
-                val nextIfNotJumping = if (it.opcode != Opcodes.GOTO) skipFixedPositionNodes(it.next) else null
-                listOfNotNull(nextIfJumping, nextIfNotJumping)
-            }
-            .toSet()
-}
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/BUILD.bazel b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/BUILD.bazel
index 03f82df0..4048a3dd 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/BUILD.bazel
+++ b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/BUILD.bazel
@@ -12,6 +12,7 @@ kt_jvm_library(
     deps = [
         "//agent/src/main/java/com/code_intelligence/jazzer/runtime",
         "@com_github_jetbrains_kotlin//:kotlin-reflect",
+        "@jacoco_internal",
         "@maven//:org_ow2_asm_asm",
         "@maven//:org_ow2_asm_asm_commons",
     ],
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/ClassInstrumentor.kt b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/ClassInstrumentor.kt
index 2977f15f..7921ad3b 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/ClassInstrumentor.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/ClassInstrumentor.kt
@@ -24,7 +24,7 @@ class ClassInstrumentor constructor(bytecode: ByteArray) {
         private set
 
     fun coverage() {
-        instrumentedBytecode = AFLCoverageMapInstrumentor().instrument(instrumentedBytecode)
+        instrumentedBytecode = EdgeCoverageInstrumentor.instrument(instrumentedBytecode)
     }
 
     fun traceDataFlow(instrumentations: Set<InstrumentationType>) {
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt
new file mode 100644
index 00000000..0306f667
--- /dev/null
+++ b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt
@@ -0,0 +1,215 @@
+// Copyright 2021 Code Intelligence GmbH
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.code_intelligence.jazzer.instrumentor
+
+import com.code_intelligence.jazzer.runtime.CoverageMap
+import org.jacoco.core.internal.flow.ClassProbesAdapter
+import org.jacoco.core.internal.flow.ClassProbesVisitor
+import org.jacoco.core.internal.instr.ClassInstrumenter
+import org.jacoco.core.internal.instr.IProbeArrayStrategy
+import org.jacoco.core.internal.instr.IProbeInserterFactory
+import org.jacoco.core.internal.instr.InstrSupport
+import org.jacoco.core.internal.instr.ProbeInserter
+import org.objectweb.asm.ClassReader
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.ClassWriter
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+import kotlin.math.max
+
+object EdgeCoverageInstrumentor : Instrumentor {
+    override fun instrument(bytecode: ByteArray): ByteArray {
+        val reader = InstrSupport.classReaderFor(bytecode)
+        val writer = ClassWriter(reader, 0)
+        val strategy = EdgeCoverageProbeArrayStrategy()
+        val version = InstrSupport.getMajorVersion(reader)
+        val visitor = EdgeCoverageClassProbesAdapter(
+            ClassInstrumenter(strategy, EdgeCoverageProbeInserterFactory, writer),
+            InstrSupport.needsFrames(version)
+        )
+        reader.accept(visitor, ClassReader.EXPAND_FRAMES)
+        return writer.toByteArray()
+    }
+
+    var coverageMapClass: Class<*> = CoverageMap::class.java
+        private set
+    val coverageMapInternalClassName
+        get() = coverageMapClass.name.replace('.', '/')
+
+    fun setCoverageMapClassForTestingOnly(coverageMap: Class<*>) {
+        coverageMapClass = coverageMap
+    }
+    val isTesting
+        get() = coverageMapClass != CoverageMap::class.java
+
+    fun resetNextGlobalEdgeIdForTestingOnly() {
+        nextGlobalEdgeId = 0
+    }
+}
+
+private var nextGlobalEdgeId = 0
+
+private fun nextEdgeId(): Int = nextGlobalEdgeId++
+
+/**
+ * The maximal number of stack elements used by [loadCoverageMap].
+ */
+private const val LOAD_COVERAGE_MAP_STACK_SIZE = 1
+
+/**
+ * Inject bytecode that loads the coverage map into local variable [variable].
+ */
+private fun loadCoverageMap(mv: MethodVisitor, variable: Int) {
+    mv.apply {
+        visitFieldInsn(
+            Opcodes.GETSTATIC,
+            EdgeCoverageInstrumentor.coverageMapInternalClassName,
+            "mem",
+            "Ljava/nio/ByteBuffer;"
+        )
+        // Stack: mem (maxStack: 1)
+        visitVarInsn(Opcodes.ASTORE, variable)
+    }
+}
+
+/**
+ * The maximal number of stack elements used by [instrumentControlFlowEdge].
+ */
+private const val INSTRUMENT_CONTROL_FLOW_EDGE_STACK_SIZE = 5
+
+/**
+ * Inject bytecode instrumentation on a control flow edge with ID [edgeId]. The coverage map can be loaded from local
+ * variable [variable].
+ */
+private fun instrumentControlFlowEdge(mv: MethodVisitor, edgeId: Int, variable: Int) {
+    mv.apply {
+        visitVarInsn(Opcodes.ALOAD, variable)
+        // Stack: mem
+        push(edgeId)
+        // Stack: mem | edgeId
+        visitInsn(Opcodes.DUP2)
+        // Stack: mem | edgeId | mem | edgeId
+        visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/nio/ByteBuffer", "get", "(I)B", false)
+        // Stack: mem | edgeId | counter (sign-extended to int)
+        push(0xff)
+        // Stack: mem | edgeId | counter (sign-extended to int) | 0x000000ff
+        visitInsn(Opcodes.IAND)
+        // Stack: mem | edgeId | counter (zero-extended to int)
+        push(1)
+        // Stack: mem | edgeId | counter | 1
+        visitInsn(Opcodes.IADD)
+        // Stack: mem | edgeId | counter + 1
+        visitInsn(Opcodes.DUP)
+        // Stack: mem | edgeId | counter + 1 | counter + 1
+        push(8)
+        // Stack: mem | edgeId | counter + 1 | counter + 1 | 8 (maxStack: +5)
+        visitInsn(Opcodes.ISHR)
+        // Stack: mem | edgeId | counter + 1 | 1 if the increment overflowed, 0 otherwise
+        visitInsn(Opcodes.ISUB)
+        // Stack: mem | edgeId | counter if the increment overflowed, counter + 1 otherwise
+        visitMethodInsn(Opcodes.INVOKEVIRTUAL, "java/nio/ByteBuffer", "put", "(IB)Ljava/nio/ByteBuffer;", false)
+        // Stack: mem
+        visitInsn(Opcodes.POP)
+        if (EdgeCoverageInstrumentor.isTesting) {
+            visitMethodInsn(Opcodes.INVOKESTATIC, EdgeCoverageInstrumentor.coverageMapInternalClassName, "updated", "()V", false)
+        }
+    }
+}
+
+/**
+ * The maximal number of stack elements used by [instrumentMethodEdge].
+ */
+private const val INSTRUMENT_METHOD_EDGE_STACK_SIZE = INSTRUMENT_CONTROL_FLOW_EDGE_STACK_SIZE
+
+/**
+ * Inject bytecode instrumentation right before a method invocation (if needed). The coverage map can be loaded from
+ * local variable [variable].
+ *
+ * Note: Since not every method invocation might need instrumentation, an edge ID should only be generated if needed
+ * by calling [nextEdgeId].
+ */
+private fun instrumentMethodEdge(
+    mv: MethodVisitor,
+    variable: Int,
+    internalClassName: String,
+    methodName: String,
+    descriptor: String
+) {
+    if (internalClassName.startsWith("com/code_intelligence/jazzer/") && !EdgeCoverageInstrumentor.isTesting)
+        return
+    instrumentControlFlowEdge(mv, nextEdgeId(), variable)
+}
+
+
+// The remainder of this file interfaces with classes in org.jacoco.core.internal. Changes to this part should not be
+// necessary unless JaCoCo is updated or the way we instrument for coverage changes fundamentally.
+
+private class EdgeCoverageProbeInserter(
+    access: Int,
+    name: String,
+    desc: String,
+    mv: MethodVisitor,
+    arrayStrategy: IProbeArrayStrategy,
+) : ProbeInserter(access, name, desc, mv, arrayStrategy) {
+    override fun insertProbe(id: Int) {
+        instrumentControlFlowEdge(mv, id, variable)
+    }
+
+    override fun visitMaxs(maxStack: Int, maxLocals: Int) {
+        val maxStackIncrease = max(INSTRUMENT_CONTROL_FLOW_EDGE_STACK_SIZE, INSTRUMENT_METHOD_EDGE_STACK_SIZE)
+        val newMaxStack = max(maxStack + maxStackIncrease, LOAD_COVERAGE_MAP_STACK_SIZE)
+        val newMaxLocals = maxLocals + 1
+        mv.visitMaxs(newMaxStack, newMaxLocals)
+    }
+
+    override fun visitMethodInsn(
+        opcode: Int,
+        owner: String,
+        name: String,
+        descriptor: String,
+        isInterface: Boolean
+    ) {
+        instrumentMethodEdge(mv, variable, owner, name, descriptor)
+        mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface)
+    }
+}
+
+private object EdgeCoverageProbeInserterFactory : IProbeInserterFactory {
+    override fun makeProbeInserter(
+        access: Int,
+        name: String,
+        desc: String,
+        mv: MethodVisitor,
+        arrayStrategy: IProbeArrayStrategy,
+    ) = EdgeCoverageProbeInserter(access, name, desc, mv, arrayStrategy)
+}
+
+private class EdgeCoverageClassProbesAdapter(cv: ClassProbesVisitor, trackFrames: Boolean) :
+    ClassProbesAdapter(cv, trackFrames) {
+    override fun nextId(): Int = nextEdgeId()
+}
+
+private class EdgeCoverageProbeArrayStrategy : IProbeArrayStrategy {
+    override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
+        loadCoverageMap(mv, variable)
+        return LOAD_COVERAGE_MAP_STACK_SIZE
+    }
+
+    override fun addMembers(cv: ClassVisitor, probeCount: Int) {}
+}
+
+private fun MethodVisitor.push(value: Int) {
+    InstrSupport.push(this, value)
+}
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/runtime/CoverageMap.java b/agent/src/main/java/com/code_intelligence/jazzer/runtime/CoverageMap.java
index bafee76f..903f416f 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/runtime/CoverageMap.java
+++ b/agent/src/main/java/com/code_intelligence/jazzer/runtime/CoverageMap.java
@@ -22,25 +22,7 @@ import java.nio.ByteBuffer;
  * native code.
  */
 final public class CoverageMap {
-  // This needs to be a power of two to ensure two indices XORed together
-  // don't overflow the buffer.
   public static int SIZE = 65536;
   public static ByteBuffer mem = ByteBuffer.allocateDirect(SIZE);
-  @SuppressWarnings("unused") public static int prev_location = 0; // is used in byte code directly
 
-  private static int nextPowerOfTwo(int minSize) {
-    int nextPowerOfTwo = 1;
-    // Cannot represent 2^31 in a signed int.
-    for (int log2 = 0; log2 < 30; log2++) {
-      nextPowerOfTwo <<= 1;
-      if (nextPowerOfTwo >= minSize)
-        break;
-    }
-    return nextPowerOfTwo;
-  }
-
-  public static void reinit(int minSize) {
-    SIZE = nextPowerOfTwo(minSize);
-    mem = ByteBuffer.allocateDirect(SIZE);
-  }
 }
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);
   }
 
diff --git a/format.sh b/format.sh
index b8af889e..87fb5aca 100755
--- a/format.sh
+++ b/format.sh
@@ -3,7 +3,7 @@ find -name '*.cpp' -o -name '*.h' -o -name '*.java' | xargs clang-format-11 -i
 
 # Kotlin
 # curl -sSLO https://github.com/pinterest/ktlint/releases/download/0.40.0/ktlint && chmod a+x ktlint
-ktlint -F
+ktlint -F "agent/**/*.kt" "driver/**/*.kt" "examples/**/*.kt"
 
 # BUILD files
 # go get github.com/bazelbuild/buildtools/buildifier
diff --git a/maven.bzl b/maven.bzl
index 6eb82aad..899afcf8 100644
--- a/maven.bzl
+++ b/maven.bzl
@@ -20,6 +20,7 @@ JAZZER_API_COORDINATES = "com.code-intelligence:jazzer-api:%s" % JAZZER_API_VERS
 MAVEN_ARTIFACTS = [
     "org.ow2.asm:asm:9.1",
     "org.ow2.asm:asm-commons:9.1",
+    "org.ow2.asm:asm-tree:9.1",
     maven.artifact("junit", "junit", "4.12", testonly = True),
     "org.apache.commons:commons-imaging:1.0-alpha2",
     "com.mikesamuel:json-sanitizer:1.2.1",
diff --git a/maven_install.json b/maven_install.json
index 42d6462c..aad562cd 100644
--- a/maven_install.json
+++ b/maven_install.json
@@ -1,7 +1,7 @@
 {
     "dependency_tree": {
         "__AUTOGENERATED_FILE_DO_NOT_MODIFY_THIS_FILE_MANUALLY": "THERE_IS_NO_DATA_ONLY_ZUUL",
-        "__INPUT_ARTIFACTS_HASH": -1561364904,
+        "__INPUT_ARTIFACTS_HASH": -13066241,
         "__RESOLVED_ARTIFACTS_HASH": 2026903625,
         "conflict_resolution": {},
         "dependencies": [
diff --git a/third_party/BUILD.bazel b/third_party/BUILD.bazel
index e4e22821..fd65eb9b 100644
--- a/third_party/BUILD.bazel
+++ b/third_party/BUILD.bazel
@@ -1,4 +1,6 @@
 exports_files([
     "gflags-use-double-dash-args.patch",
+    "jacoco-make-probe-inserter-subclassable.patch",
+    "jacoco_internal.BUILD",
     "libjpeg_turbo.BUILD",
 ])
diff --git a/third_party/jacoco-make-probe-inserter-subclassable.patch b/third_party/jacoco-make-probe-inserter-subclassable.patch
new file mode 100644
index 00000000..6191a00f
--- /dev/null
+++ b/third_party/jacoco-make-probe-inserter-subclassable.patch
@@ -0,0 +1,123 @@
+// SPDX-License-Identifier: EPL-2.0 and Apache-2.0
+// These patches apply to JaCoCo (https://github.com/jacoco/jacoco) and are hereby made available under the terms of the
+// Eclipse Public License 2.0 available at:
+//   http://www.eclipse.org/legal/epl-2.0
+diff --git org.jacoco.core/src/org/jacoco/core/internal/flow/LabelInfo.java org.jacoco.core/src/org/jacoco/core/internal/flow/LabelInfo.java
+index 122d62b5..5ba3cf8d 100644
+--- org.jacoco.core/src/org/jacoco/core/internal/flow/LabelInfo.java
++++ org.jacoco.core/src/org/jacoco/core/internal/flow/LabelInfo.java
+@@ -141,8 +141,7 @@ public final class LabelInfo {
+ 	 */
+ 	public static boolean needsProbe(final Label label) {
+ 		final LabelInfo info = get(label);
+-		return info != null && info.successor
+-				&& (info.multiTarget || info.methodInvocationLine);
++		return info != null && info.successor && info.multiTarget;
+ 	}
+ 
+ 	/**
+diff --git org.jacoco.core/src/org/jacoco/core/internal/instr/ClassInstrumenter.java org.jacoco.core/src/org/jacoco/core/internal/instr/ClassInstrumenter.java
+index 476c9e34..bc192dc6 100644
+--- org.jacoco.core/src/org/jacoco/core/internal/instr/ClassInstrumenter.java
++++ org.jacoco.core/src/org/jacoco/core/internal/instr/ClassInstrumenter.java
+@@ -24,6 +24,7 @@ import org.objectweb.asm.MethodVisitor;
+ public class ClassInstrumenter extends ClassProbesVisitor {
+ 
+ 	private final IProbeArrayStrategy probeArrayStrategy;
++	private final IProbeInserterFactory probeInserterFactory;
+ 
+ 	private String className;
+ 
+@@ -40,6 +41,22 @@ public class ClassInstrumenter extends ClassProbesVisitor {
+ 			final ClassVisitor cv) {
+ 		super(cv);
+ 		this.probeArrayStrategy = probeArrayStrategy;
++		this.probeInserterFactory = new IProbeInserterFactory() {
++			@Override
++			public ProbeInserter makeProbeInserter(int access, String name,
++					String desc, MethodVisitor mv,
++					IProbeArrayStrategy arrayStrategy) {
++				return new ProbeInserter(access, name, desc, mv, arrayStrategy);
++			}
++		};
++	}
++
++	public ClassInstrumenter(final IProbeArrayStrategy probeArrayStrategy,
++			final IProbeInserterFactory probeInserterFactory,
++			final ClassVisitor cv) {
++		super(cv);
++		this.probeArrayStrategy = probeArrayStrategy;
++		this.probeInserterFactory = probeInserterFactory;
+ 	}
+ 
+ 	@Override
+@@ -71,8 +88,9 @@ public class ClassInstrumenter extends ClassProbesVisitor {
+ 			return null;
+ 		}
+ 		final MethodVisitor frameEliminator = new DuplicateFrameEliminator(mv);
+-		final ProbeInserter probeVariableInserter = new ProbeInserter(access,
+-				name, desc, frameEliminator, probeArrayStrategy);
++		final ProbeInserter probeVariableInserter =
++				probeInserterFactory.makeProbeInserter(access, name, desc,
++						frameEliminator, probeArrayStrategy);
+ 		return new MethodInstrumenter(probeVariableInserter,
+ 				probeVariableInserter);
+ 	}
+diff --git org.jacoco.core/src/org/jacoco/core/internal/instr/IProbeInserterFactory.java org.jacoco.core/src/org/jacoco/core/internal/instr/IProbeInserterFactory.java
+new file mode 100644
+index 00000000..19c2a7e2
+--- /dev/null
++++ org.jacoco.core/src/org/jacoco/core/internal/instr/IProbeInserterFactory.java
+@@ -0,0 +1,8 @@
++package org.jacoco.core.internal.instr;
++
++import org.objectweb.asm.MethodVisitor;
++
++public interface IProbeInserterFactory {
++    ProbeInserter makeProbeInserter(int access, String name, String desc,
++            MethodVisitor mv, IProbeArrayStrategy arrayStrategy);
++}
+diff --git org.jacoco.core/src/org/jacoco/core/internal/instr/InstrSupport.java org.jacoco.core/src/org/jacoco/core/internal/instr/InstrSupport.java
+index 71808ac8..3df93f63 100644
+--- org.jacoco.core/src/org/jacoco/core/internal/instr/InstrSupport.java
++++ org.jacoco.core/src/org/jacoco/core/internal/instr/InstrSupport.java
+@@ -78,7 +78,7 @@ public final class InstrSupport {
+ 	 * Data type of the field that stores coverage information for a class (
+ 	 * <code>boolean[]</code>).
+ 	 */
+-	public static final String DATAFIELD_DESC = "[Z";
++	public static final String DATAFIELD_DESC = "java/nio/ByteBuffer";
+ 
+ 	// === Init Method ===
+ 
+diff --git org.jacoco.core/src/org/jacoco/core/internal/instr/ProbeInserter.java org.jacoco.core/src/org/jacoco/core/internal/instr/ProbeInserter.java
+index 0f5b99ff..ba5daa6d 100644
+--- org.jacoco.core/src/org/jacoco/core/internal/instr/ProbeInserter.java
++++ org.jacoco.core/src/org/jacoco/core/internal/instr/ProbeInserter.java
+@@ -25,7 +25,7 @@ import org.objectweb.asm.TypePath;
+  * addition the probe array has to be retrieved at the beginning of the method
+  * and stored in a local variable.
+  */
+-class ProbeInserter extends MethodVisitor implements IProbeInserter {
++public class ProbeInserter extends MethodVisitor implements IProbeInserter {
+ 
+ 	private final IProbeArrayStrategy arrayStrategy;
+ 
+@@ -36,7 +36,7 @@ class ProbeInserter extends MethodVisitor implements IProbeInserter {
+ 	private final boolean clinit;
+ 
+ 	/** Position of the inserted variable. */
+-	private final int variable;
++	protected final int variable;
+ 
+ 	/** Maximum stack usage of the code to access the probe array. */
+ 	private int accessorStackSize;
+@@ -56,7 +56,7 @@ class ProbeInserter extends MethodVisitor implements IProbeInserter {
+ 	 *            callback to create the code that retrieves the reference to
+ 	 *            the probe array
+ 	 */
+-	ProbeInserter(final int access, final String name, final String desc,
++	public ProbeInserter(final int access, final String name, final String desc,
+ 			final MethodVisitor mv, final IProbeArrayStrategy arrayStrategy) {
+ 		super(InstrSupport.ASM_API_VERSION, mv);
+ 		this.clinit = InstrSupport.CLINIT_NAME.equals(name);
diff --git a/third_party/jacoco_internal.BUILD b/third_party/jacoco_internal.BUILD
new file mode 100644
index 00000000..586f8a25
--- /dev/null
+++ b/third_party/jacoco_internal.BUILD
@@ -0,0 +1,15 @@
+java_library(
+    name = "jacoco_internal",
+    srcs = glob([
+        "org.jacoco.core/src/org/jacoco/core/**/*.java",
+    ]),
+    javacopts = [
+        "-Xep:EqualsHashCode:WARN",
+    ],
+    deps = [
+        "@maven//:org_ow2_asm_asm",
+        "@maven//:org_ow2_asm_asm_commons",
+        "@maven//:org_ow2_asm_asm_tree",
+    ],
+    visibility = ["//visibility:public"],
+)