From 72fb549bb5bdac36e1eef31430ac7a1d258529fe Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Thu, 3 Aug 2023 13:47:58 +0300
Subject: [PATCH 01/13] Add working version of branch coverage

---
 utbot-instrumentation/build.gradle.kts        |   5 +
 .../org/jacoco/core/internal/instr/Api.kt     |  28 ++++
 .../instr/MyClassFieldProbeArrayStrategy.kt   |  93 ++++++++++++
 .../internal/instr/MyClassInstrumenter.kt     |  25 ++++
 .../internal/instr/MyMethodInstrumenter.kt    | 134 ++++++++++++++++++
 .../core/internal/instr/MyProbeInserter.kt    |  18 +++
 .../coverage/BranchCoverageInstrumentation.kt |  70 +++++++++
 .../coverage/CoverageInstrumentation.kt       |  10 +-
 .../instrumenter/Instrumenter.kt              |  18 +++
 9 files changed, 398 insertions(+), 3 deletions(-)
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassInstrumenter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodInstrumenter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyProbeInserter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt

diff --git a/utbot-instrumentation/build.gradle.kts b/utbot-instrumentation/build.gradle.kts
index 718aa6b29b..581354f2bf 100644
--- a/utbot-instrumentation/build.gradle.kts
+++ b/utbot-instrumentation/build.gradle.kts
@@ -14,6 +14,8 @@ val rdVersion: String by rootProject
 val mockitoVersion: String by rootProject
 val mockitoInlineVersion: String by rootProject
 
+val jacocoVersion: String by rootProject
+
 plugins {
     id("com.github.johnrengelman.shadow") version "7.1.2"
     id("java")
@@ -52,6 +54,7 @@ dependencies {
 
     implementation("org.ow2.asm:asm:$asmVersion")
     implementation("org.ow2.asm:asm-commons:$asmVersion")
+    implementation("org.ow2.asm:asm-util:$asmVersion")
     implementation("io.github.microutils:kotlin-logging:$kotlinLoggingVersion")
 
     implementation("com.jetbrains.rd:rd-framework:$rdVersion")
@@ -62,6 +65,8 @@ dependencies {
     implementation("org.mockito:mockito-core:$mockitoVersion")
     implementation("org.mockito:mockito-inline:$mockitoInlineVersion")
 
+    implementation("org.jacoco:org.jacoco.core:$jacocoVersion")
+
     implementation(project(":utbot-spring-commons-api"))
     if (projectType == springEdition || projectType == ultimateEdition) {
         fetchSpringCommonsJar(project(":utbot-spring-commons", configuration = "springCommonsJar"))
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
new file mode 100644
index 0000000000..2ae5e23b80
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -0,0 +1,28 @@
+package org.jacoco.core.internal.instr
+
+import org.jacoco.core.internal.flow.ClassProbesAdapter
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.ClassWriter
+import org.objectweb.asm.util.TraceClassVisitor
+import java.io.PrintWriter
+import java.io.StringWriter
+
+val sw: StringWriter = StringWriter()
+
+fun createJacocoClassVisitorForCollectionBranchCoverage(writer: ClassWriter): ClassVisitor {
+    val strategy = NoneProbeArrayStrategy()
+    val tcv = TraceClassVisitor(writer, PrintWriter(sw))
+    return ClassProbesAdapter(
+        MyClassInstrumenter(strategy, tcv),
+        true
+    )
+}
+
+fun createJacocoClassVisitorForBytecodeInstrumentation(writer: ClassWriter, className: String): ClassVisitor {
+    val strategy = MyClassFieldProbeArrayStrategy(className)
+    val tcv = TraceClassVisitor(writer, PrintWriter(sw))
+    return ClassProbesAdapter(
+        ClassInstrumenter(strategy, tcv),
+        true
+    )
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt
new file mode 100644
index 0000000000..521b0a37f9
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt
@@ -0,0 +1,93 @@
+package org.jacoco.core.internal.instr
+
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.Label
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+import org.utbot.instrumentation.Settings
+
+/**
+ * Inspired by [org.jacoco.core.internal.instr.ClassFieldProbeArrayStrategy]
+ */
+class MyClassFieldProbeArrayStrategy(private val className: String) : IProbeArrayStrategy {
+
+    private val FRAME_STACK_EMPTY = arrayOf<Any>(0)
+
+    private val FRAME_LOCALS_EMPTY = arrayOfNulls<Any>(0)
+
+    override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
+        mv.visitMethodInsn(
+            Opcodes.INVOKESTATIC,
+            className,
+            // TODO change method name
+            InstrSupport.INITMETHOD_NAME,
+            InstrSupport.INITMETHOD_DESC,
+            false
+        )
+        mv.visitVarInsn(Opcodes.ASTORE, variable)
+        return 1
+    }
+
+    override fun addMembers(cv: ClassVisitor, probeCount: Int) {
+        createDataField(cv)
+        createInitMethod(cv, probeCount)
+    }
+
+    private fun createDataField(cv: ClassVisitor) {
+        cv.visitField(
+            Opcodes.ACC_PUBLIC + Opcodes.ACC_STATIC + Opcodes.ACC_SYNTHETIC,
+            Settings.PROBES_ARRAY_NAME,
+            Settings.PROBES_ARRAY_DESC,
+            null,
+            null
+        )
+    }
+
+    private fun createInitMethod(cv: ClassVisitor, probeCount: Int) {
+        val mv = cv.visitMethod(
+            InstrSupport.INITMETHOD_ACC,
+            // TODO change method name
+            InstrSupport.INITMETHOD_NAME,
+            InstrSupport.INITMETHOD_DESC,
+            null,
+            null
+        )
+
+        mv.visitCode()
+        mv.visitFieldInsn(
+            Opcodes.GETSTATIC,
+            className,
+            Settings.PROBES_ARRAY_NAME,
+            Settings.PROBES_ARRAY_DESC
+        )
+        val alreadyInitialized = Label()
+        mv.visitJumpInsn(Opcodes.IFNONNULL, alreadyInitialized)
+        mv.visitIntInsn(Opcodes.BIPUSH, probeCount)
+        mv.visitIntInsn(Opcodes.NEWARRAY, Opcodes.T_BOOLEAN)
+        mv.visitFieldInsn(
+            Opcodes.PUTSTATIC,
+            className,
+            Settings.PROBES_ARRAY_NAME,
+            Settings.PROBES_ARRAY_DESC
+        )
+        mv.visitFrame(
+            Opcodes.F_NEW,
+            0,
+            FRAME_LOCALS_EMPTY,
+            0,
+            FRAME_STACK_EMPTY
+        )
+        mv.visitLabel(alreadyInitialized)
+        mv.visitFieldInsn(
+            Opcodes.GETSTATIC,
+            className,
+            Settings.PROBES_ARRAY_NAME,
+            Settings.PROBES_ARRAY_DESC
+        )
+        mv.visitInsn(Opcodes.ARETURN)
+
+        mv.visitMaxs(1, 0)
+        mv.visitEnd()
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassInstrumenter.kt
new file mode 100644
index 0000000000..0a4a9d0f2b
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassInstrumenter.kt
@@ -0,0 +1,25 @@
+package org.jacoco.core.internal.instr
+
+import org.jacoco.core.internal.flow.MethodProbesVisitor
+import org.objectweb.asm.ClassVisitor
+
+class MyClassInstrumenter(
+    probeArrayStrategy: IProbeArrayStrategy,
+    cv: ClassVisitor
+) : ClassInstrumenter(probeArrayStrategy, cv) {
+
+    override fun visitMethod(
+        access: Int,
+        name: String,
+        desc: String,
+        signature: String?,
+        exceptions: Array<out String>?
+    ): MethodProbesVisitor {
+        val mv = cv.visitMethod(access, name, desc, signature, exceptions)
+
+        val frameEliminator = DuplicateFrameEliminator(mv)
+        val probeVariableInserter = MyProbeInserter(name, frameEliminator)
+        return MyMethodInstrumenter(probeVariableInserter, probeVariableInserter)
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodInstrumenter.kt
new file mode 100644
index 0000000000..3039b35897
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodInstrumenter.kt
@@ -0,0 +1,134 @@
+package org.jacoco.core.internal.instr
+
+import org.jacoco.core.internal.flow.IFrame
+import org.jacoco.core.internal.flow.LabelInfo
+import org.jacoco.core.internal.flow.MethodProbesVisitor
+import org.objectweb.asm.Label
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+
+internal class MyMethodInstrumenter(
+    mv: MethodVisitor,
+    private val probeInserter: IProbeInserter
+) : MethodProbesVisitor(mv) {
+
+    override fun visitProbe(probeId: Int) {
+        probeInserter.insertProbe(probeId)
+    }
+
+    override fun visitInsnWithProbe(opcode: Int, probeId: Int) {
+        probeInserter.insertProbe(probeId)
+        mv.visitInsn(opcode)
+    }
+
+    override fun visitJumpInsnWithProbe(opcode: Int, label: Label?, probeId: Int, frame: IFrame) {
+        if (opcode == Opcodes.GOTO) {
+            probeInserter.insertProbe(probeId)
+            mv.visitJumpInsn(Opcodes.GOTO, label)
+        } else {
+            val intermediate = Label()
+            mv.visitJumpInsn(getInverted(opcode), intermediate)
+            probeInserter.insertProbe(probeId)
+            mv.visitJumpInsn(Opcodes.GOTO, label)
+            mv.visitLabel(intermediate)
+            frame.accept(mv)
+        }
+    }
+
+    private fun getInverted(opcode: Int): Int =
+        when (opcode) {
+            Opcodes.IFEQ -> Opcodes.IFNE
+            Opcodes.IFNE -> Opcodes.IFEQ
+            Opcodes.IFLT -> Opcodes.IFGE
+            Opcodes.IFGE -> Opcodes.IFLT
+            Opcodes.IFGT -> Opcodes.IFLE
+            Opcodes.IFLE -> Opcodes.IFGT
+            Opcodes.IF_ICMPEQ -> Opcodes.IF_ICMPNE
+            Opcodes.IF_ICMPNE -> Opcodes.IF_ICMPEQ
+            Opcodes.IF_ICMPLT -> Opcodes.IF_ICMPGE
+            Opcodes.IF_ICMPGE -> Opcodes.IF_ICMPLT
+            Opcodes.IF_ICMPGT -> Opcodes.IF_ICMPLE
+            Opcodes.IF_ICMPLE -> Opcodes.IF_ICMPGT
+            Opcodes.IF_ACMPEQ -> Opcodes.IF_ACMPNE
+            Opcodes.IF_ACMPNE -> Opcodes.IF_ACMPEQ
+            Opcodes.IFNULL -> Opcodes.IFNONNULL
+            Opcodes.IFNONNULL -> Opcodes.IFNULL
+            else -> throw IllegalArgumentException()
+        }
+
+    override fun visitTableSwitchInsnWithProbes(
+        min: Int,
+        max: Int,
+        dflt: Label,
+        labels: Array<Label>,
+        frame: IFrame
+    ) {
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        val newDflt: Label = createIntermediate(dflt)
+        val newLabels: Array<Label?> = createIntermediates(labels)
+        mv.visitTableSwitchInsn(min, max, newDflt, *newLabels)
+
+        insertIntermediateProbes(dflt, labels, frame)
+    }
+
+    override fun visitLookupSwitchInsnWithProbes(
+        dflt: Label,
+        keys: IntArray,
+        labels: Array<Label>,
+        frame: IFrame
+    ) {
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        val newDflt: Label = createIntermediate(dflt)
+        val newLabels: Array<Label?> = createIntermediates(labels)
+        mv.visitLookupSwitchInsn(newDflt, keys, newLabels)
+
+        insertIntermediateProbes(dflt, labels, frame)
+    }
+
+    private fun createIntermediates(labels: Array<Label>): Array<Label?> {
+        val intermediates = arrayOfNulls<Label>(labels.size)
+        for (i in labels.indices) {
+            intermediates[i] = createIntermediate(labels[i])
+        }
+        return intermediates
+    }
+
+    private fun createIntermediate(label: Label): Label {
+        val intermediate: Label
+        if (LabelInfo.getProbeId(label) == LabelInfo.NO_PROBE) {
+            intermediate = label
+        } else {
+            if (LabelInfo.isDone(label)) {
+                intermediate = LabelInfo.getIntermediateLabel(label)
+            } else {
+                intermediate = Label()
+                LabelInfo.setIntermediateLabel(label, intermediate)
+                LabelInfo.setDone(label)
+            }
+        }
+        return intermediate
+    }
+
+    private fun insertIntermediateProbe(label: Label, frame: IFrame) {
+        val probeId = LabelInfo.getProbeId(label)
+        if (probeId != LabelInfo.NO_PROBE && !LabelInfo.isDone(label)) {
+            mv.visitLabel(LabelInfo.getIntermediateLabel(label))
+            frame.accept(mv)
+            probeInserter.insertProbe(probeId)
+            mv.visitJumpInsn(Opcodes.GOTO, label)
+            LabelInfo.setDone(label)
+        }
+    }
+
+    private fun insertIntermediateProbes(dflt: Label, labels: Array<Label>, frame: IFrame) {
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        insertIntermediateProbe(dflt, frame)
+        for (l in labels) {
+            insertIntermediateProbe(l, frame)
+        }
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyProbeInserter.kt
new file mode 100644
index 0000000000..ea2b07e862
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyProbeInserter.kt
@@ -0,0 +1,18 @@
+package org.jacoco.core.internal.instr
+
+import org.objectweb.asm.MethodVisitor
+
+class MyProbeInserter(
+    private val name: String,
+    mv: MethodVisitor,
+) : IProbeInserter, MethodVisitor(InstrSupport.ASM_API_VERSION, mv) {
+
+    companion object {
+        val methodToProbes = mutableMapOf<String, MutableList<Int>>()
+    }
+
+    override fun insertProbe(id: Int) {
+        methodToProbes.getOrPut(name) { mutableListOf() }.add(id)
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
new file mode 100644
index 0000000000..30122de90f
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
@@ -0,0 +1,70 @@
+package org.utbot.instrumentation.instrumentation.coverage
+
+import org.jacoco.core.internal.instr.createJacocoClassVisitorForBytecodeInstrumentation
+import org.jacoco.core.internal.instr.sw
+import org.objectweb.asm.ClassReader
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.ClassWriter
+import org.utbot.common.withAccessibility
+import org.utbot.instrumentation.Settings
+import org.utbot.instrumentation.instrumentation.ArgumentList
+import org.utbot.instrumentation.instrumentation.Instrumentation
+import org.utbot.instrumentation.instrumentation.InvokeInstrumentation
+import org.utbot.instrumentation.instrumentation.instrumenter.ClassVisitorBuilder
+import org.utbot.instrumentation.instrumentation.instrumenter.Instrumenter
+import org.utbot.instrumentation.util.NoProbesArrayException
+import java.security.ProtectionDomain
+
+class BranchCoverageInstrumentation : CoverageInstrumentation() {
+
+    private val invokeInstrumentation = InvokeInstrumentation()
+
+    override fun invoke(
+        clazz: Class<*>,
+        methodSignature: String,
+        arguments: ArgumentList,
+        parameters: Any?
+    ): Result<*> {
+        val probesFieldName: String = Settings.PROBES_ARRAY_NAME
+        val visitedLinesField = clazz.fields.firstOrNull { it.name == probesFieldName }
+            ?: throw NoProbesArrayException(clazz, Settings.PROBES_ARRAY_NAME)
+
+        return visitedLinesField.withAccessibility {
+            invokeInstrumentation.invoke(clazz, methodSignature, arguments, parameters)
+        }
+    }
+
+    override fun <T : Any> methodToCollectCoverage(clazz: Class<out T>): Map<String, IntRange> {
+        return Instrumenter(clazz).computeMapOfRangesForBranchCoverage()
+    }
+
+    override fun transform(
+        loader: ClassLoader?,
+        className: String,
+        classBeingRedefined: Class<*>?,
+        protectionDomain: ProtectionDomain?,
+        classfileBuffer: ByteArray
+    ): ByteArray {
+        val instrumenter = Instrumenter(classfileBuffer)
+
+        instrumenter.visitClass(object : ClassVisitorBuilder<ClassVisitor> {
+            override val writerFlags: Int
+                get() = 0
+
+            override val readerParsingOptions: Int
+                get() = ClassReader.EXPAND_FRAMES
+
+            override fun build(writer: ClassWriter): ClassVisitor =
+                createJacocoClassVisitorForBytecodeInstrumentation(writer, className)
+        })
+
+        val result = sw.toString()
+        println(result)
+
+        return instrumenter.classByteCode
+    }
+
+    object Factory : Instrumentation.Factory<Result<*>, BranchCoverageInstrumentation> {
+        override fun create(): BranchCoverageInstrumentation = BranchCoverageInstrumentation()
+    }
+}
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
index ab3238a97b..b9f6729345 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
@@ -22,7 +22,7 @@ data class CoverageInfo(
 /**
  * This instrumentation allows collecting coverage after several calls.
  */
-class CoverageInstrumentation : Instrumentation<Result<*>> {
+open class CoverageInstrumentation : Instrumentation<Result<*>> {
     private val invokeWithStatics = InvokeWithStaticsInstrumentation()
 
     /**
@@ -53,7 +53,7 @@ class CoverageInstrumentation : Instrumentation<Result<*>> {
     /**
      * Collects coverage from the given [clazz] via reflection.
      */
-    fun <T : Any> collectCoverageInfo(clazz: Class<out T>): CoverageInfo {
+    open fun <T : Any> collectCoverageInfo(clazz: Class<out T>): CoverageInfo {
         val probesFieldName: String = Settings.PROBES_ARRAY_NAME
         val visitedLinesField = clazz.fields.firstOrNull { it.name == probesFieldName }
             ?: throw NoProbesArrayException(clazz, Settings.PROBES_ARRAY_NAME)
@@ -62,7 +62,7 @@ class CoverageInstrumentation : Instrumentation<Result<*>> {
             val visitedLines = visitedLinesField.get(null) as? BooleanArray
                 ?: throw CastProbesArrayException()
 
-            val methodToInstrRange = Instrumenter(clazz).computeMapOfRanges()
+            val methodToInstrRange = methodToCollectCoverage(clazz)
 
             val res = CoverageInfo(methodToInstrRange, visitedLines.mapIndexed { idx, b ->
                 if (b) idx else null
@@ -74,6 +74,10 @@ class CoverageInstrumentation : Instrumentation<Result<*>> {
         }
     }
 
+    protected open fun <T : Any> methodToCollectCoverage(clazz: Class<out T>): Map<String, IntRange> {
+        return Instrumenter(clazz).computeMapOfRanges()
+    }
+
     /**
      * Transforms bytecode such way that it becomes possible to get a coverage.
      *
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
index dfb403e0ac..fcde31717d 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
@@ -1,5 +1,6 @@
 package org.utbot.instrumentation.instrumentation.instrumenter
 
+import org.jacoco.core.internal.instr.*
 import org.objectweb.asm.ClassReader
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
@@ -49,6 +50,22 @@ class Instrumenter(classByteCode: ByteArray, val classLoader: ClassLoader? = nul
         return methodToListOfProbesInserter.methodToProbes.mapValues { (_, probes) -> (probes.first()..probes.last()) }
     }
 
+    fun computeMapOfRangesForBranchCoverage(): Map<String, IntRange> {
+
+        visitClass(object : ClassVisitorBuilder<ClassVisitor> {
+            override val writerFlags: Int
+                get() = 0
+
+            override val readerParsingOptions: Int
+                get() = ClassReader.EXPAND_FRAMES
+
+            override fun build(writer: ClassWriter): ClassVisitor =
+                createJacocoClassVisitorForCollectionBranchCoverage(writer)
+        })
+
+        return MyProbeInserter.methodToProbes.mapValues { (_, probes) -> (probes.first()..probes.last()) }
+    }
+
     fun addField(instanceFieldInitializer: InstanceFieldInitializer) {
         visitClass { writer -> AddFieldAdapter(writer, instanceFieldInitializer) }
     }
@@ -82,6 +99,7 @@ private class TunedClassWriter(
     override fun getClassLoader(): ClassLoader {
         return HandlerClassesLoader
     }
+
     override fun getCommonSuperClass(type1: String, type2: String): String {
         try {
             val info1 = typeInfo(type1)

From 0c5cfd662c834869a804bf57d28f3dd01380d0fa Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Thu, 3 Aug 2023 17:21:06 +0300
Subject: [PATCH 02/13] Refactor

---
 utbot-instrumentation/build.gradle.kts        |   2 +-
 .../org/jacoco/core/internal/instr/Api.kt     |  14 +-
 .../instr/MyClassFieldProbeArrayStrategy.kt   |  76 +---------
 .../internal/instr/MyMethodInstrumenter.kt    | 134 ------------------
 .../internal/instr/MyMethodToProbesVisitor.kt |  17 +++
 .../core/internal/instr/MyProbeInserter.kt    |  18 ---
 ...strumenter.kt => NoneClassInstrumenter.kt} |   9 +-
 .../core/internal/instr/NoneProbeInserter.kt  |  19 +++
 .../coverage/BranchCoverageInstrumentation.kt |  37 ++---
 .../coverage/CoverageInstrumentation.kt       |   4 +-
 .../coverage/StaticArrayFieldInitializer.kt   |  24 ++++
 .../instrumenter/Instrumenter.kt              |  14 +-
 12 files changed, 96 insertions(+), 272 deletions(-)
 delete mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodInstrumenter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodToProbesVisitor.kt
 delete mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyProbeInserter.kt
 rename utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/{MyClassInstrumenter.kt => NoneClassInstrumenter.kt} (65%)
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/StaticArrayFieldInitializer.kt

diff --git a/utbot-instrumentation/build.gradle.kts b/utbot-instrumentation/build.gradle.kts
index 581354f2bf..4191b25b6e 100644
--- a/utbot-instrumentation/build.gradle.kts
+++ b/utbot-instrumentation/build.gradle.kts
@@ -54,7 +54,7 @@ dependencies {
 
     implementation("org.ow2.asm:asm:$asmVersion")
     implementation("org.ow2.asm:asm-commons:$asmVersion")
-    implementation("org.ow2.asm:asm-util:$asmVersion")
+    implementation("org.ow2.asm:asm-util:$asmVersion") // TODO: remove
     implementation("io.github.microutils:kotlin-logging:$kotlinLoggingVersion")
 
     implementation("com.jetbrains.rd:rd-framework:$rdVersion")
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
index 2ae5e23b80..52d7cc0fff 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -1,28 +1,24 @@
 package org.jacoco.core.internal.instr
 
 import org.jacoco.core.internal.flow.ClassProbesAdapter
-import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
 import org.objectweb.asm.util.TraceClassVisitor
+import org.utbot.instrumentation.Settings
 import java.io.PrintWriter
 import java.io.StringWriter
 
 val sw: StringWriter = StringWriter()
 
-fun createJacocoClassVisitorForCollectionBranchCoverage(writer: ClassWriter): ClassVisitor {
+fun createJacocoClassVisitorForCollectionBranchCoverage(writer: ClassWriter): MyMethodToProbesVisitor {
     val strategy = NoneProbeArrayStrategy()
     val tcv = TraceClassVisitor(writer, PrintWriter(sw))
-    return ClassProbesAdapter(
-        MyClassInstrumenter(strategy, tcv),
-        true
-    )
+    return MyMethodToProbesVisitor(Settings.ASM_API, strategy, tcv)
 }
 
-fun createJacocoClassVisitorForBytecodeInstrumentation(writer: ClassWriter, className: String): ClassVisitor {
+fun createJacocoClassVisitorForBytecodeInstrumentation(writer: ClassWriter, className: String): ClassProbesAdapter {
     val strategy = MyClassFieldProbeArrayStrategy(className)
-    val tcv = TraceClassVisitor(writer, PrintWriter(sw))
     return ClassProbesAdapter(
-        ClassInstrumenter(strategy, tcv),
+        ClassInstrumenter(strategy, writer),
         true
     )
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt
index 521b0a37f9..ba420b0588 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt
@@ -1,93 +1,19 @@
 package org.jacoco.core.internal.instr
 
 import org.objectweb.asm.ClassVisitor
-import org.objectweb.asm.Label
 import org.objectweb.asm.MethodVisitor
 import org.objectweb.asm.Opcodes
 import org.utbot.instrumentation.Settings
 
-/**
- * Inspired by [org.jacoco.core.internal.instr.ClassFieldProbeArrayStrategy]
- */
 class MyClassFieldProbeArrayStrategy(private val className: String) : IProbeArrayStrategy {
 
-    private val FRAME_STACK_EMPTY = arrayOf<Any>(0)
-
-    private val FRAME_LOCALS_EMPTY = arrayOfNulls<Any>(0)
-
     override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
-        mv.visitMethodInsn(
-            Opcodes.INVOKESTATIC,
-            className,
-            // TODO change method name
-            InstrSupport.INITMETHOD_NAME,
-            InstrSupport.INITMETHOD_DESC,
-            false
-        )
+        mv.visitFieldInsn(Opcodes.GETSTATIC, className, Settings.PROBES_ARRAY_NAME, Settings.PROBES_ARRAY_DESC)
         mv.visitVarInsn(Opcodes.ASTORE, variable)
         return 1
     }
 
     override fun addMembers(cv: ClassVisitor, probeCount: Int) {
-        createDataField(cv)
-        createInitMethod(cv, probeCount)
-    }
-
-    private fun createDataField(cv: ClassVisitor) {
-        cv.visitField(
-            Opcodes.ACC_PUBLIC + Opcodes.ACC_STATIC + Opcodes.ACC_SYNTHETIC,
-            Settings.PROBES_ARRAY_NAME,
-            Settings.PROBES_ARRAY_DESC,
-            null,
-            null
-        )
-    }
-
-    private fun createInitMethod(cv: ClassVisitor, probeCount: Int) {
-        val mv = cv.visitMethod(
-            InstrSupport.INITMETHOD_ACC,
-            // TODO change method name
-            InstrSupport.INITMETHOD_NAME,
-            InstrSupport.INITMETHOD_DESC,
-            null,
-            null
-        )
-
-        mv.visitCode()
-        mv.visitFieldInsn(
-            Opcodes.GETSTATIC,
-            className,
-            Settings.PROBES_ARRAY_NAME,
-            Settings.PROBES_ARRAY_DESC
-        )
-        val alreadyInitialized = Label()
-        mv.visitJumpInsn(Opcodes.IFNONNULL, alreadyInitialized)
-        mv.visitIntInsn(Opcodes.BIPUSH, probeCount)
-        mv.visitIntInsn(Opcodes.NEWARRAY, Opcodes.T_BOOLEAN)
-        mv.visitFieldInsn(
-            Opcodes.PUTSTATIC,
-            className,
-            Settings.PROBES_ARRAY_NAME,
-            Settings.PROBES_ARRAY_DESC
-        )
-        mv.visitFrame(
-            Opcodes.F_NEW,
-            0,
-            FRAME_LOCALS_EMPTY,
-            0,
-            FRAME_STACK_EMPTY
-        )
-        mv.visitLabel(alreadyInitialized)
-        mv.visitFieldInsn(
-            Opcodes.GETSTATIC,
-            className,
-            Settings.PROBES_ARRAY_NAME,
-            Settings.PROBES_ARRAY_DESC
-        )
-        mv.visitInsn(Opcodes.ARETURN)
-
-        mv.visitMaxs(1, 0)
-        mv.visitEnd()
     }
 
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodInstrumenter.kt
deleted file mode 100644
index 3039b35897..0000000000
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodInstrumenter.kt
+++ /dev/null
@@ -1,134 +0,0 @@
-package org.jacoco.core.internal.instr
-
-import org.jacoco.core.internal.flow.IFrame
-import org.jacoco.core.internal.flow.LabelInfo
-import org.jacoco.core.internal.flow.MethodProbesVisitor
-import org.objectweb.asm.Label
-import org.objectweb.asm.MethodVisitor
-import org.objectweb.asm.Opcodes
-
-internal class MyMethodInstrumenter(
-    mv: MethodVisitor,
-    private val probeInserter: IProbeInserter
-) : MethodProbesVisitor(mv) {
-
-    override fun visitProbe(probeId: Int) {
-        probeInserter.insertProbe(probeId)
-    }
-
-    override fun visitInsnWithProbe(opcode: Int, probeId: Int) {
-        probeInserter.insertProbe(probeId)
-        mv.visitInsn(opcode)
-    }
-
-    override fun visitJumpInsnWithProbe(opcode: Int, label: Label?, probeId: Int, frame: IFrame) {
-        if (opcode == Opcodes.GOTO) {
-            probeInserter.insertProbe(probeId)
-            mv.visitJumpInsn(Opcodes.GOTO, label)
-        } else {
-            val intermediate = Label()
-            mv.visitJumpInsn(getInverted(opcode), intermediate)
-            probeInserter.insertProbe(probeId)
-            mv.visitJumpInsn(Opcodes.GOTO, label)
-            mv.visitLabel(intermediate)
-            frame.accept(mv)
-        }
-    }
-
-    private fun getInverted(opcode: Int): Int =
-        when (opcode) {
-            Opcodes.IFEQ -> Opcodes.IFNE
-            Opcodes.IFNE -> Opcodes.IFEQ
-            Opcodes.IFLT -> Opcodes.IFGE
-            Opcodes.IFGE -> Opcodes.IFLT
-            Opcodes.IFGT -> Opcodes.IFLE
-            Opcodes.IFLE -> Opcodes.IFGT
-            Opcodes.IF_ICMPEQ -> Opcodes.IF_ICMPNE
-            Opcodes.IF_ICMPNE -> Opcodes.IF_ICMPEQ
-            Opcodes.IF_ICMPLT -> Opcodes.IF_ICMPGE
-            Opcodes.IF_ICMPGE -> Opcodes.IF_ICMPLT
-            Opcodes.IF_ICMPGT -> Opcodes.IF_ICMPLE
-            Opcodes.IF_ICMPLE -> Opcodes.IF_ICMPGT
-            Opcodes.IF_ACMPEQ -> Opcodes.IF_ACMPNE
-            Opcodes.IF_ACMPNE -> Opcodes.IF_ACMPEQ
-            Opcodes.IFNULL -> Opcodes.IFNONNULL
-            Opcodes.IFNONNULL -> Opcodes.IFNULL
-            else -> throw IllegalArgumentException()
-        }
-
-    override fun visitTableSwitchInsnWithProbes(
-        min: Int,
-        max: Int,
-        dflt: Label,
-        labels: Array<Label>,
-        frame: IFrame
-    ) {
-        LabelInfo.resetDone(dflt)
-        LabelInfo.resetDone(labels)
-        val newDflt: Label = createIntermediate(dflt)
-        val newLabels: Array<Label?> = createIntermediates(labels)
-        mv.visitTableSwitchInsn(min, max, newDflt, *newLabels)
-
-        insertIntermediateProbes(dflt, labels, frame)
-    }
-
-    override fun visitLookupSwitchInsnWithProbes(
-        dflt: Label,
-        keys: IntArray,
-        labels: Array<Label>,
-        frame: IFrame
-    ) {
-        LabelInfo.resetDone(dflt)
-        LabelInfo.resetDone(labels)
-        val newDflt: Label = createIntermediate(dflt)
-        val newLabels: Array<Label?> = createIntermediates(labels)
-        mv.visitLookupSwitchInsn(newDflt, keys, newLabels)
-
-        insertIntermediateProbes(dflt, labels, frame)
-    }
-
-    private fun createIntermediates(labels: Array<Label>): Array<Label?> {
-        val intermediates = arrayOfNulls<Label>(labels.size)
-        for (i in labels.indices) {
-            intermediates[i] = createIntermediate(labels[i])
-        }
-        return intermediates
-    }
-
-    private fun createIntermediate(label: Label): Label {
-        val intermediate: Label
-        if (LabelInfo.getProbeId(label) == LabelInfo.NO_PROBE) {
-            intermediate = label
-        } else {
-            if (LabelInfo.isDone(label)) {
-                intermediate = LabelInfo.getIntermediateLabel(label)
-            } else {
-                intermediate = Label()
-                LabelInfo.setIntermediateLabel(label, intermediate)
-                LabelInfo.setDone(label)
-            }
-        }
-        return intermediate
-    }
-
-    private fun insertIntermediateProbe(label: Label, frame: IFrame) {
-        val probeId = LabelInfo.getProbeId(label)
-        if (probeId != LabelInfo.NO_PROBE && !LabelInfo.isDone(label)) {
-            mv.visitLabel(LabelInfo.getIntermediateLabel(label))
-            frame.accept(mv)
-            probeInserter.insertProbe(probeId)
-            mv.visitJumpInsn(Opcodes.GOTO, label)
-            LabelInfo.setDone(label)
-        }
-    }
-
-    private fun insertIntermediateProbes(dflt: Label, labels: Array<Label>, frame: IFrame) {
-        LabelInfo.resetDone(dflt)
-        LabelInfo.resetDone(labels)
-        insertIntermediateProbe(dflt, frame)
-        for (l in labels) {
-            insertIntermediateProbe(l, frame)
-        }
-    }
-
-}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodToProbesVisitor.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodToProbesVisitor.kt
new file mode 100644
index 0000000000..cb4dddeafd
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodToProbesVisitor.kt
@@ -0,0 +1,17 @@
+package org.jacoco.core.internal.instr
+
+import org.jacoco.core.internal.flow.ClassProbesAdapter
+import org.objectweb.asm.ClassVisitor
+
+class MyMethodToProbesVisitor(
+    api: Int,
+    strategy: IProbeArrayStrategy,
+    writer: ClassVisitor,
+    val methodToProbes: MutableMap<String, MutableList<Int>> = mutableMapOf()
+) : ClassVisitor(
+    api,
+    ClassProbesAdapter(
+        NoneClassInstrumenter(strategy, writer, methodToProbes),
+        true
+    )
+)
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyProbeInserter.kt
deleted file mode 100644
index ea2b07e862..0000000000
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyProbeInserter.kt
+++ /dev/null
@@ -1,18 +0,0 @@
-package org.jacoco.core.internal.instr
-
-import org.objectweb.asm.MethodVisitor
-
-class MyProbeInserter(
-    private val name: String,
-    mv: MethodVisitor,
-) : IProbeInserter, MethodVisitor(InstrSupport.ASM_API_VERSION, mv) {
-
-    companion object {
-        val methodToProbes = mutableMapOf<String, MutableList<Int>>()
-    }
-
-    override fun insertProbe(id: Int) {
-        methodToProbes.getOrPut(name) { mutableListOf() }.add(id)
-    }
-
-}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
similarity index 65%
rename from utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassInstrumenter.kt
rename to utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
index 0a4a9d0f2b..ebc263bca9 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
@@ -3,9 +3,10 @@ package org.jacoco.core.internal.instr
 import org.jacoco.core.internal.flow.MethodProbesVisitor
 import org.objectweb.asm.ClassVisitor
 
-class MyClassInstrumenter(
+class NoneClassInstrumenter(
     probeArrayStrategy: IProbeArrayStrategy,
-    cv: ClassVisitor
+    cv: ClassVisitor,
+    private val methodToProbes: MutableMap<String, MutableList<Int>>,
 ) : ClassInstrumenter(probeArrayStrategy, cv) {
 
     override fun visitMethod(
@@ -18,8 +19,8 @@ class MyClassInstrumenter(
         val mv = cv.visitMethod(access, name, desc, signature, exceptions)
 
         val frameEliminator = DuplicateFrameEliminator(mv)
-        val probeVariableInserter = MyProbeInserter(name, frameEliminator)
-        return MyMethodInstrumenter(probeVariableInserter, probeVariableInserter)
+        val probeVariableInserter = NoneProbeInserter(name, methodToProbes, frameEliminator)
+        return MethodInstrumenter(probeVariableInserter, probeVariableInserter)
     }
 
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
new file mode 100644
index 0000000000..8529306167
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
@@ -0,0 +1,19 @@
+package org.jacoco.core.internal.instr
+
+import org.objectweb.asm.MethodVisitor
+import org.utbot.instrumentation.Settings
+
+/**
+ * The probe inserter does not emit any code at all. This is used to collect method probes.
+ */
+class NoneProbeInserter(
+    private val methodName: String,
+    private val methodToProbes: MutableMap<String, MutableList<Int>>,
+    mv: MethodVisitor,
+) : IProbeInserter, MethodVisitor(Settings.ASM_API, mv) {
+
+    override fun insertProbe(id: Int) {
+        methodToProbes.getOrPut(methodName) { mutableListOf() }.add(id)
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
index 30122de90f..037b32e1a3 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
@@ -1,40 +1,20 @@
 package org.utbot.instrumentation.instrumentation.coverage
 
+import org.jacoco.core.internal.flow.ClassProbesAdapter
 import org.jacoco.core.internal.instr.createJacocoClassVisitorForBytecodeInstrumentation
 import org.jacoco.core.internal.instr.sw
 import org.objectweb.asm.ClassReader
-import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
 import org.utbot.common.withAccessibility
 import org.utbot.instrumentation.Settings
-import org.utbot.instrumentation.instrumentation.ArgumentList
 import org.utbot.instrumentation.instrumentation.Instrumentation
-import org.utbot.instrumentation.instrumentation.InvokeInstrumentation
 import org.utbot.instrumentation.instrumentation.instrumenter.ClassVisitorBuilder
 import org.utbot.instrumentation.instrumentation.instrumenter.Instrumenter
-import org.utbot.instrumentation.util.NoProbesArrayException
 import java.security.ProtectionDomain
 
 class BranchCoverageInstrumentation : CoverageInstrumentation() {
 
-    private val invokeInstrumentation = InvokeInstrumentation()
-
-    override fun invoke(
-        clazz: Class<*>,
-        methodSignature: String,
-        arguments: ArgumentList,
-        parameters: Any?
-    ): Result<*> {
-        val probesFieldName: String = Settings.PROBES_ARRAY_NAME
-        val visitedLinesField = clazz.fields.firstOrNull { it.name == probesFieldName }
-            ?: throw NoProbesArrayException(clazz, Settings.PROBES_ARRAY_NAME)
-
-        return visitedLinesField.withAccessibility {
-            invokeInstrumentation.invoke(clazz, methodSignature, arguments, parameters)
-        }
-    }
-
-    override fun <T : Any> methodToCollectCoverage(clazz: Class<out T>): Map<String, IntRange> {
+    override fun <T : Any> computeMapOfRanges(clazz: Class<out T>): Map<String, IntRange> {
         return Instrumenter(clazz).computeMapOfRangesForBranchCoverage()
     }
 
@@ -47,17 +27,26 @@ class BranchCoverageInstrumentation : CoverageInstrumentation() {
     ): ByteArray {
         val instrumenter = Instrumenter(classfileBuffer)
 
-        instrumenter.visitClass(object : ClassVisitorBuilder<ClassVisitor> {
+        val cv = instrumenter.visitClass(object : ClassVisitorBuilder<ClassProbesAdapter> {
             override val writerFlags: Int
                 get() = 0
 
             override val readerParsingOptions: Int
                 get() = ClassReader.EXPAND_FRAMES
 
-            override fun build(writer: ClassWriter): ClassVisitor =
+            override fun build(writer: ClassWriter): ClassProbesAdapter =
                 createJacocoClassVisitorForBytecodeInstrumentation(writer, className)
         })
 
+        // TODO exceptions
+        val probesCountField = ClassProbesAdapter::class.java.declaredFields.firstOrNull { it.name == "counter" }
+            ?: throw IllegalStateException()
+        val probeCount = probesCountField.withAccessibility {
+            this.get(cv) as? Int ?: throw ClassCastException()
+        }
+
+        instrumenter.addStaticField(StaticArrayFieldInitializer(className, Settings.PROBES_ARRAY_NAME, probeCount))
+
         val result = sw.toString()
         println(result)
 
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
index b9f6729345..5b63455ec4 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
@@ -62,7 +62,7 @@ open class CoverageInstrumentation : Instrumentation<Result<*>> {
             val visitedLines = visitedLinesField.get(null) as? BooleanArray
                 ?: throw CastProbesArrayException()
 
-            val methodToInstrRange = methodToCollectCoverage(clazz)
+            val methodToInstrRange = computeMapOfRanges(clazz)
 
             val res = CoverageInfo(methodToInstrRange, visitedLines.mapIndexed { idx, b ->
                 if (b) idx else null
@@ -74,7 +74,7 @@ open class CoverageInstrumentation : Instrumentation<Result<*>> {
         }
     }
 
-    protected open fun <T : Any> methodToCollectCoverage(clazz: Class<out T>): Map<String, IntRange> {
+    protected open fun <T : Any> computeMapOfRanges(clazz: Class<out T>): Map<String, IntRange> {
         return Instrumenter(clazz).computeMapOfRanges()
     }
 
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/StaticArrayFieldInitializer.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/StaticArrayFieldInitializer.kt
new file mode 100644
index 0000000000..cc4f9411ab
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/StaticArrayFieldInitializer.kt
@@ -0,0 +1,24 @@
+package org.utbot.instrumentation.instrumentation.coverage
+
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+import org.utbot.instrumentation.Settings
+import org.utbot.instrumentation.instrumentation.instrumenter.visitors.util.StaticFieldInitializer
+
+class StaticArrayFieldInitializer(
+    private val className: String,
+    override val name: String,
+    private val probesCount: Int
+) : StaticFieldInitializer {
+
+    override val descriptor: String = Settings.PROBES_ARRAY_DESC
+    override val signature: String? = null
+
+    override fun initField(mv: MethodVisitor): MethodVisitor {
+        mv.visitLdcInsn(probesCount)
+        mv.visitIntInsn(Opcodes.NEWARRAY, Opcodes.T_BOOLEAN)
+        mv.visitFieldInsn(Opcodes.PUTSTATIC, className, name, Settings.PROBES_ARRAY_DESC)
+        return mv
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
index fcde31717d..2020d96fba 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
@@ -5,11 +5,13 @@ import org.objectweb.asm.ClassReader
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
 import org.objectweb.asm.Opcodes
+import org.objectweb.asm.util.TraceClassVisitor
 import org.utbot.instrumentation.instrumentation.instrumenter.visitors.MethodToProbesVisitor
 import org.utbot.instrumentation.instrumentation.instrumenter.visitors.util.*
 import org.utbot.instrumentation.process.HandlerClassesLoader
 import java.io.IOException
 import java.io.InputStream
+import java.io.PrintWriter
 
 
 // TODO: handle with flags EXPAND_FRAMES, etc.
@@ -51,19 +53,18 @@ class Instrumenter(classByteCode: ByteArray, val classLoader: ClassLoader? = nul
     }
 
     fun computeMapOfRangesForBranchCoverage(): Map<String, IntRange> {
-
-        visitClass(object : ClassVisitorBuilder<ClassVisitor> {
+        val methodToListOfProbesInserter = visitClass(object : ClassVisitorBuilder<MyMethodToProbesVisitor> {
             override val writerFlags: Int
                 get() = 0
 
             override val readerParsingOptions: Int
                 get() = ClassReader.EXPAND_FRAMES
 
-            override fun build(writer: ClassWriter): ClassVisitor =
+            override fun build(writer: ClassWriter): MyMethodToProbesVisitor =
                 createJacocoClassVisitorForCollectionBranchCoverage(writer)
         })
 
-        return MyProbeInserter.methodToProbes.mapValues { (_, probes) -> (probes.first()..probes.last()) }
+        return methodToListOfProbesInserter.methodToProbes.mapValues { (_, probes) -> (probes.first()..probes.last()) }
     }
 
     fun addField(instanceFieldInitializer: InstanceFieldInitializer) {
@@ -71,7 +72,10 @@ class Instrumenter(classByteCode: ByteArray, val classLoader: ClassLoader? = nul
     }
 
     fun addStaticField(staticFieldInitializer: StaticFieldInitializer) {
-        visitClass { writer -> AddStaticFieldAdapter(writer, staticFieldInitializer) }
+        visitClass { writer ->
+            val tcv = TraceClassVisitor(writer, PrintWriter(sw))
+            AddStaticFieldAdapter(tcv, staticFieldInitializer)
+        }
     }
 
     fun visitInstructions(instructionVisitor: IInstructionVisitor, methodName: String? = null) {

From 032f8de7744f819115e62faa1184e57d6730fbcd Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Fri, 4 Aug 2023 16:33:13 +0300
Subject: [PATCH 03/13] Add working version of branch trace instrumentation

---
 .../MethodProbesCollector.kt}                 |   7 +-
 .../internal/flow/TraceClassProbesAdapter.kt  |  61 +++++
 .../internal/flow/TraceMethodProbesAdapter.kt | 253 ++++++++++++++++++
 .../org/jacoco/core/internal/instr/Api.kt     |  34 ++-
 ...egy.kt => ClassFieldProbeArrayStrategy.kt} |   2 +-
 .../core/internal/instr/NoneProbeInserter.kt  |   2 +-
 .../internal/instr/TraceClassInstrumenter.kt  |  27 ++
 .../core/internal/instr/TraceProbeInserter.kt |  40 +++
 .../core/internal/instr/TraceStrategy.kt      |  30 +++
 .../coverage/BranchCoverageInstrumentation.kt |   4 +-
 .../et/ExecutionBranchTraceInstrumentation.kt |  35 +++
 .../et/ExecutionTraceInstrumentation.kt       |   6 +-
 .../instrumenter/Instrumenter.kt              |  13 +-
 .../util/InstructionVisitorMethodAdapter.kt   |   2 +-
 14 files changed, 488 insertions(+), 28 deletions(-)
 rename utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/{instr/MyMethodToProbesVisitor.kt => flow/MethodProbesCollector.kt} (64%)
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
 rename utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/{MyClassFieldProbeArrayStrategy.kt => ClassFieldProbeArrayStrategy.kt} (85%)
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceStrategy.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionBranchTraceInstrumentation.kt

diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodToProbesVisitor.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/MethodProbesCollector.kt
similarity index 64%
rename from utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodToProbesVisitor.kt
rename to utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/MethodProbesCollector.kt
index cb4dddeafd..c9b06b52b5 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyMethodToProbesVisitor.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/MethodProbesCollector.kt
@@ -1,9 +1,10 @@
-package org.jacoco.core.internal.instr
+package org.jacoco.core.internal.flow
 
-import org.jacoco.core.internal.flow.ClassProbesAdapter
+import org.jacoco.core.internal.instr.IProbeArrayStrategy
+import org.jacoco.core.internal.instr.NoneClassInstrumenter
 import org.objectweb.asm.ClassVisitor
 
-class MyMethodToProbesVisitor(
+class MethodProbesCollector(
     api: Int,
     strategy: IProbeArrayStrategy,
     writer: ClassVisitor,
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
new file mode 100644
index 0000000000..e314acf0a3
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
@@ -0,0 +1,61 @@
+package org.jacoco.core.internal.flow
+
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.commons.AnalyzerAdapter
+import org.utbot.instrumentation.Settings
+import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
+
+class TraceClassProbesAdapter(
+    val cv: ClassProbesVisitor,
+    private val trackFrames: Boolean,
+    private val className: String,
+    private val storage: ProcessingStorage
+) : IProbeIdGenerator, ClassVisitor(Settings.ASM_API, cv) {
+
+    private val EMPTY_METHOD_PROBES_VISITOR: MethodProbesVisitor = object : MethodProbesVisitor() {}
+    private var counter: Int = 0
+
+    override fun visitMethod(
+        access: Int, name: String,
+        desc: String, signature: String?,
+        exceptions: Array<String?>?
+    ): MethodVisitor {
+        val methodProbes: MethodProbesVisitor = (cv as ClassProbesVisitor).visitMethod(
+            access, name, desc,
+            signature, exceptions
+        ) ?: EMPTY_METHOD_PROBES_VISITOR
+
+        return object : MethodSanitizer(
+            null, access, name, desc, signature,
+            exceptions
+        ) {
+            override fun visitEnd() {
+                super.visitEnd()
+                LabelFlowAnalyzer.markLabels(this)
+                val probesAdapter = TraceMethodProbesAdapter(
+                    methodProbes, this@TraceClassProbesAdapter,
+                    storage, className, name, desc
+                )
+                if (trackFrames) {
+                    val analyzer = AnalyzerAdapter(
+                        this@TraceClassProbesAdapter.className, access, name, desc,
+                        probesAdapter
+                    )
+                    probesAdapter.setAnalyzer(analyzer)
+                    methodProbes.accept(this, analyzer)
+                } else {
+                    methodProbes.accept(this, probesAdapter)
+                }
+            }
+        }
+    }
+
+    override fun visitEnd() {
+        (cv as ClassProbesVisitor).visitTotalProbeCount(counter)
+        super.visitEnd()
+    }
+
+    override fun nextId(): Int = counter++
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
new file mode 100644
index 0000000000..eedf84c924
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
@@ -0,0 +1,253 @@
+package org.jacoco.core.internal.flow
+
+import org.objectweb.asm.Label
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+import org.objectweb.asm.commons.AnalyzerAdapter
+import org.utbot.instrumentation.Settings
+import org.utbot.instrumentation.instrumentation.et.CommonInstruction
+import org.utbot.instrumentation.instrumentation.et.ExplicitThrowInstruction
+import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
+import org.utbot.instrumentation.instrumentation.et.ReturnInstruction
+import org.utbot.instrumentation.instrumentation.instrumenter.visitors.util.InstructionVisitorMethodAdapter.Companion.returnInsns
+
+class TraceMethodProbesAdapter(
+    private val probesVisitor: MethodProbesVisitor,
+    private val idGenerator: IProbeIdGenerator,
+    private val storage: ProcessingStorage,
+    private val className: String,
+    methodName: String,
+    descriptor: String,
+) : MethodVisitor(Settings.ASM_API, probesVisitor) {
+
+    private var analyzer: AnalyzerAdapter? = null
+    private val tryCatchProbeLabels: HashMap<Label, Label> = hashMapOf()
+    private var currentLineNumber: Int = 0
+    private val currentMethodSignature: String = methodName + descriptor
+
+    fun setAnalyzer(analyzer: AnalyzerAdapter) {
+        this.analyzer = analyzer
+    }
+
+    override fun visitTryCatchBlock(start: Label, end: Label, handler: Label, type: String) {
+        probesVisitor.visitTryCatchBlock(getTryCatchLabel(start), getTryCatchLabel(end), handler, type)
+    }
+
+    private fun getTryCatchLabel(l: Label): Label {
+        var label: Label = l
+        if (tryCatchProbeLabels.containsKey(label)) {
+            label = tryCatchProbeLabels[label]!!
+        } else if (LabelInfo.needsProbe(label)) {
+            // If a probe will be inserted before the label, we'll need to use a
+            // different label to define the range of the try-catch block.
+            val probeLabel = Label()
+            LabelInfo.setSuccessor(probeLabel)
+            tryCatchProbeLabels[label] = probeLabel
+            label = probeLabel
+        }
+        return label
+    }
+
+    override fun visitLabel(label: Label?) {
+        if (LabelInfo.needsProbe(label)) {
+            if (tryCatchProbeLabels.containsKey(label)) {
+                probesVisitor.visitLabel(tryCatchProbeLabels[label])
+            }
+            val localId = idGenerator.nextId()
+            val id = storage.computeId(className, localId)
+            val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
+            storage.addInstruction(id, instruction)
+            probesVisitor.visitProbe(localId)
+        }
+        probesVisitor.visitLabel(label)
+    }
+
+    override fun visitLineNumber(line: Int, start: Label?) {
+        currentLineNumber = line
+        super.visitLineNumber(line, start)
+    }
+
+    override fun visitInsn(opcode: Int) {
+        when (opcode) {
+            in returnInsns -> {
+                val localId = idGenerator.nextId()
+                val id = storage.computeId(className, localId)
+                val instruction = ReturnInstruction(currentLineNumber, currentMethodSignature)
+                storage.addInstruction(id, instruction)
+                probesVisitor.visitInsnWithProbe(
+                    opcode,
+                    localId
+                )
+            }
+
+            Opcodes.ATHROW -> {
+                val localId = idGenerator.nextId()
+                val id = storage.computeId(className, localId)
+                val instruction = ExplicitThrowInstruction(currentLineNumber, currentMethodSignature)
+                storage.addInstruction(id, instruction)
+                probesVisitor.visitInsnWithProbe(
+                    opcode,
+                    localId
+                )
+            }
+
+            else -> probesVisitor.visitInsn(opcode)
+        }
+    }
+
+    override fun visitJumpInsn(opcode: Int, label: Label?) {
+        if (LabelInfo.isMultiTarget(label)) {
+            val localId = idGenerator.nextId()
+            val id = storage.computeId(className, localId)
+            val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
+            storage.addInstruction(id, instruction)
+            probesVisitor.visitJumpInsnWithProbe(opcode, label, localId, frame(jumpPopCount(opcode)))
+        } else {
+            probesVisitor.visitJumpInsn(opcode, label)
+        }
+    }
+
+    private fun jumpPopCount(opcode: Int): Int {
+        return when (opcode) {
+            Opcodes.GOTO -> 0
+
+            Opcodes.IFEQ,
+            Opcodes.IFNE,
+            Opcodes.IFLT,
+            Opcodes.IFGE,
+            Opcodes.IFGT,
+            Opcodes.IFLE,
+            Opcodes.IFNULL,
+            Opcodes.IFNONNULL -> 1
+
+            else -> 2
+        }
+    }
+
+    override fun visitLookupSwitchInsn(dflt: Label, keys: IntArray, labels: Array<Label>) {
+        if (markLabels(dflt, labels)) {
+            visitLookupSwitchInsnWithProbes(
+                dflt, keys, labels,
+                frame(1)
+            )
+        } else {
+            probesVisitor.visitLookupSwitchInsn(dflt, keys, labels)
+        }
+    }
+
+    override fun visitTableSwitchInsn(min: Int, max: Int, dflt: Label, vararg labels: Label) {
+        if (markLabels(dflt, labels)) {
+            visitTableSwitchInsnWithProbes(
+                min, max, dflt, labels,
+                frame(1)
+            )
+        } else {
+            probesVisitor.visitTableSwitchInsn(min, max, dflt, *labels)
+        }
+    }
+
+    private fun markLabels(dflt: Label, labels: Array<out Label>): Boolean {
+        var probe = false
+        LabelInfo.resetDone(labels)
+        if (LabelInfo.isMultiTarget(dflt)) {
+            LabelInfo.setProbeId(dflt, idGenerator.nextId())
+            probe = true
+        }
+        LabelInfo.setDone(dflt)
+        for (l in labels) {
+            if (LabelInfo.isMultiTarget(l) && !LabelInfo.isDone(l)) {
+                LabelInfo.setProbeId(l, idGenerator.nextId())
+                probe = true
+            }
+            LabelInfo.setDone(l)
+        }
+        return probe
+    }
+
+    private fun frame(popCount: Int): IFrame {
+        return FrameSnapshot.create(analyzer, popCount)
+    }
+
+    // === MethodProbesVisitor ===
+
+    private fun visitLookupSwitchInsnWithProbes(dflt: Label, keys: IntArray?, labels: Array<Label>, frame: IFrame) {
+        // 1. Calculate intermediate labels:
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        val newDflt = createIntermediate(dflt)
+        val newLabels = createIntermediates(labels)
+        probesVisitor.visitLookupSwitchInsn(newDflt, keys, newLabels)
+
+        // 2. Insert probes:
+        insertIntermediateProbes(dflt, labels, frame)
+    }
+
+    private fun visitTableSwitchInsnWithProbes(
+        min: Int,
+        max: Int,
+        dflt: Label,
+        labels: Array<out Label>,
+        frame: IFrame
+    ) {
+        // 1. Calculate intermediate labels:
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        val newDflt = createIntermediate(dflt)
+        val newLabels = createIntermediates(labels)
+        probesVisitor.visitTableSwitchInsn(min, max, newDflt, *newLabels)
+
+        // 2. Insert probes:
+        insertIntermediateProbes(dflt, labels, frame)
+    }
+
+
+    // === MethodInstrumenter ===
+
+    private fun createIntermediates(labels: Array<out Label>): Array<Label?> {
+        val intermediates = arrayOfNulls<Label>(labels.size)
+        for (i in labels.indices) {
+            intermediates[i] = createIntermediate(labels[i])
+        }
+        return intermediates
+    }
+
+    private fun createIntermediate(label: Label): Label {
+        val intermediate: Label
+        if (LabelInfo.getProbeId(label) == LabelInfo.NO_PROBE) {
+            intermediate = label
+        } else {
+            if (LabelInfo.isDone(label)) {
+                intermediate = LabelInfo.getIntermediateLabel(label)
+            } else {
+                intermediate = Label()
+                LabelInfo.setIntermediateLabel(label, intermediate)
+                LabelInfo.setDone(label)
+            }
+        }
+        return intermediate
+    }
+
+    private fun insertIntermediateProbe(label: Label, frame: IFrame) {
+        val localId = LabelInfo.getProbeId(label)
+        if (localId != LabelInfo.NO_PROBE && !LabelInfo.isDone(label)) {
+            probesVisitor.visitLabel(LabelInfo.getIntermediateLabel(label))
+            frame.accept(probesVisitor)
+            val id = storage.computeId(className, localId)
+            val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
+            storage.addInstruction(id, instruction)
+            probesVisitor.visitProbe(localId)
+            probesVisitor.visitJumpInsn(Opcodes.GOTO, label)
+            LabelInfo.setDone(label)
+        }
+    }
+
+    private fun insertIntermediateProbes(dflt: Label, labels: Array<out Label>, frame: IFrame) {
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        insertIntermediateProbe(dflt, frame)
+        for (l in labels) {
+            insertIntermediateProbe(l, frame)
+        }
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
index 52d7cc0fff..aff0e5b88a 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -1,24 +1,44 @@
 package org.jacoco.core.internal.instr
 
 import org.jacoco.core.internal.flow.ClassProbesAdapter
+import org.jacoco.core.internal.flow.MethodProbesCollector
+import org.jacoco.core.internal.flow.TraceClassProbesAdapter
 import org.objectweb.asm.ClassWriter
 import org.objectweb.asm.util.TraceClassVisitor
 import org.utbot.instrumentation.Settings
+import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
 import java.io.PrintWriter
 import java.io.StringWriter
 
 val sw: StringWriter = StringWriter()
 
-fun createJacocoClassVisitorForCollectionBranchCoverage(writer: ClassWriter): MyMethodToProbesVisitor {
+fun createClassVisitorForBranchCoverageInstrumentation(writer: ClassWriter, className: String): ClassProbesAdapter {
+    val strategy = ClassFieldProbeArrayStrategy(className)
+    return ClassProbesAdapter(
+        ClassInstrumenter(strategy, writer),
+        true
+    )
+}
+
+fun createClassVisitorForComputeMapOfRangesForBranchCoverage(writer: ClassWriter): MethodProbesCollector {
     val strategy = NoneProbeArrayStrategy()
     val tcv = TraceClassVisitor(writer, PrintWriter(sw))
-    return MyMethodToProbesVisitor(Settings.ASM_API, strategy, tcv)
+    return MethodProbesCollector(Settings.ASM_API, strategy, tcv)
 }
 
-fun createJacocoClassVisitorForBytecodeInstrumentation(writer: ClassWriter, className: String): ClassProbesAdapter {
-    val strategy = MyClassFieldProbeArrayStrategy(className)
-    return ClassProbesAdapter(
-        ClassInstrumenter(strategy, writer),
-        true
+fun createClassVisitorForTracingBranchInstructions(
+    className: String,
+    storage: ProcessingStorage,
+    writer: ClassWriter
+): TraceClassProbesAdapter {
+    val strategy = TraceStrategy()
+    val tcv = TraceClassVisitor(writer, PrintWriter(sw))
+    return TraceClassProbesAdapter(
+        TraceClassInstrumenter(strategy, tcv) { id ->
+            storage.computeId(className, id)
+        },
+        true,
+        className,
+        storage
     )
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt
similarity index 85%
rename from utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt
rename to utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt
index ba420b0588..639a010739 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/MyClassFieldProbeArrayStrategy.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt
@@ -5,7 +5,7 @@ import org.objectweb.asm.MethodVisitor
 import org.objectweb.asm.Opcodes
 import org.utbot.instrumentation.Settings
 
-class MyClassFieldProbeArrayStrategy(private val className: String) : IProbeArrayStrategy {
+class ClassFieldProbeArrayStrategy(private val className: String) : IProbeArrayStrategy {
 
     override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
         mv.visitFieldInsn(Opcodes.GETSTATIC, className, Settings.PROBES_ARRAY_NAME, Settings.PROBES_ARRAY_DESC)
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
index 8529306167..67777762ac 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
@@ -4,7 +4,7 @@ import org.objectweb.asm.MethodVisitor
 import org.utbot.instrumentation.Settings
 
 /**
- * The probe inserter does not emit any code at all. This is used to collect method probes.
+ * This inserter does not emit any code at all. This is used to collect method probes.
  */
 class NoneProbeInserter(
     private val methodName: String,
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
new file mode 100644
index 0000000000..04b39dd76e
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
@@ -0,0 +1,27 @@
+package org.jacoco.core.internal.instr
+
+import org.jacoco.core.internal.flow.MethodProbesVisitor
+import org.objectweb.asm.ClassVisitor
+
+class TraceClassInstrumenter(
+    private val probeArrayStrategy: IProbeArrayStrategy,
+    cv: ClassVisitor,
+    private val nextIdGenerator: (id: Int) -> Long
+) : ClassInstrumenter(probeArrayStrategy, cv) {
+
+    override fun visitMethod(
+        access: Int,
+        name: String,
+        desc: String,
+        signature: String?,
+        exceptions: Array<out String>?
+    ): MethodProbesVisitor {
+        val mv = cv.visitMethod(access, name, desc, signature, exceptions)
+
+        val frameEliminator = DuplicateFrameEliminator(mv)
+        val probeVariableInserter = TraceProbeInserter(access,
+            name, desc, frameEliminator, probeArrayStrategy, nextIdGenerator)
+        return MethodInstrumenter(probeVariableInserter, probeVariableInserter)
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
new file mode 100644
index 0000000000..0d2ae55e18
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
@@ -0,0 +1,40 @@
+package org.jacoco.core.internal.instr
+
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+import org.objectweb.asm.Type
+import org.utbot.instrumentation.instrumentation.et.RuntimeTraceStorage
+import kotlin.reflect.jvm.javaMethod
+
+internal class TraceProbeInserter(
+    access: Int,
+    name: String,
+    desc: String,
+    mv: MethodVisitor,
+    arrayStrategy: IProbeArrayStrategy,
+    private val nextIdGenerator: (id: Int) -> Long,
+) : ProbeInserter(access, name, desc, mv, arrayStrategy) {
+
+    private val internalName = Type.getInternalName(RuntimeTraceStorage::class.java)
+    private val visitMethodDescriptor = Type.getMethodDescriptor(RuntimeTraceStorage::visit.javaMethod)
+
+    private val clinit: Boolean
+    private val variable: Int
+
+    init {
+        clinit = InstrSupport.CLINIT_NAME == name
+        var pos = if (Opcodes.ACC_STATIC and access == 0) 1 else 0
+        for (t in Type.getArgumentTypes(desc)) {
+            pos += t.size
+        }
+        variable = pos
+    }
+
+    override fun insertProbe(localId: Int) {
+        mv.visitVarInsn(Opcodes.ILOAD, variable)
+        val id = nextIdGenerator.invoke(localId)
+        mv.visitLdcInsn(id)
+        mv.visitMethodInsn(Opcodes.INVOKESTATIC, internalName, "visit", visitMethodDescriptor, false)
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceStrategy.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceStrategy.kt
new file mode 100644
index 0000000000..6abe249292
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceStrategy.kt
@@ -0,0 +1,30 @@
+package org.jacoco.core.internal.instr
+
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+import org.objectweb.asm.Type
+import org.utbot.instrumentation.instrumentation.et.RuntimeTraceStorage
+import kotlin.reflect.jvm.javaField
+
+class TraceStrategy : IProbeArrayStrategy {
+
+    private val internalName = Type.getInternalName(RuntimeTraceStorage::class.java)
+    private val counterCallIdName = RuntimeTraceStorage::`$__counter_call_id__`.javaField!!.name
+
+    override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
+        // int variable = RuntimeTraceStorage.$__counter_call_id__++
+        mv.visitFieldInsn(Opcodes.GETSTATIC, internalName, counterCallIdName, RuntimeTraceStorage.DESC_CALL_ID_COUNTER)
+        mv.visitInsn(Opcodes.ICONST_1)
+        mv.visitInsn(Opcodes.IADD)
+        mv.visitInsn(Opcodes.DUP)
+
+        mv.visitFieldInsn(Opcodes.PUTSTATIC, internalName, counterCallIdName, RuntimeTraceStorage.DESC_CALL_ID_COUNTER)
+        mv.visitVarInsn(Opcodes.ISTORE, variable)
+        return 2
+    }
+
+    override fun addMembers(cv: ClassVisitor?, probeCount: Int) {
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
index 037b32e1a3..4e6c0ef16f 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
@@ -1,7 +1,7 @@
 package org.utbot.instrumentation.instrumentation.coverage
 
 import org.jacoco.core.internal.flow.ClassProbesAdapter
-import org.jacoco.core.internal.instr.createJacocoClassVisitorForBytecodeInstrumentation
+import org.jacoco.core.internal.instr.createClassVisitorForBranchCoverageInstrumentation
 import org.jacoco.core.internal.instr.sw
 import org.objectweb.asm.ClassReader
 import org.objectweb.asm.ClassWriter
@@ -35,7 +35,7 @@ class BranchCoverageInstrumentation : CoverageInstrumentation() {
                 get() = ClassReader.EXPAND_FRAMES
 
             override fun build(writer: ClassWriter): ClassProbesAdapter =
-                createJacocoClassVisitorForBytecodeInstrumentation(writer, className)
+                createClassVisitorForBranchCoverageInstrumentation(writer, className)
         })
 
         // TODO exceptions
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionBranchTraceInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionBranchTraceInstrumentation.kt
new file mode 100644
index 0000000000..0761024037
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionBranchTraceInstrumentation.kt
@@ -0,0 +1,35 @@
+package org.utbot.instrumentation.instrumentation.et
+
+import org.jacoco.core.internal.flow.TraceClassProbesAdapter
+import org.jacoco.core.internal.instr.createClassVisitorForTracingBranchInstructions
+import org.objectweb.asm.ClassWriter
+import org.utbot.instrumentation.instrumentation.Instrumentation
+import org.utbot.instrumentation.instrumentation.instrumenter.ClassVisitorBuilder
+import org.utbot.instrumentation.instrumentation.instrumenter.Instrumenter
+import java.security.ProtectionDomain
+
+class ExecutionBranchTraceInstrumentation : ExecutionTraceInstrumentation() {
+
+    override fun transform(
+        loader: ClassLoader?,
+        className: String,
+        classBeingRedefined: Class<*>?,
+        protectionDomain: ProtectionDomain?,
+        classfileBuffer: ByteArray
+    ): ByteArray {
+        val instrumenter = Instrumenter(classfileBuffer)
+
+        traceHandler.registerClass(className)
+        instrumenter.visitClass(object : ClassVisitorBuilder<TraceClassProbesAdapter> {
+
+            override fun build(writer: ClassWriter): TraceClassProbesAdapter =
+                createClassVisitorForTracingBranchInstructions(className, traceHandler.processingStorage, writer)
+        })
+
+        return instrumenter.classByteCode
+    }
+
+    object Factory : Instrumentation.Factory<Trace, ExecutionBranchTraceInstrumentation> {
+        override fun create(): ExecutionBranchTraceInstrumentation = ExecutionBranchTraceInstrumentation()
+    }
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionTraceInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionTraceInstrumentation.kt
index f62eff9b47..2b139496e6 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionTraceInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionTraceInstrumentation.kt
@@ -1,19 +1,19 @@
 package org.utbot.instrumentation.instrumentation.et
 
+import org.utbot.framework.plugin.api.FieldId
 import org.utbot.instrumentation.instrumentation.ArgumentList
 import org.utbot.instrumentation.instrumentation.Instrumentation
 import org.utbot.instrumentation.instrumentation.InvokeWithStaticsInstrumentation
 import org.utbot.instrumentation.instrumentation.instrumenter.Instrumenter
 import java.security.ProtectionDomain
-import org.utbot.framework.plugin.api.FieldId
 
 /**
  * This instrumentation allows to get execution trace during each call.
  */
 
-class ExecutionTraceInstrumentation : Instrumentation<Trace> {
+open class ExecutionTraceInstrumentation : Instrumentation<Trace> {
     private val invokeWithStatics = InvokeWithStaticsInstrumentation()
-    private val traceHandler = TraceHandler()
+    protected val traceHandler = TraceHandler()
 
     /**
      * Invokes a method with the given [methodSignature], the declaring class of which is [clazz], with the supplied
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
index 2020d96fba..d1d6caa3a4 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
@@ -53,16 +53,9 @@ class Instrumenter(classByteCode: ByteArray, val classLoader: ClassLoader? = nul
     }
 
     fun computeMapOfRangesForBranchCoverage(): Map<String, IntRange> {
-        val methodToListOfProbesInserter = visitClass(object : ClassVisitorBuilder<MyMethodToProbesVisitor> {
-            override val writerFlags: Int
-                get() = 0
-
-            override val readerParsingOptions: Int
-                get() = ClassReader.EXPAND_FRAMES
-
-            override fun build(writer: ClassWriter): MyMethodToProbesVisitor =
-                createJacocoClassVisitorForCollectionBranchCoverage(writer)
-        })
+        val methodToListOfProbesInserter = visitClass { writer ->
+            createClassVisitorForComputeMapOfRangesForBranchCoverage(writer)
+        }
 
         return methodToListOfProbesInserter.methodToProbes.mapValues { (_, probes) -> (probes.first()..probes.last()) }
     }
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/visitors/util/InstructionVisitorMethodAdapter.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/visitors/util/InstructionVisitorMethodAdapter.kt
index bfff5b8014..5bd120e93f 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/visitors/util/InstructionVisitorMethodAdapter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/visitors/util/InstructionVisitorMethodAdapter.kt
@@ -16,7 +16,7 @@ open class InstructionVisitorMethodAdapter(
     var lvs: LocalVariablesSorter? = null
 
     companion object {
-        private val returnInsns = setOf(
+        val returnInsns = setOf(
             Opcodes.IRETURN,
             Opcodes.LRETURN,
             Opcodes.FRETURN,

From f46d08e85e85bcadc8962efd849436542da9e147 Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Sat, 5 Aug 2023 15:39:37 +0300
Subject: [PATCH 04/13] Refactor

---
 .../core/internal/flow/TraceClassProbesAdapter.kt  | 13 +------------
 .../core/internal/flow/TraceMethodProbesAdapter.kt |  8 +-------
 .../kotlin/org/jacoco/core/internal/instr/Api.kt   | 11 ++++++-----
 .../internal/instr/ClassFieldProbeArrayStrategy.kt |  2 ++
 .../coverage/BranchCoverageInstrumentation.kt      | 14 ++++----------
 5 files changed, 14 insertions(+), 34 deletions(-)

diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
index e314acf0a3..7423502641 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
@@ -2,13 +2,11 @@ package org.jacoco.core.internal.flow
 
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.MethodVisitor
-import org.objectweb.asm.commons.AnalyzerAdapter
 import org.utbot.instrumentation.Settings
 import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
 
 class TraceClassProbesAdapter(
     val cv: ClassProbesVisitor,
-    private val trackFrames: Boolean,
     private val className: String,
     private val storage: ProcessingStorage
 ) : IProbeIdGenerator, ClassVisitor(Settings.ASM_API, cv) {
@@ -37,16 +35,7 @@ class TraceClassProbesAdapter(
                     methodProbes, this@TraceClassProbesAdapter,
                     storage, className, name, desc
                 )
-                if (trackFrames) {
-                    val analyzer = AnalyzerAdapter(
-                        this@TraceClassProbesAdapter.className, access, name, desc,
-                        probesAdapter
-                    )
-                    probesAdapter.setAnalyzer(analyzer)
-                    methodProbes.accept(this, analyzer)
-                } else {
-                    methodProbes.accept(this, probesAdapter)
-                }
+                methodProbes.accept(this, probesAdapter)
             }
         }
     }
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
index eedf84c924..0b1df6dfe7 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
@@ -3,7 +3,6 @@ package org.jacoco.core.internal.flow
 import org.objectweb.asm.Label
 import org.objectweb.asm.MethodVisitor
 import org.objectweb.asm.Opcodes
-import org.objectweb.asm.commons.AnalyzerAdapter
 import org.utbot.instrumentation.Settings
 import org.utbot.instrumentation.instrumentation.et.CommonInstruction
 import org.utbot.instrumentation.instrumentation.et.ExplicitThrowInstruction
@@ -20,15 +19,10 @@ class TraceMethodProbesAdapter(
     descriptor: String,
 ) : MethodVisitor(Settings.ASM_API, probesVisitor) {
 
-    private var analyzer: AnalyzerAdapter? = null
     private val tryCatchProbeLabels: HashMap<Label, Label> = hashMapOf()
     private var currentLineNumber: Int = 0
     private val currentMethodSignature: String = methodName + descriptor
 
-    fun setAnalyzer(analyzer: AnalyzerAdapter) {
-        this.analyzer = analyzer
-    }
-
     override fun visitTryCatchBlock(start: Label, end: Label, handler: Label, type: String) {
         probesVisitor.visitTryCatchBlock(getTryCatchLabel(start), getTryCatchLabel(end), handler, type)
     }
@@ -165,7 +159,7 @@ class TraceMethodProbesAdapter(
     }
 
     private fun frame(popCount: Int): IFrame {
-        return FrameSnapshot.create(analyzer, popCount)
+        return FrameSnapshot.create(null, popCount)
     }
 
     // === MethodProbesVisitor ===
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
index aff0e5b88a..cf177b3d0c 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -10,19 +10,21 @@ import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
 import java.io.PrintWriter
 import java.io.StringWriter
 
+// TODO: remove
 val sw: StringWriter = StringWriter()
 
 fun createClassVisitorForBranchCoverageInstrumentation(writer: ClassWriter, className: String): ClassProbesAdapter {
     val strategy = ClassFieldProbeArrayStrategy(className)
+    val tcv = TraceClassVisitor(writer, PrintWriter(sw)) // TODO: remove
     return ClassProbesAdapter(
-        ClassInstrumenter(strategy, writer),
-        true
+        ClassInstrumenter(strategy, tcv),
+        false
     )
 }
 
 fun createClassVisitorForComputeMapOfRangesForBranchCoverage(writer: ClassWriter): MethodProbesCollector {
     val strategy = NoneProbeArrayStrategy()
-    val tcv = TraceClassVisitor(writer, PrintWriter(sw))
+    val tcv = TraceClassVisitor(writer, PrintWriter(sw)) // TODO: remove
     return MethodProbesCollector(Settings.ASM_API, strategy, tcv)
 }
 
@@ -32,12 +34,11 @@ fun createClassVisitorForTracingBranchInstructions(
     writer: ClassWriter
 ): TraceClassProbesAdapter {
     val strategy = TraceStrategy()
-    val tcv = TraceClassVisitor(writer, PrintWriter(sw))
+    val tcv = TraceClassVisitor(writer, PrintWriter(sw)) // TODO: remove
     return TraceClassProbesAdapter(
         TraceClassInstrumenter(strategy, tcv) { id ->
             storage.computeId(className, id)
         },
-        true,
         className,
         storage
     )
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt
index 639a010739..339fbe7df0 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt
@@ -5,6 +5,7 @@ import org.objectweb.asm.MethodVisitor
 import org.objectweb.asm.Opcodes
 import org.utbot.instrumentation.Settings
 
+// TODO: rename
 class ClassFieldProbeArrayStrategy(private val className: String) : IProbeArrayStrategy {
 
     override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
@@ -14,6 +15,7 @@ class ClassFieldProbeArrayStrategy(private val className: String) : IProbeArrayS
     }
 
     override fun addMembers(cv: ClassVisitor, probeCount: Int) {
+        // nothing to do
     }
 
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
index 4e6c0ef16f..2ebd6844f3 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
@@ -27,16 +27,9 @@ class BranchCoverageInstrumentation : CoverageInstrumentation() {
     ): ByteArray {
         val instrumenter = Instrumenter(classfileBuffer)
 
-        val cv = instrumenter.visitClass(object : ClassVisitorBuilder<ClassProbesAdapter> {
-            override val writerFlags: Int
-                get() = 0
-
-            override val readerParsingOptions: Int
-                get() = ClassReader.EXPAND_FRAMES
-
-            override fun build(writer: ClassWriter): ClassProbesAdapter =
-                createClassVisitorForBranchCoverageInstrumentation(writer, className)
-        })
+        val cv = instrumenter.visitClass { writer ->
+            createClassVisitorForBranchCoverageInstrumentation(writer, className)
+        }
 
         // TODO exceptions
         val probesCountField = ClassProbesAdapter::class.java.declaredFields.firstOrNull { it.name == "counter" }
@@ -47,6 +40,7 @@ class BranchCoverageInstrumentation : CoverageInstrumentation() {
 
         instrumenter.addStaticField(StaticArrayFieldInitializer(className, Settings.PROBES_ARRAY_NAME, probeCount))
 
+        // TODO remove
         val result = sw.toString()
         println(result)
 

From 7acd40d7616cf5f679e4578ba225e2c875fdbc16 Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Sat, 5 Aug 2023 18:16:35 +0300
Subject: [PATCH 05/13] Refactor

---
 .../internal/flow/TraceClassProbesAdapter.kt  |  50 ----
 .../internal/flow/TraceMethodProbesAdapter.kt | 247 ------------------
 .../org/jacoco/core/internal/instr/Api.kt     |  10 +-
 .../internal/instr/TraceClassInstrumenter.kt  |  10 +-
 .../internal/instr/TraceMethodInstrumenter.kt | 149 +++++++++++
 .../et/ExecutionBranchTraceInstrumentation.kt |  11 +-
 6 files changed, 163 insertions(+), 314 deletions(-)
 delete mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
 delete mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt

diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
deleted file mode 100644
index 7423502641..0000000000
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceClassProbesAdapter.kt
+++ /dev/null
@@ -1,50 +0,0 @@
-package org.jacoco.core.internal.flow
-
-import org.objectweb.asm.ClassVisitor
-import org.objectweb.asm.MethodVisitor
-import org.utbot.instrumentation.Settings
-import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
-
-class TraceClassProbesAdapter(
-    val cv: ClassProbesVisitor,
-    private val className: String,
-    private val storage: ProcessingStorage
-) : IProbeIdGenerator, ClassVisitor(Settings.ASM_API, cv) {
-
-    private val EMPTY_METHOD_PROBES_VISITOR: MethodProbesVisitor = object : MethodProbesVisitor() {}
-    private var counter: Int = 0
-
-    override fun visitMethod(
-        access: Int, name: String,
-        desc: String, signature: String?,
-        exceptions: Array<String?>?
-    ): MethodVisitor {
-        val methodProbes: MethodProbesVisitor = (cv as ClassProbesVisitor).visitMethod(
-            access, name, desc,
-            signature, exceptions
-        ) ?: EMPTY_METHOD_PROBES_VISITOR
-
-        return object : MethodSanitizer(
-            null, access, name, desc, signature,
-            exceptions
-        ) {
-            override fun visitEnd() {
-                super.visitEnd()
-                LabelFlowAnalyzer.markLabels(this)
-                val probesAdapter = TraceMethodProbesAdapter(
-                    methodProbes, this@TraceClassProbesAdapter,
-                    storage, className, name, desc
-                )
-                methodProbes.accept(this, probesAdapter)
-            }
-        }
-    }
-
-    override fun visitEnd() {
-        (cv as ClassProbesVisitor).visitTotalProbeCount(counter)
-        super.visitEnd()
-    }
-
-    override fun nextId(): Int = counter++
-
-}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
deleted file mode 100644
index 0b1df6dfe7..0000000000
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/TraceMethodProbesAdapter.kt
+++ /dev/null
@@ -1,247 +0,0 @@
-package org.jacoco.core.internal.flow
-
-import org.objectweb.asm.Label
-import org.objectweb.asm.MethodVisitor
-import org.objectweb.asm.Opcodes
-import org.utbot.instrumentation.Settings
-import org.utbot.instrumentation.instrumentation.et.CommonInstruction
-import org.utbot.instrumentation.instrumentation.et.ExplicitThrowInstruction
-import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
-import org.utbot.instrumentation.instrumentation.et.ReturnInstruction
-import org.utbot.instrumentation.instrumentation.instrumenter.visitors.util.InstructionVisitorMethodAdapter.Companion.returnInsns
-
-class TraceMethodProbesAdapter(
-    private val probesVisitor: MethodProbesVisitor,
-    private val idGenerator: IProbeIdGenerator,
-    private val storage: ProcessingStorage,
-    private val className: String,
-    methodName: String,
-    descriptor: String,
-) : MethodVisitor(Settings.ASM_API, probesVisitor) {
-
-    private val tryCatchProbeLabels: HashMap<Label, Label> = hashMapOf()
-    private var currentLineNumber: Int = 0
-    private val currentMethodSignature: String = methodName + descriptor
-
-    override fun visitTryCatchBlock(start: Label, end: Label, handler: Label, type: String) {
-        probesVisitor.visitTryCatchBlock(getTryCatchLabel(start), getTryCatchLabel(end), handler, type)
-    }
-
-    private fun getTryCatchLabel(l: Label): Label {
-        var label: Label = l
-        if (tryCatchProbeLabels.containsKey(label)) {
-            label = tryCatchProbeLabels[label]!!
-        } else if (LabelInfo.needsProbe(label)) {
-            // If a probe will be inserted before the label, we'll need to use a
-            // different label to define the range of the try-catch block.
-            val probeLabel = Label()
-            LabelInfo.setSuccessor(probeLabel)
-            tryCatchProbeLabels[label] = probeLabel
-            label = probeLabel
-        }
-        return label
-    }
-
-    override fun visitLabel(label: Label?) {
-        if (LabelInfo.needsProbe(label)) {
-            if (tryCatchProbeLabels.containsKey(label)) {
-                probesVisitor.visitLabel(tryCatchProbeLabels[label])
-            }
-            val localId = idGenerator.nextId()
-            val id = storage.computeId(className, localId)
-            val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
-            storage.addInstruction(id, instruction)
-            probesVisitor.visitProbe(localId)
-        }
-        probesVisitor.visitLabel(label)
-    }
-
-    override fun visitLineNumber(line: Int, start: Label?) {
-        currentLineNumber = line
-        super.visitLineNumber(line, start)
-    }
-
-    override fun visitInsn(opcode: Int) {
-        when (opcode) {
-            in returnInsns -> {
-                val localId = idGenerator.nextId()
-                val id = storage.computeId(className, localId)
-                val instruction = ReturnInstruction(currentLineNumber, currentMethodSignature)
-                storage.addInstruction(id, instruction)
-                probesVisitor.visitInsnWithProbe(
-                    opcode,
-                    localId
-                )
-            }
-
-            Opcodes.ATHROW -> {
-                val localId = idGenerator.nextId()
-                val id = storage.computeId(className, localId)
-                val instruction = ExplicitThrowInstruction(currentLineNumber, currentMethodSignature)
-                storage.addInstruction(id, instruction)
-                probesVisitor.visitInsnWithProbe(
-                    opcode,
-                    localId
-                )
-            }
-
-            else -> probesVisitor.visitInsn(opcode)
-        }
-    }
-
-    override fun visitJumpInsn(opcode: Int, label: Label?) {
-        if (LabelInfo.isMultiTarget(label)) {
-            val localId = idGenerator.nextId()
-            val id = storage.computeId(className, localId)
-            val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
-            storage.addInstruction(id, instruction)
-            probesVisitor.visitJumpInsnWithProbe(opcode, label, localId, frame(jumpPopCount(opcode)))
-        } else {
-            probesVisitor.visitJumpInsn(opcode, label)
-        }
-    }
-
-    private fun jumpPopCount(opcode: Int): Int {
-        return when (opcode) {
-            Opcodes.GOTO -> 0
-
-            Opcodes.IFEQ,
-            Opcodes.IFNE,
-            Opcodes.IFLT,
-            Opcodes.IFGE,
-            Opcodes.IFGT,
-            Opcodes.IFLE,
-            Opcodes.IFNULL,
-            Opcodes.IFNONNULL -> 1
-
-            else -> 2
-        }
-    }
-
-    override fun visitLookupSwitchInsn(dflt: Label, keys: IntArray, labels: Array<Label>) {
-        if (markLabels(dflt, labels)) {
-            visitLookupSwitchInsnWithProbes(
-                dflt, keys, labels,
-                frame(1)
-            )
-        } else {
-            probesVisitor.visitLookupSwitchInsn(dflt, keys, labels)
-        }
-    }
-
-    override fun visitTableSwitchInsn(min: Int, max: Int, dflt: Label, vararg labels: Label) {
-        if (markLabels(dflt, labels)) {
-            visitTableSwitchInsnWithProbes(
-                min, max, dflt, labels,
-                frame(1)
-            )
-        } else {
-            probesVisitor.visitTableSwitchInsn(min, max, dflt, *labels)
-        }
-    }
-
-    private fun markLabels(dflt: Label, labels: Array<out Label>): Boolean {
-        var probe = false
-        LabelInfo.resetDone(labels)
-        if (LabelInfo.isMultiTarget(dflt)) {
-            LabelInfo.setProbeId(dflt, idGenerator.nextId())
-            probe = true
-        }
-        LabelInfo.setDone(dflt)
-        for (l in labels) {
-            if (LabelInfo.isMultiTarget(l) && !LabelInfo.isDone(l)) {
-                LabelInfo.setProbeId(l, idGenerator.nextId())
-                probe = true
-            }
-            LabelInfo.setDone(l)
-        }
-        return probe
-    }
-
-    private fun frame(popCount: Int): IFrame {
-        return FrameSnapshot.create(null, popCount)
-    }
-
-    // === MethodProbesVisitor ===
-
-    private fun visitLookupSwitchInsnWithProbes(dflt: Label, keys: IntArray?, labels: Array<Label>, frame: IFrame) {
-        // 1. Calculate intermediate labels:
-        LabelInfo.resetDone(dflt)
-        LabelInfo.resetDone(labels)
-        val newDflt = createIntermediate(dflt)
-        val newLabels = createIntermediates(labels)
-        probesVisitor.visitLookupSwitchInsn(newDflt, keys, newLabels)
-
-        // 2. Insert probes:
-        insertIntermediateProbes(dflt, labels, frame)
-    }
-
-    private fun visitTableSwitchInsnWithProbes(
-        min: Int,
-        max: Int,
-        dflt: Label,
-        labels: Array<out Label>,
-        frame: IFrame
-    ) {
-        // 1. Calculate intermediate labels:
-        LabelInfo.resetDone(dflt)
-        LabelInfo.resetDone(labels)
-        val newDflt = createIntermediate(dflt)
-        val newLabels = createIntermediates(labels)
-        probesVisitor.visitTableSwitchInsn(min, max, newDflt, *newLabels)
-
-        // 2. Insert probes:
-        insertIntermediateProbes(dflt, labels, frame)
-    }
-
-
-    // === MethodInstrumenter ===
-
-    private fun createIntermediates(labels: Array<out Label>): Array<Label?> {
-        val intermediates = arrayOfNulls<Label>(labels.size)
-        for (i in labels.indices) {
-            intermediates[i] = createIntermediate(labels[i])
-        }
-        return intermediates
-    }
-
-    private fun createIntermediate(label: Label): Label {
-        val intermediate: Label
-        if (LabelInfo.getProbeId(label) == LabelInfo.NO_PROBE) {
-            intermediate = label
-        } else {
-            if (LabelInfo.isDone(label)) {
-                intermediate = LabelInfo.getIntermediateLabel(label)
-            } else {
-                intermediate = Label()
-                LabelInfo.setIntermediateLabel(label, intermediate)
-                LabelInfo.setDone(label)
-            }
-        }
-        return intermediate
-    }
-
-    private fun insertIntermediateProbe(label: Label, frame: IFrame) {
-        val localId = LabelInfo.getProbeId(label)
-        if (localId != LabelInfo.NO_PROBE && !LabelInfo.isDone(label)) {
-            probesVisitor.visitLabel(LabelInfo.getIntermediateLabel(label))
-            frame.accept(probesVisitor)
-            val id = storage.computeId(className, localId)
-            val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
-            storage.addInstruction(id, instruction)
-            probesVisitor.visitProbe(localId)
-            probesVisitor.visitJumpInsn(Opcodes.GOTO, label)
-            LabelInfo.setDone(label)
-        }
-    }
-
-    private fun insertIntermediateProbes(dflt: Label, labels: Array<out Label>, frame: IFrame) {
-        LabelInfo.resetDone(dflt)
-        LabelInfo.resetDone(labels)
-        insertIntermediateProbe(dflt, frame)
-        for (l in labels) {
-            insertIntermediateProbe(l, frame)
-        }
-    }
-
-}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
index cf177b3d0c..b21f965e25 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -2,7 +2,6 @@ package org.jacoco.core.internal.instr
 
 import org.jacoco.core.internal.flow.ClassProbesAdapter
 import org.jacoco.core.internal.flow.MethodProbesCollector
-import org.jacoco.core.internal.flow.TraceClassProbesAdapter
 import org.objectweb.asm.ClassWriter
 import org.objectweb.asm.util.TraceClassVisitor
 import org.utbot.instrumentation.Settings
@@ -32,14 +31,13 @@ fun createClassVisitorForTracingBranchInstructions(
     className: String,
     storage: ProcessingStorage,
     writer: ClassWriter
-): TraceClassProbesAdapter {
+): ClassProbesAdapter {
     val strategy = TraceStrategy()
     val tcv = TraceClassVisitor(writer, PrintWriter(sw)) // TODO: remove
-    return TraceClassProbesAdapter(
-        TraceClassInstrumenter(strategy, tcv) { id ->
+    return ClassProbesAdapter(
+        TraceClassInstrumenter(className, storage, strategy, tcv) { id ->
             storage.computeId(className, id)
         },
-        className,
-        storage
+        false
     )
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
index 04b39dd76e..09710811c9 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
@@ -2,8 +2,11 @@ package org.jacoco.core.internal.instr
 
 import org.jacoco.core.internal.flow.MethodProbesVisitor
 import org.objectweb.asm.ClassVisitor
+import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
 
 class TraceClassInstrumenter(
+    private val className: String,
+    private val storage: ProcessingStorage,
     private val probeArrayStrategy: IProbeArrayStrategy,
     cv: ClassVisitor,
     private val nextIdGenerator: (id: Int) -> Long
@@ -19,9 +22,10 @@ class TraceClassInstrumenter(
         val mv = cv.visitMethod(access, name, desc, signature, exceptions)
 
         val frameEliminator = DuplicateFrameEliminator(mv)
-        val probeVariableInserter = TraceProbeInserter(access,
-            name, desc, frameEliminator, probeArrayStrategy, nextIdGenerator)
-        return MethodInstrumenter(probeVariableInserter, probeVariableInserter)
+        val probeVariableInserter = TraceProbeInserter(
+            access, name, desc, frameEliminator, probeArrayStrategy, nextIdGenerator
+        )
+        return TraceMethodInstrumenter(className, storage, name, desc, probeVariableInserter, probeVariableInserter)
     }
 
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
new file mode 100644
index 0000000000..603f9b31a7
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
@@ -0,0 +1,149 @@
+package org.jacoco.core.internal.instr
+
+import org.jacoco.core.internal.flow.IFrame
+import org.jacoco.core.internal.flow.LabelInfo
+import org.objectweb.asm.Label
+import org.objectweb.asm.MethodVisitor
+import org.objectweb.asm.Opcodes
+import org.utbot.instrumentation.instrumentation.et.CommonInstruction
+import org.utbot.instrumentation.instrumentation.et.ExplicitThrowInstruction
+import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
+import org.utbot.instrumentation.instrumentation.et.ReturnInstruction
+import org.utbot.instrumentation.instrumentation.instrumenter.visitors.util.InstructionVisitorMethodAdapter.Companion.returnInsns
+
+internal class TraceMethodInstrumenter(
+    private val className: String,
+    private val storage: ProcessingStorage,
+    methodName: String,
+    descriptor: String,
+    mv: MethodVisitor,
+    private val probeInserter: IProbeInserter
+) : MethodInstrumenter(mv, probeInserter) {
+
+    private var currentLineNumber: Int = 0
+    private val currentMethodSignature: String = methodName + descriptor
+
+    // === MethodVisitor ===
+
+    override fun visitLineNumber(line: Int, start: Label?) {
+        currentLineNumber = line
+        super.visitLineNumber(line, start)
+    }
+
+    // === MethodInstrumenter ===
+
+    override fun visitProbe(localId: Int) {
+        val id = storage.computeId(className, localId)
+        val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
+        storage.addInstruction(id, instruction)
+        super.visitProbe(localId)
+    }
+
+    override fun visitInsnWithProbe(opcode: Int, localId: Int) {
+        when (opcode) {
+            in returnInsns -> {
+                val id = storage.computeId(className, localId)
+                val instruction = ReturnInstruction(currentLineNumber, currentMethodSignature)
+                storage.addInstruction(id, instruction)
+            }
+
+            Opcodes.ATHROW -> {
+                val id = storage.computeId(className, localId)
+                val instruction = ExplicitThrowInstruction(currentLineNumber, currentMethodSignature)
+                storage.addInstruction(id, instruction)
+            }
+        }
+        super.visitInsnWithProbe(opcode, localId)
+    }
+
+    override fun visitJumpInsnWithProbe(opcode: Int, label: Label?, localId: Int, frame: IFrame?) {
+        val id = storage.computeId(className, localId)
+        val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
+        storage.addInstruction(id, instruction)
+        super.visitJumpInsnWithProbe(opcode, label, localId, frame)
+    }
+
+    override fun visitTableSwitchInsnWithProbes(
+        min: Int, max: Int,
+        dflt: Label, labels: Array<Label>, frame: IFrame
+    ) {
+        // 1. Calculate intermediate labels:
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        val newDflt = createIntermediate(dflt)
+        val newLabels = createIntermediates(labels)
+        mv.visitTableSwitchInsn(min, max, newDflt, *newLabels)
+
+        // 2. Insert probes:
+        insertIntermediateProbes(dflt, labels, frame)
+    }
+
+    override fun visitLookupSwitchInsnWithProbes(
+        dflt: Label,
+        keys: IntArray?, labels: Array<Label>, frame: IFrame
+    ) {
+        // 1. Calculate intermediate labels:
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        val newDflt = createIntermediate(dflt)
+        val newLabels = createIntermediates(labels)
+        mv.visitLookupSwitchInsn(newDflt, keys, newLabels)
+
+        // 2. Insert probes:
+        insertIntermediateProbes(dflt, labels, frame)
+    }
+
+    private fun createIntermediates(labels: Array<Label>): Array<Label?> {
+        val intermediates = arrayOfNulls<Label>(labels.size)
+        for (i in labels.indices) {
+            intermediates[i] = createIntermediate(labels[i])
+        }
+        return intermediates
+    }
+
+    private fun createIntermediate(label: Label): Label {
+        val intermediate: Label
+        if (LabelInfo.getProbeId(label) == LabelInfo.NO_PROBE) {
+            intermediate = label
+        } else {
+            if (LabelInfo.isDone(label)) {
+                intermediate = LabelInfo.getIntermediateLabel(label)
+            } else {
+                intermediate = Label()
+                LabelInfo.setIntermediateLabel(label, intermediate)
+                LabelInfo.setDone(label)
+            }
+        }
+        return intermediate
+    }
+
+    private fun insertIntermediateProbe(
+        label: Label,
+        frame: IFrame
+    ) {
+        val localId = LabelInfo.getProbeId(label)
+        if (localId != LabelInfo.NO_PROBE && !LabelInfo.isDone(label)) {
+            mv.visitLabel(LabelInfo.getIntermediateLabel(label))
+            frame.accept(mv)
+            val id = storage.computeId(className, localId)
+            val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
+            storage.addInstruction(id, instruction)
+            probeInserter.insertProbe(localId)
+            mv.visitJumpInsn(Opcodes.GOTO, label)
+            LabelInfo.setDone(label)
+        }
+    }
+
+    private fun insertIntermediateProbes(
+        dflt: Label,
+        labels: Array<Label>, frame: IFrame
+    ) {
+        LabelInfo.resetDone(dflt)
+        LabelInfo.resetDone(labels)
+        insertIntermediateProbe(dflt, frame)
+        for (l in labels) {
+            insertIntermediateProbe(l, frame)
+        }
+    }
+
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionBranchTraceInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionBranchTraceInstrumentation.kt
index 0761024037..8797eaee02 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionBranchTraceInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/et/ExecutionBranchTraceInstrumentation.kt
@@ -1,10 +1,7 @@
 package org.utbot.instrumentation.instrumentation.et
 
-import org.jacoco.core.internal.flow.TraceClassProbesAdapter
 import org.jacoco.core.internal.instr.createClassVisitorForTracingBranchInstructions
-import org.objectweb.asm.ClassWriter
 import org.utbot.instrumentation.instrumentation.Instrumentation
-import org.utbot.instrumentation.instrumentation.instrumenter.ClassVisitorBuilder
 import org.utbot.instrumentation.instrumentation.instrumenter.Instrumenter
 import java.security.ProtectionDomain
 
@@ -20,11 +17,9 @@ class ExecutionBranchTraceInstrumentation : ExecutionTraceInstrumentation() {
         val instrumenter = Instrumenter(classfileBuffer)
 
         traceHandler.registerClass(className)
-        instrumenter.visitClass(object : ClassVisitorBuilder<TraceClassProbesAdapter> {
-
-            override fun build(writer: ClassWriter): TraceClassProbesAdapter =
-                createClassVisitorForTracingBranchInstructions(className, traceHandler.processingStorage, writer)
-        })
+        instrumenter.visitClass { writer ->
+            createClassVisitorForTracingBranchInstructions(className, traceHandler.processingStorage, writer)
+        }
 
         return instrumenter.classByteCode
     }

From e734db2826cc00f333e7608fc52e75003aa949ef Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Mon, 7 Aug 2023 20:30:50 +0300
Subject: [PATCH 06/13] Refactor

---
 .../internal/flow/MethodProbesCollector.kt    | 18 ---------
 .../org/jacoco/core/internal/instr/Api.kt     | 40 +++++++++++--------
 .../internal/instr/NoneClassInstrumenter.kt   |  4 ++
 .../internal/instr/TraceClassInstrumenter.kt  | 13 ++++--
 .../internal/instr/TraceMethodInstrumenter.kt | 36 +++++++----------
 .../core/internal/instr/TraceProbeInserter.kt |  5 +--
 .../core/internal/instr/TraceStrategy.kt      |  2 +-
 .../coverage/BranchCoverageInstrumentation.kt | 21 ++++------
 .../util/InstrumentationException.kt          | 23 ++++++++---
 9 files changed, 77 insertions(+), 85 deletions(-)
 delete mode 100644 utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/MethodProbesCollector.kt

diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/MethodProbesCollector.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/MethodProbesCollector.kt
deleted file mode 100644
index c9b06b52b5..0000000000
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/flow/MethodProbesCollector.kt
+++ /dev/null
@@ -1,18 +0,0 @@
-package org.jacoco.core.internal.flow
-
-import org.jacoco.core.internal.instr.IProbeArrayStrategy
-import org.jacoco.core.internal.instr.NoneClassInstrumenter
-import org.objectweb.asm.ClassVisitor
-
-class MethodProbesCollector(
-    api: Int,
-    strategy: IProbeArrayStrategy,
-    writer: ClassVisitor,
-    val methodToProbes: MutableMap<String, MutableList<Int>> = mutableMapOf()
-) : ClassVisitor(
-    api,
-    ClassProbesAdapter(
-        NoneClassInstrumenter(strategy, writer, methodToProbes),
-        true
-    )
-)
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
index b21f965e25..4ca2e32733 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -1,30 +1,39 @@
 package org.jacoco.core.internal.instr
 
 import org.jacoco.core.internal.flow.ClassProbesAdapter
-import org.jacoco.core.internal.flow.MethodProbesCollector
+import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
-import org.objectweb.asm.util.TraceClassVisitor
 import org.utbot.instrumentation.Settings
 import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
-import java.io.PrintWriter
-import java.io.StringWriter
 
-// TODO: remove
-val sw: StringWriter = StringWriter()
-
-fun createClassVisitorForBranchCoverageInstrumentation(writer: ClassWriter, className: String): ClassProbesAdapter {
+fun createClassVisitorForBranchCoverageInstrumentation(
+    writer: ClassWriter,
+    className: String
+): ClassProbesAdapter {
     val strategy = ClassFieldProbeArrayStrategy(className)
-    val tcv = TraceClassVisitor(writer, PrintWriter(sw)) // TODO: remove
     return ClassProbesAdapter(
-        ClassInstrumenter(strategy, tcv),
+        ClassInstrumenter(strategy, writer),
         false
     )
 }
 
-fun createClassVisitorForComputeMapOfRangesForBranchCoverage(writer: ClassWriter): MethodProbesCollector {
+class MethodProbesCollector(
+    strategy: IProbeArrayStrategy,
+    writer: ClassVisitor,
+    val methodToProbes: MutableMap<String, MutableList<Int>> = mutableMapOf()
+) : ClassVisitor(
+    Settings.ASM_API,
+    ClassProbesAdapter(
+        NoneClassInstrumenter(strategy, writer, methodToProbes),
+        false
+    )
+)
+
+fun createClassVisitorForComputeMapOfRangesForBranchCoverage(
+    writer: ClassWriter
+): MethodProbesCollector {
     val strategy = NoneProbeArrayStrategy()
-    val tcv = TraceClassVisitor(writer, PrintWriter(sw)) // TODO: remove
-    return MethodProbesCollector(Settings.ASM_API, strategy, tcv)
+    return MethodProbesCollector(strategy, writer)
 }
 
 fun createClassVisitorForTracingBranchInstructions(
@@ -33,10 +42,9 @@ fun createClassVisitorForTracingBranchInstructions(
     writer: ClassWriter
 ): ClassProbesAdapter {
     val strategy = TraceStrategy()
-    val tcv = TraceClassVisitor(writer, PrintWriter(sw)) // TODO: remove
     return ClassProbesAdapter(
-        TraceClassInstrumenter(className, storage, strategy, tcv) { id ->
-            storage.computeId(className, id)
+        TraceClassInstrumenter(strategy, writer, storage) { localId ->
+            storage.computeId(className, localId)
         },
         false
     )
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
index ebc263bca9..cc8951c9a2 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
@@ -3,6 +3,10 @@ package org.jacoco.core.internal.instr
 import org.jacoco.core.internal.flow.MethodProbesVisitor
 import org.objectweb.asm.ClassVisitor
 
+/**
+ * Just a copy of [ClassInstrumenter] with one overridden method.
+ * This class is used to replace probe inserter.
+ */
 class NoneClassInstrumenter(
     probeArrayStrategy: IProbeArrayStrategy,
     cv: ClassVisitor,
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
index 09710811c9..5c7bd025ac 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
@@ -4,12 +4,15 @@ import org.jacoco.core.internal.flow.MethodProbesVisitor
 import org.objectweb.asm.ClassVisitor
 import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
 
+/**
+ * Just a copy of [ClassInstrumenter] with one overridden method.
+ * This class is used to replace probe inserter and method instrumenter.
+ */
 class TraceClassInstrumenter(
-    private val className: String,
-    private val storage: ProcessingStorage,
     private val probeArrayStrategy: IProbeArrayStrategy,
     cv: ClassVisitor,
-    private val nextIdGenerator: (id: Int) -> Long
+    private val storage: ProcessingStorage,
+    private val nextIdGenerator: (localId: Int) -> Long
 ) : ClassInstrumenter(probeArrayStrategy, cv) {
 
     override fun visitMethod(
@@ -25,7 +28,9 @@ class TraceClassInstrumenter(
         val probeVariableInserter = TraceProbeInserter(
             access, name, desc, frameEliminator, probeArrayStrategy, nextIdGenerator
         )
-        return TraceMethodInstrumenter(className, storage, name, desc, probeVariableInserter, probeVariableInserter)
+        return TraceMethodInstrumenter(
+            name, desc, probeVariableInserter, probeVariableInserter, storage, nextIdGenerator
+        )
     }
 
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
index 603f9b31a7..a49e5e3262 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
@@ -5,19 +5,16 @@ import org.jacoco.core.internal.flow.LabelInfo
 import org.objectweb.asm.Label
 import org.objectweb.asm.MethodVisitor
 import org.objectweb.asm.Opcodes
-import org.utbot.instrumentation.instrumentation.et.CommonInstruction
-import org.utbot.instrumentation.instrumentation.et.ExplicitThrowInstruction
-import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
-import org.utbot.instrumentation.instrumentation.et.ReturnInstruction
+import org.utbot.instrumentation.instrumentation.et.*
 import org.utbot.instrumentation.instrumentation.instrumenter.visitors.util.InstructionVisitorMethodAdapter.Companion.returnInsns
 
 internal class TraceMethodInstrumenter(
-    private val className: String,
-    private val storage: ProcessingStorage,
     methodName: String,
     descriptor: String,
     mv: MethodVisitor,
-    private val probeInserter: IProbeInserter
+    private val probeInserter: IProbeInserter,
+    private val storage: ProcessingStorage,
+    private val nextIdGenerator: (localId: Int) -> Long
 ) : MethodInstrumenter(mv, probeInserter) {
 
     private var currentLineNumber: Int = 0
@@ -32,34 +29,31 @@ internal class TraceMethodInstrumenter(
 
     // === MethodInstrumenter ===
 
+    private fun storeInstruction(localId: Int, instructionData: InstructionData) {
+        val id = nextIdGenerator.invoke(localId)
+        storage.addInstruction(id, instructionData)
+    }
+
     override fun visitProbe(localId: Int) {
-        val id = storage.computeId(className, localId)
-        val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
-        storage.addInstruction(id, instruction)
+        storeInstruction(localId, CommonInstruction(currentLineNumber, currentMethodSignature))
         super.visitProbe(localId)
     }
 
     override fun visitInsnWithProbe(opcode: Int, localId: Int) {
         when (opcode) {
             in returnInsns -> {
-                val id = storage.computeId(className, localId)
-                val instruction = ReturnInstruction(currentLineNumber, currentMethodSignature)
-                storage.addInstruction(id, instruction)
+                storeInstruction(localId, ReturnInstruction(currentLineNumber, currentMethodSignature))
             }
 
             Opcodes.ATHROW -> {
-                val id = storage.computeId(className, localId)
-                val instruction = ExplicitThrowInstruction(currentLineNumber, currentMethodSignature)
-                storage.addInstruction(id, instruction)
+                storeInstruction(localId, ExplicitThrowInstruction(currentLineNumber, currentMethodSignature))
             }
         }
         super.visitInsnWithProbe(opcode, localId)
     }
 
     override fun visitJumpInsnWithProbe(opcode: Int, label: Label?, localId: Int, frame: IFrame?) {
-        val id = storage.computeId(className, localId)
-        val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
-        storage.addInstruction(id, instruction)
+        storeInstruction(localId, CommonInstruction(currentLineNumber, currentMethodSignature))
         super.visitJumpInsnWithProbe(opcode, label, localId, frame)
     }
 
@@ -125,9 +119,7 @@ internal class TraceMethodInstrumenter(
         if (localId != LabelInfo.NO_PROBE && !LabelInfo.isDone(label)) {
             mv.visitLabel(LabelInfo.getIntermediateLabel(label))
             frame.accept(mv)
-            val id = storage.computeId(className, localId)
-            val instruction = CommonInstruction(currentLineNumber, currentMethodSignature)
-            storage.addInstruction(id, instruction)
+            storeInstruction(localId, CommonInstruction(currentLineNumber, currentMethodSignature))
             probeInserter.insertProbe(localId)
             mv.visitJumpInsn(Opcodes.GOTO, label)
             LabelInfo.setDone(label)
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
index 0d2ae55e18..22077bb8bd 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
@@ -12,17 +12,14 @@ internal class TraceProbeInserter(
     desc: String,
     mv: MethodVisitor,
     arrayStrategy: IProbeArrayStrategy,
-    private val nextIdGenerator: (id: Int) -> Long,
+    private val nextIdGenerator: (localId: Int) -> Long,
 ) : ProbeInserter(access, name, desc, mv, arrayStrategy) {
 
     private val internalName = Type.getInternalName(RuntimeTraceStorage::class.java)
     private val visitMethodDescriptor = Type.getMethodDescriptor(RuntimeTraceStorage::visit.javaMethod)
-
-    private val clinit: Boolean
     private val variable: Int
 
     init {
-        clinit = InstrSupport.CLINIT_NAME == name
         var pos = if (Opcodes.ACC_STATIC and access == 0) 1 else 0
         for (t in Type.getArgumentTypes(desc)) {
             pos += t.size
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceStrategy.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceStrategy.kt
index 6abe249292..fe6b2a6937 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceStrategy.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceStrategy.kt
@@ -18,13 +18,13 @@ class TraceStrategy : IProbeArrayStrategy {
         mv.visitInsn(Opcodes.ICONST_1)
         mv.visitInsn(Opcodes.IADD)
         mv.visitInsn(Opcodes.DUP)
-
         mv.visitFieldInsn(Opcodes.PUTSTATIC, internalName, counterCallIdName, RuntimeTraceStorage.DESC_CALL_ID_COUNTER)
         mv.visitVarInsn(Opcodes.ISTORE, variable)
         return 2
     }
 
     override fun addMembers(cv: ClassVisitor?, probeCount: Int) {
+        // nothing to do
     }
 
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
index 2ebd6844f3..02127d43e0 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
@@ -2,14 +2,12 @@ package org.utbot.instrumentation.instrumentation.coverage
 
 import org.jacoco.core.internal.flow.ClassProbesAdapter
 import org.jacoco.core.internal.instr.createClassVisitorForBranchCoverageInstrumentation
-import org.jacoco.core.internal.instr.sw
-import org.objectweb.asm.ClassReader
-import org.objectweb.asm.ClassWriter
 import org.utbot.common.withAccessibility
 import org.utbot.instrumentation.Settings
 import org.utbot.instrumentation.instrumentation.Instrumentation
-import org.utbot.instrumentation.instrumentation.instrumenter.ClassVisitorBuilder
 import org.utbot.instrumentation.instrumentation.instrumenter.Instrumenter
+import org.utbot.instrumentation.util.CastProbeCounterException
+import org.utbot.instrumentation.util.NoProbeCounterException
 import java.security.ProtectionDomain
 
 class BranchCoverageInstrumentation : CoverageInstrumentation() {
@@ -31,18 +29,13 @@ class BranchCoverageInstrumentation : CoverageInstrumentation() {
             createClassVisitorForBranchCoverageInstrumentation(writer, className)
         }
 
-        // TODO exceptions
-        val probesCountField = ClassProbesAdapter::class.java.declaredFields.firstOrNull { it.name == "counter" }
-            ?: throw IllegalStateException()
-        val probeCount = probesCountField.withAccessibility {
-            this.get(cv) as? Int ?: throw ClassCastException()
+        val probeCounterField = ClassProbesAdapter::class.java.declaredFields.firstOrNull { it.name == "counter" }
+            ?: throw NoProbeCounterException(ClassProbesAdapter::class.java, "counter")
+        val probeCounter = probeCounterField.withAccessibility {
+            this.get(cv) as? Int ?: throw CastProbeCounterException()
         }
 
-        instrumenter.addStaticField(StaticArrayFieldInitializer(className, Settings.PROBES_ARRAY_NAME, probeCount))
-
-        // TODO remove
-        val result = sw.toString()
-        println(result)
+        instrumenter.addStaticField(StaticArrayFieldInitializer(className, Settings.PROBES_ARRAY_NAME, probeCounter))
 
         return instrumenter.classByteCode
     }
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/util/InstrumentationException.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/util/InstrumentationException.kt
index 8a617f0168..2596e12e5f 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/util/InstrumentationException.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/util/InstrumentationException.kt
@@ -9,22 +9,33 @@ import org.utbot.framework.plugin.api.InstrumentedProcessDeathException
  */
 open class InstrumentationException(msg: String?, cause: Throwable? = null) : Exception(msg, cause)
 
-class NoProbesArrayException(clazz: Class<*>, arrayName: String) :
+abstract class NoFieldException(clazz: Class<*>, name: String, fieldName: String) :
     InstrumentationException(
-        "No probes array found\n\t" +
-                "Clazz: $clazz\n\t" +
-                "Probes array name: $arrayName\n\t" +
-                "All fields: ${clazz.fields.joinToString { it.name }}"
+        buildString {
+            appendLine("No $name found")
+            appendLine("\tClazz: $clazz")
+            appendLine("\t${name.capitalize()} name: $fieldName")
+            append("All fields: ${clazz.declaredFields.joinToString { it.name }}")
+        }
     )
 
+class NoProbesArrayException(clazz: Class<*>, arrayName: String) :
+    NoFieldException(clazz, "probes array", arrayName)
+
+class NoProbeCounterException(clazz: Class<*>, counterName: String) :
+    NoFieldException(clazz, "probe counter", counterName)
+
 class CastProbesArrayException :
     InstrumentationException("Can't cast probes array to Boolean array")
 
+class CastProbeCounterException :
+    InstrumentationException("Can't cast probe counter to Int")
+
 /**
  * this exception is thrown only in main process.
  * currently it means that {e: Throwable} happened in instrumented process,
  * but instrumented process still can operate and not dead.
  * on instrumented process death - [InstrumentedProcessDeathException] is thrown
-*/
+ */
 class InstrumentedProcessError(e: Throwable) :
     InstrumentationException("Error in the instrumented process |> ${e.stackTraceToString()}", e)
\ No newline at end of file

From 218f61cd7364327a5541a3ed61b02fa02763abe0 Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Mon, 7 Aug 2023 20:37:35 +0300
Subject: [PATCH 07/13] Delete extra dependency

---
 utbot-instrumentation/build.gradle.kts                       | 1 -
 .../instrumentation/instrumenter/Instrumenter.kt             | 5 +----
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/utbot-instrumentation/build.gradle.kts b/utbot-instrumentation/build.gradle.kts
index 4191b25b6e..4a1fb97969 100644
--- a/utbot-instrumentation/build.gradle.kts
+++ b/utbot-instrumentation/build.gradle.kts
@@ -54,7 +54,6 @@ dependencies {
 
     implementation("org.ow2.asm:asm:$asmVersion")
     implementation("org.ow2.asm:asm-commons:$asmVersion")
-    implementation("org.ow2.asm:asm-util:$asmVersion") // TODO: remove
     implementation("io.github.microutils:kotlin-logging:$kotlinLoggingVersion")
 
     implementation("com.jetbrains.rd:rd-framework:$rdVersion")
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
index d1d6caa3a4..9e7f6c4163 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
@@ -5,13 +5,11 @@ import org.objectweb.asm.ClassReader
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.ClassWriter
 import org.objectweb.asm.Opcodes
-import org.objectweb.asm.util.TraceClassVisitor
 import org.utbot.instrumentation.instrumentation.instrumenter.visitors.MethodToProbesVisitor
 import org.utbot.instrumentation.instrumentation.instrumenter.visitors.util.*
 import org.utbot.instrumentation.process.HandlerClassesLoader
 import java.io.IOException
 import java.io.InputStream
-import java.io.PrintWriter
 
 
 // TODO: handle with flags EXPAND_FRAMES, etc.
@@ -66,8 +64,7 @@ class Instrumenter(classByteCode: ByteArray, val classLoader: ClassLoader? = nul
 
     fun addStaticField(staticFieldInitializer: StaticFieldInitializer) {
         visitClass { writer ->
-            val tcv = TraceClassVisitor(writer, PrintWriter(sw))
-            AddStaticFieldAdapter(tcv, staticFieldInitializer)
+            AddStaticFieldAdapter(writer, staticFieldInitializer)
         }
     }
 

From fb003cb56f63133f9830238c879ce70e00504631 Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Tue, 8 Aug 2023 16:31:10 +0300
Subject: [PATCH 08/13] Add a base class CoverageInstrumentation

---
 .../framework/coverage/CoverageCalculator.kt  |  6 +--
 .../org/utbot/examples/TestConstructors.kt    |  4 +-
 ...TestInstructionCoverageInstrumentation.kt} | 20 ++++-----
 .../org/utbot/examples/TestStaticMethods.kt   |  8 ++--
 .../utbot/examples/TestWithInstrumentation.kt |  6 +--
 .../org/utbot/examples/benchmark/Benchmark.kt |  4 +-
 .../examples/benchmark/BenchmarkFibonacci.kt  |  4 +-
 .../benchmark/TestBenchmarkClasses.kt         |  4 +-
 .../utbot/examples/jacoco/TestSameAsJaCoCo.kt |  4 +-
 .../org/jacoco/core/internal/instr/Api.kt     |  2 +-
 .../internal/instr/NoneClassInstrumenter.kt   |  2 +-
 .../core/internal/instr/NoneProbeInserter.kt  |  7 ++-
 ...ArrayStrategy.kt => ProbeArrayStrategy.kt} |  3 +-
 .../coverage/BranchCoverageInstrumentation.kt |  8 ++++
 .../coverage/CoverageInstrumentation.kt       | 45 ++++---------------
 .../InstructionCoverageInstrumentation.kt     | 41 +++++++++++++++++
 .../instrumenter/Instrumenter.kt              |  2 +-
 17 files changed, 96 insertions(+), 74 deletions(-)
 rename utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/{TestCoverageInstrumentation.kt => TestInstructionCoverageInstrumentation.kt} (92%)
 rename utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/{ClassFieldProbeArrayStrategy.kt => ProbeArrayStrategy.kt} (84%)
 create mode 100644 utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/InstructionCoverageInstrumentation.kt

diff --git a/utbot-framework/src/main/kotlin/org/utbot/framework/coverage/CoverageCalculator.kt b/utbot-framework/src/main/kotlin/org/utbot/framework/coverage/CoverageCalculator.kt
index c662da9b31..4bc34a4cc3 100644
--- a/utbot-framework/src/main/kotlin/org/utbot/framework/coverage/CoverageCalculator.kt
+++ b/utbot-framework/src/main/kotlin/org/utbot/framework/coverage/CoverageCalculator.kt
@@ -1,19 +1,19 @@
 package org.utbot.framework.coverage
 
+import kotlinx.coroutines.runBlocking
 import org.utbot.framework.plugin.api.ExecutableId
 import org.utbot.framework.plugin.api.UtValueExecution
 import org.utbot.framework.plugin.api.util.jClass
 import org.utbot.instrumentation.ConcreteExecutor
 import org.utbot.instrumentation.instrumentation.coverage.CoverageInfo
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
 import org.utbot.instrumentation.util.StaticEnvironment
-import kotlinx.coroutines.runBlocking
 
 fun methodCoverage(executable: ExecutableId, executions: List<UtValueExecution<*>>, classpath: String): Coverage {
     val methodSignature = executable.signature
     val classId = executable.classId
-    return ConcreteExecutor(CoverageInstrumentation.Factory, classpath).let { executor ->
+    return ConcreteExecutor(InstructionCoverageInstrumentation.Factory, classpath).let { executor ->
         for (execution in executions) {
             val args = execution.stateBefore.params.map { it.value }.toMutableList()
             val caller = execution.stateBefore.caller
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestConstructors.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestConstructors.kt
index fea5acdb17..fecbe02f8f 100644
--- a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestConstructors.kt
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestConstructors.kt
@@ -6,7 +6,7 @@ import org.utbot.framework.plugin.api.util.signature
 import org.utbot.instrumentation.ConcreteExecutor
 import org.utbot.instrumentation.execute
 import org.utbot.instrumentation.instrumentation.InvokeInstrumentation
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
 import org.utbot.instrumentation.instrumentation.et.ExecutionTraceInstrumentation
 import org.utbot.instrumentation.instrumentation.et.convert
@@ -86,7 +86,7 @@ class TestConstructors {
     @Test
     fun testCoverageConstructor() {
         withInstrumentation(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             CLASSPATH
         ) { executor ->
             val constructors = ClassWithMultipleConstructors::class.constructors
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestCoverageInstrumentation.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestInstructionCoverageInstrumentation.kt
similarity index 92%
rename from utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestCoverageInstrumentation.kt
rename to utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestInstructionCoverageInstrumentation.kt
index 565ca75f12..47a8f70475 100644
--- a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestCoverageInstrumentation.kt
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestInstructionCoverageInstrumentation.kt
@@ -8,7 +8,7 @@ import org.utbot.framework.plugin.api.util.fieldId
 import org.utbot.framework.plugin.api.util.signature
 import org.utbot.instrumentation.ConcreteExecutor
 import org.utbot.instrumentation.execute
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
 import org.utbot.instrumentation.util.InstrumentedProcessError
 import org.utbot.instrumentation.util.StaticEnvironment
@@ -18,13 +18,13 @@ import org.junit.jupiter.api.Assertions.assertTrue
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.assertThrows
 
-class TestCoverageInstrumentation {
+class TestInstructionCoverageInstrumentation {
     lateinit var utContext: AutoCloseable
 
     @Test
     fun testCatchTargetException() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             ExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val testObject = ExampleClass()
@@ -41,7 +41,7 @@ class TestCoverageInstrumentation {
     @Test
     fun testIfBranches() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             ExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val testObject = ExampleClass()
@@ -63,7 +63,7 @@ class TestCoverageInstrumentation {
     @Test
     fun testWrongArgumentsException() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             ExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val testObject = ExampleClass()
@@ -86,7 +86,7 @@ class TestCoverageInstrumentation {
     @Test
     fun testMultipleRunsInsideCoverage() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             ExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val testObject = ExampleClass()
@@ -121,7 +121,7 @@ class TestCoverageInstrumentation {
     @Test
     fun testSameResult() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             ExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val testObject = ExampleClass()
@@ -143,7 +143,7 @@ class TestCoverageInstrumentation {
     @Test
     fun testResult() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             ExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val testObject = ExampleClass()
@@ -160,7 +160,7 @@ class TestCoverageInstrumentation {
     @Test
     fun testEmptyMethod() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             ExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val testObject = ExampleClass()
@@ -176,7 +176,7 @@ class TestCoverageInstrumentation {
     @Test
     fun testTernaryOperator() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             StaticSubstitutionExamples::class.java.protectionDomain.codeSource.location.path
         ).use {
             val testObject = StaticSubstitutionExamples()
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestStaticMethods.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestStaticMethods.kt
index 0a75292636..b60ea80260 100644
--- a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestStaticMethods.kt
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestStaticMethods.kt
@@ -4,7 +4,7 @@ import org.utbot.examples.samples.staticenvironment.StaticExampleClass
 import org.utbot.framework.plugin.api.util.signature
 import org.utbot.instrumentation.ConcreteExecutor
 import org.utbot.instrumentation.execute
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Assertions.assertTrue
@@ -17,7 +17,7 @@ class TestStaticMethods {
     @Test
     fun testStaticMethodCall() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             StaticExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val res1 = it.execute(StaticExampleClass::inc, arrayOf())
@@ -44,7 +44,7 @@ class TestStaticMethods {
     @Test
     fun testNullableMethod() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             StaticExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val res1 = it.execute(
@@ -75,7 +75,7 @@ class TestStaticMethods {
     @Test
     fun testNullableMethodWithoutAnnotations() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             StaticExampleClass::class.java.protectionDomain.codeSource.location.path
         ).use {
             val res1 = it.execute(
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestWithInstrumentation.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestWithInstrumentation.kt
index ae00e0f7a8..dbb95e5df1 100644
--- a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestWithInstrumentation.kt
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestWithInstrumentation.kt
@@ -6,7 +6,7 @@ import org.utbot.examples.samples.staticenvironment.StaticExampleClass
 import org.utbot.framework.plugin.api.util.signature
 import org.utbot.instrumentation.execute
 import org.utbot.instrumentation.instrumentation.InvokeInstrumentation
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
 import org.utbot.instrumentation.withInstrumentation
 import kotlin.reflect.full.declaredMembers
@@ -20,7 +20,7 @@ class TestWithInstrumentation {
     @Test
     fun testStaticMethodCall() {
         withInstrumentation(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             StaticExampleClass::class.java.protectionDomain.codeSource.location.path
         ) { executor ->
             val res1 = executor.execute(StaticExampleClass::inc, arrayOf())
@@ -70,7 +70,7 @@ class TestWithInstrumentation {
     @Test
     fun testInnerClasses() {
         withInstrumentation(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             ClassWithInnerClasses::class.java.protectionDomain.codeSource.location.path
         ) { executor ->
             val innerClazz = ClassWithInnerClasses.InnerClass::class.java
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/Benchmark.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/Benchmark.kt
index d4e9ba2028..43ec154e52 100644
--- a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/Benchmark.kt
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/Benchmark.kt
@@ -2,7 +2,7 @@ package org.utbot.examples.benchmark
 
 import org.utbot.instrumentation.ConcreteExecutor
 import org.utbot.instrumentation.instrumentation.InvokeInstrumentation
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
 import org.utbot.instrumentation.util.Isolated
 import kotlin.system.measureNanoTime
@@ -12,7 +12,7 @@ import org.junit.jupiter.api.Assertions.assertEquals
 fun getBasicCoverageTime(count: Int): Double {
     var time: Long
     ConcreteExecutor(
-        CoverageInstrumentation.Factory,
+        InstructionCoverageInstrumentation.Factory,
         Repeater::class.java.protectionDomain.codeSource.location.path
     ).use { executor ->
         val dc0 = Repeater(", ")
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/BenchmarkFibonacci.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/BenchmarkFibonacci.kt
index 07d788e4ce..88b3205ca2 100644
--- a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/BenchmarkFibonacci.kt
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/BenchmarkFibonacci.kt
@@ -5,7 +5,7 @@ import org.utbot.framework.plugin.api.util.UtContext
 import org.utbot.framework.plugin.api.util.withUtContext
 import org.utbot.instrumentation.ConcreteExecutor
 import org.utbot.instrumentation.instrumentation.InvokeInstrumentation
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
 import org.utbot.instrumentation.util.Isolated
 import kotlin.system.measureNanoTime
@@ -13,7 +13,7 @@ import kotlin.system.measureNanoTime
 fun getBasicCoverageTime_fib(count: Int): Double {
     var time: Long
     ConcreteExecutor(
-        CoverageInstrumentation.Factory,
+        InstructionCoverageInstrumentation.Factory,
         Fibonacci::class.java.protectionDomain.codeSource.location.path
     ).use {
         val fib = Isolated(Fibonacci::calc, it)
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/TestBenchmarkClasses.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/TestBenchmarkClasses.kt
index f189907c5d..31d0556f02 100644
--- a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/TestBenchmarkClasses.kt
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/benchmark/TestBenchmarkClasses.kt
@@ -4,7 +4,7 @@ import org.utbot.examples.samples.benchmark.Fibonacci
 import org.utbot.instrumentation.ConcreteExecutor
 import org.utbot.instrumentation.execute
 import org.utbot.instrumentation.instrumentation.InvokeInstrumentation
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import java.math.BigInteger
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Disabled
@@ -18,7 +18,7 @@ class TestBenchmarkClasses {
     @Disabled("Ask Sergey to check")
     fun testRepeater() {
         ConcreteExecutor(
-            CoverageInstrumentation.Factory,
+            InstructionCoverageInstrumentation.Factory,
             Repeater::class.java.protectionDomain.codeSource.location.path
         ).use {
             val dc0 = Repeater(", ")
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/jacoco/TestSameAsJaCoCo.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/jacoco/TestSameAsJaCoCo.kt
index a289e1d5fa..5c946c0e3a 100644
--- a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/jacoco/TestSameAsJaCoCo.kt
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/jacoco/TestSameAsJaCoCo.kt
@@ -8,7 +8,7 @@ import org.utbot.framework.plugin.api.util.signature
 import org.utbot.instrumentation.execute
 import org.utbot.instrumentation.instrumentation.ArgumentList
 import org.utbot.instrumentation.instrumentation.coverage.CoverageInfo
-import org.utbot.instrumentation.instrumentation.coverage.CoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.InstructionCoverageInstrumentation
 import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
 import org.utbot.instrumentation.withInstrumentation
 import java.io.InputStream
@@ -123,7 +123,7 @@ private fun methodCoverageWithJaCoCo(kClass: KClass<*>, method: KCallable<*>, ex
 
 private fun methodCoverage(kClass: KClass<*>, method: KCallable<*>, executions: List<ArgumentList>): Pair<Int, Int> {
     return withInstrumentation(
-        CoverageInstrumentation.Factory,
+        InstructionCoverageInstrumentation.Factory,
         kClass.java.protectionDomain.codeSource.location.path
     ) { executor ->
         for (execution in executions) {
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
index 4ca2e32733..e34d040112 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -10,7 +10,7 @@ fun createClassVisitorForBranchCoverageInstrumentation(
     writer: ClassWriter,
     className: String
 ): ClassProbesAdapter {
-    val strategy = ClassFieldProbeArrayStrategy(className)
+    val strategy = ProbeArrayStrategy(className)
     return ClassProbesAdapter(
         ClassInstrumenter(strategy, writer),
         false
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
index cc8951c9a2..6fada8d8b2 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneClassInstrumenter.kt
@@ -23,7 +23,7 @@ class NoneClassInstrumenter(
         val mv = cv.visitMethod(access, name, desc, signature, exceptions)
 
         val frameEliminator = DuplicateFrameEliminator(mv)
-        val probeVariableInserter = NoneProbeInserter(name, methodToProbes, frameEliminator)
+        val probeVariableInserter = NoneProbeInserter(name, desc, methodToProbes, frameEliminator)
         return MethodInstrumenter(probeVariableInserter, probeVariableInserter)
     }
 
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
index 67777762ac..cabce1c099 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/NoneProbeInserter.kt
@@ -7,13 +7,16 @@ import org.utbot.instrumentation.Settings
  * This inserter does not emit any code at all. This is used to collect method probes.
  */
 class NoneProbeInserter(
-    private val methodName: String,
+    methodName: String,
+    descriptor: String,
     private val methodToProbes: MutableMap<String, MutableList<Int>>,
     mv: MethodVisitor,
 ) : IProbeInserter, MethodVisitor(Settings.ASM_API, mv) {
 
+    private val currentMethodSignature: String = methodName + descriptor
+
     override fun insertProbe(id: Int) {
-        methodToProbes.getOrPut(methodName) { mutableListOf() }.add(id)
+        methodToProbes.getOrPut(currentMethodSignature) { mutableListOf() }.add(id)
     }
 
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ProbeArrayStrategy.kt
similarity index 84%
rename from utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt
rename to utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ProbeArrayStrategy.kt
index 339fbe7df0..1057b4302f 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ClassFieldProbeArrayStrategy.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/ProbeArrayStrategy.kt
@@ -5,8 +5,7 @@ import org.objectweb.asm.MethodVisitor
 import org.objectweb.asm.Opcodes
 import org.utbot.instrumentation.Settings
 
-// TODO: rename
-class ClassFieldProbeArrayStrategy(private val className: String) : IProbeArrayStrategy {
+class ProbeArrayStrategy(private val className: String) : IProbeArrayStrategy {
 
     override fun storeInstance(mv: MethodVisitor, clinit: Boolean, variable: Int): Int {
         mv.visitFieldInsn(Opcodes.GETSTATIC, className, Settings.PROBES_ARRAY_NAME, Settings.PROBES_ARRAY_DESC)
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
index 02127d43e0..de4aae7ef3 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/BranchCoverageInstrumentation.kt
@@ -10,12 +10,20 @@ import org.utbot.instrumentation.util.CastProbeCounterException
 import org.utbot.instrumentation.util.NoProbeCounterException
 import java.security.ProtectionDomain
 
+/**
+ * This instrumentation allows collecting branch coverage after several calls.
+ */
 class BranchCoverageInstrumentation : CoverageInstrumentation() {
 
     override fun <T : Any> computeMapOfRanges(clazz: Class<out T>): Map<String, IntRange> {
         return Instrumenter(clazz).computeMapOfRangesForBranchCoverage()
     }
 
+    /**
+     * Transforms bytecode such way that it becomes possible to get a branch coverage by JaCoCo.
+     *
+     * Adds set of instructions which marks the executed instruction as completed. Uses static boolean array for this.
+     */
     override fun transform(
         loader: ClassLoader?,
         className: String,
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
index 5b63455ec4..811c01eafd 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/CoverageInstrumentation.kt
@@ -2,27 +2,22 @@ package org.utbot.instrumentation.instrumentation.coverage
 
 import kotlinx.coroutines.runBlocking
 import org.utbot.common.withAccessibility
+import org.utbot.framework.plugin.api.FieldId
 import org.utbot.instrumentation.ConcreteExecutor
 import org.utbot.instrumentation.Settings
 import org.utbot.instrumentation.instrumentation.ArgumentList
 import org.utbot.instrumentation.instrumentation.Instrumentation
 import org.utbot.instrumentation.instrumentation.InvokeWithStaticsInstrumentation
-import org.utbot.instrumentation.instrumentation.instrumenter.Instrumenter
+import org.utbot.instrumentation.process.generated.CollectCoverageParams
 import org.utbot.instrumentation.util.CastProbesArrayException
 import org.utbot.instrumentation.util.NoProbesArrayException
-import java.security.ProtectionDomain
-import org.utbot.framework.plugin.api.FieldId
-import org.utbot.instrumentation.process.generated.CollectCoverageParams
 
 data class CoverageInfo(
     val methodToInstrRange: Map<String, IntRange>,
     val visitedInstrs: List<Int>
 )
 
-/**
- * This instrumentation allows collecting coverage after several calls.
- */
-open class CoverageInstrumentation : Instrumentation<Result<*>> {
+abstract class CoverageInstrumentation : Instrumentation<Result<*>> {
     private val invokeWithStatics = InvokeWithStaticsInstrumentation()
 
     /**
@@ -53,7 +48,7 @@ open class CoverageInstrumentation : Instrumentation<Result<*>> {
     /**
      * Collects coverage from the given [clazz] via reflection.
      */
-    open fun <T : Any> collectCoverageInfo(clazz: Class<out T>): CoverageInfo {
+    fun <T : Any> collectCoverageInfo(clazz: Class<out T>): CoverageInfo {
         val probesFieldName: String = Settings.PROBES_ARRAY_NAME
         val visitedLinesField = clazz.fields.firstOrNull { it.name == probesFieldName }
             ?: throw NoProbesArrayException(clazz, Settings.PROBES_ARRAY_NAME)
@@ -74,34 +69,8 @@ open class CoverageInstrumentation : Instrumentation<Result<*>> {
         }
     }
 
-    protected open fun <T : Any> computeMapOfRanges(clazz: Class<out T>): Map<String, IntRange> {
-        return Instrumenter(clazz).computeMapOfRanges()
-    }
+    abstract fun <T : Any> computeMapOfRanges(clazz: Class<out T>): Map<String, IntRange>
 
-    /**
-     * Transforms bytecode such way that it becomes possible to get a coverage.
-     *
-     * Adds set of instructions which marks the executed instruction as completed. Uses static boolean array for this.
-     */
-    override fun transform(
-        loader: ClassLoader?,
-        className: String,
-        classBeingRedefined: Class<*>?,
-        protectionDomain: ProtectionDomain?,
-        classfileBuffer: ByteArray
-    ): ByteArray {
-        val instrumenter = Instrumenter(classfileBuffer)
-
-        val staticArrayStrategy = StaticArrayStrategy(className, Settings.PROBES_ARRAY_NAME)
-        instrumenter.visitInstructions(staticArrayStrategy)
-        instrumenter.addStaticField(staticArrayStrategy)
-
-        return instrumenter.classByteCode
-    }
-
-    object Factory : Instrumentation.Factory<Result<*>, CoverageInstrumentation> {
-        override fun create(): CoverageInstrumentation = CoverageInstrumentation()
-    }
 }
 
 /**
@@ -111,6 +80,8 @@ fun ConcreteExecutor<Result<*>, CoverageInstrumentation>.collectCoverage(clazz:
     withProcess {
         val clazzByteArray = kryoHelper.writeObject(clazz)
 
-        kryoHelper.readObject(instrumentedProcessModel.collectCoverage.startSuspending(lifetime, CollectCoverageParams(clazzByteArray)).coverageInfo)
+        kryoHelper.readObject(
+            instrumentedProcessModel.collectCoverage.startSuspending(lifetime, CollectCoverageParams(clazzByteArray)).coverageInfo
+        )
     }
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/InstructionCoverageInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/InstructionCoverageInstrumentation.kt
new file mode 100644
index 0000000000..e89336a886
--- /dev/null
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/coverage/InstructionCoverageInstrumentation.kt
@@ -0,0 +1,41 @@
+package org.utbot.instrumentation.instrumentation.coverage
+
+import org.utbot.instrumentation.Settings
+import org.utbot.instrumentation.instrumentation.Instrumentation
+import org.utbot.instrumentation.instrumentation.instrumenter.Instrumenter
+import java.security.ProtectionDomain
+
+/**
+ * This instrumentation allows collecting instruction coverage after several calls.
+ */
+open class InstructionCoverageInstrumentation : CoverageInstrumentation() {
+
+    override fun <T : Any> computeMapOfRanges(clazz: Class<out T>): Map<String, IntRange> {
+        return Instrumenter(clazz).computeMapOfRangesForInstructionCoverage()
+    }
+
+    /**
+     * Transforms bytecode such way that it becomes possible to get an instruction coverage.
+     *
+     * Adds set of instructions which marks the executed instruction as completed. Uses static boolean array for this.
+     */
+    override fun transform(
+        loader: ClassLoader?,
+        className: String,
+        classBeingRedefined: Class<*>?,
+        protectionDomain: ProtectionDomain?,
+        classfileBuffer: ByteArray
+    ): ByteArray {
+        val instrumenter = Instrumenter(classfileBuffer)
+
+        val staticArrayStrategy = StaticArrayStrategy(className, Settings.PROBES_ARRAY_NAME)
+        instrumenter.visitInstructions(staticArrayStrategy)
+        instrumenter.addStaticField(staticArrayStrategy)
+
+        return instrumenter.classByteCode
+    }
+
+    object Factory : Instrumentation.Factory<Result<*>, InstructionCoverageInstrumentation> {
+        override fun create(): InstructionCoverageInstrumentation = InstructionCoverageInstrumentation()
+    }
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
index 9e7f6c4163..43831435c2 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/instrumenter/Instrumenter.kt
@@ -36,7 +36,7 @@ class Instrumenter(classByteCode: ByteArray, val classLoader: ClassLoader? = nul
         return classVisitor
     }
 
-    fun computeMapOfRanges(methodName: String? = null): Map<String, IntRange> {
+    fun computeMapOfRangesForInstructionCoverage(methodName: String? = null): Map<String, IntRange> {
         val methodToListOfProbesInserter = MethodToProbesVisitor()
 
         visitClass(object : ClassVisitorBuilder<InstructionVisitorAdapter> {

From e186e1708447095faa19098d6097895f92aa2bf4 Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Wed, 9 Aug 2023 21:40:44 +0300
Subject: [PATCH 09/13] Fix bug

---
 .../org/jacoco/core/internal/instr/Api.kt     | 17 ++++++++++++---
 .../internal/instr/TraceClassInstrumenter.kt  |  6 +++---
 .../internal/instr/TraceMethodInstrumenter.kt | 21 ++++++++++++-------
 .../core/internal/instr/TraceProbeInserter.kt | 11 +++++++---
 4 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
index e34d040112..f4e7741fd5 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -36,16 +36,27 @@ fun createClassVisitorForComputeMapOfRangesForBranchCoverage(
     return MethodProbesCollector(strategy, writer)
 }
 
+class ProbeIdGenerator(private val f: (Int) -> Long) {
+
+    private var localId: Int = 0
+
+    fun currentId(): Long = f(localId)
+
+    fun nextId(): Long = f(localId++)
+
+}
+
 fun createClassVisitorForTracingBranchInstructions(
     className: String,
     storage: ProcessingStorage,
     writer: ClassWriter
 ): ClassProbesAdapter {
     val strategy = TraceStrategy()
+    val probeIdGenerator = ProbeIdGenerator { localId ->
+        storage.computeId(className, localId)
+    }
     return ClassProbesAdapter(
-        TraceClassInstrumenter(strategy, writer, storage) { localId ->
-            storage.computeId(className, localId)
-        },
+        TraceClassInstrumenter(strategy, writer, storage, probeIdGenerator),
         false
     )
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
index 5c7bd025ac..259991406e 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
@@ -12,7 +12,7 @@ class TraceClassInstrumenter(
     private val probeArrayStrategy: IProbeArrayStrategy,
     cv: ClassVisitor,
     private val storage: ProcessingStorage,
-    private val nextIdGenerator: (localId: Int) -> Long
+    private val probeIdGenerator: ProbeIdGenerator
 ) : ClassInstrumenter(probeArrayStrategy, cv) {
 
     override fun visitMethod(
@@ -26,10 +26,10 @@ class TraceClassInstrumenter(
 
         val frameEliminator = DuplicateFrameEliminator(mv)
         val probeVariableInserter = TraceProbeInserter(
-            access, name, desc, frameEliminator, probeArrayStrategy, nextIdGenerator
+            access, name, desc, frameEliminator, probeArrayStrategy, probeIdGenerator
         )
         return TraceMethodInstrumenter(
-            name, desc, probeVariableInserter, probeVariableInserter, storage, nextIdGenerator
+            name, desc, probeVariableInserter, probeVariableInserter, storage, probeIdGenerator
         )
     }
 
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
index a49e5e3262..284c57536c 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
@@ -14,7 +14,7 @@ internal class TraceMethodInstrumenter(
     mv: MethodVisitor,
     private val probeInserter: IProbeInserter,
     private val storage: ProcessingStorage,
-    private val nextIdGenerator: (localId: Int) -> Long
+    private val probeIdGenerator: ProbeIdGenerator
 ) : MethodInstrumenter(mv, probeInserter) {
 
     private var currentLineNumber: Int = 0
@@ -22,6 +22,11 @@ internal class TraceMethodInstrumenter(
 
     // === MethodVisitor ===
 
+    override fun visitCode() {
+        storeInstruction(CommonInstruction(currentLineNumber, currentMethodSignature))
+        super.visitCode()
+    }
+
     override fun visitLineNumber(line: Int, start: Label?) {
         currentLineNumber = line
         super.visitLineNumber(line, start)
@@ -29,31 +34,31 @@ internal class TraceMethodInstrumenter(
 
     // === MethodInstrumenter ===
 
-    private fun storeInstruction(localId: Int, instructionData: InstructionData) {
-        val id = nextIdGenerator.invoke(localId)
+    private fun storeInstruction(instructionData: InstructionData) {
+        val id = probeIdGenerator.currentId()
         storage.addInstruction(id, instructionData)
     }
 
     override fun visitProbe(localId: Int) {
-        storeInstruction(localId, CommonInstruction(currentLineNumber, currentMethodSignature))
+        storeInstruction(CommonInstruction(currentLineNumber, currentMethodSignature))
         super.visitProbe(localId)
     }
 
     override fun visitInsnWithProbe(opcode: Int, localId: Int) {
         when (opcode) {
             in returnInsns -> {
-                storeInstruction(localId, ReturnInstruction(currentLineNumber, currentMethodSignature))
+                storeInstruction(ReturnInstruction(currentLineNumber, currentMethodSignature))
             }
 
             Opcodes.ATHROW -> {
-                storeInstruction(localId, ExplicitThrowInstruction(currentLineNumber, currentMethodSignature))
+                storeInstruction(ExplicitThrowInstruction(currentLineNumber, currentMethodSignature))
             }
         }
         super.visitInsnWithProbe(opcode, localId)
     }
 
     override fun visitJumpInsnWithProbe(opcode: Int, label: Label?, localId: Int, frame: IFrame?) {
-        storeInstruction(localId, CommonInstruction(currentLineNumber, currentMethodSignature))
+        storeInstruction(CommonInstruction(currentLineNumber, currentMethodSignature))
         super.visitJumpInsnWithProbe(opcode, label, localId, frame)
     }
 
@@ -119,7 +124,7 @@ internal class TraceMethodInstrumenter(
         if (localId != LabelInfo.NO_PROBE && !LabelInfo.isDone(label)) {
             mv.visitLabel(LabelInfo.getIntermediateLabel(label))
             frame.accept(mv)
-            storeInstruction(localId, CommonInstruction(currentLineNumber, currentMethodSignature))
+            storeInstruction(CommonInstruction(currentLineNumber, currentMethodSignature))
             probeInserter.insertProbe(localId)
             mv.visitJumpInsn(Opcodes.GOTO, label)
             LabelInfo.setDone(label)
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
index 22077bb8bd..f6824752c6 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
@@ -12,7 +12,7 @@ internal class TraceProbeInserter(
     desc: String,
     mv: MethodVisitor,
     arrayStrategy: IProbeArrayStrategy,
-    private val nextIdGenerator: (localId: Int) -> Long,
+    private val probeIdGenerator: ProbeIdGenerator,
 ) : ProbeInserter(access, name, desc, mv, arrayStrategy) {
 
     private val internalName = Type.getInternalName(RuntimeTraceStorage::class.java)
@@ -27,11 +27,16 @@ internal class TraceProbeInserter(
         variable = pos
     }
 
-    override fun insertProbe(localId: Int) {
+    override fun insertProbe(ignored: Int) {
         mv.visitVarInsn(Opcodes.ILOAD, variable)
-        val id = nextIdGenerator.invoke(localId)
+        val id = probeIdGenerator.nextId()
         mv.visitLdcInsn(id)
         mv.visitMethodInsn(Opcodes.INVOKESTATIC, internalName, "visit", visitMethodDescriptor, false)
     }
 
+    override fun visitCode() {
+        super.visitCode()
+        insertProbe(0)
+    }
+
 }
\ No newline at end of file

From 0babb3af2a312a2c8e698a957f794dcc12819687 Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Thu, 10 Aug 2023 12:28:57 +0300
Subject: [PATCH 10/13] Add flag to select trace instrumentation

---
 .../kotlin/org/utbot/framework/UtSettings.kt  | 18 ++++++++++
 .../agent/DynamicClassTransformer.kt          |  3 ++
 .../SimpleUtExecutionInstrumentation.kt       | 12 ++++++-
 .../process/InstrumentedProcessMain.kt        |  1 +
 .../instrumentation/rd/InstrumentedProcess.kt |  3 +-
 .../InstrumentedProcessModel.Generated.kt     | 34 +++++++++++--------
 .../rd/models/InstrumentedProcessModel.kt     |  1 +
 7 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/utbot-framework-api/src/main/kotlin/org/utbot/framework/UtSettings.kt b/utbot-framework-api/src/main/kotlin/org/utbot/framework/UtSettings.kt
index 97f456ecd4..b0a74ad1d6 100644
--- a/utbot-framework-api/src/main/kotlin/org/utbot/framework/UtSettings.kt
+++ b/utbot-framework-api/src/main/kotlin/org/utbot/framework/UtSettings.kt
@@ -475,6 +475,11 @@ object UtSettings : AbstractSettings(logger, defaultKeyForSettingsPath, defaultS
      */
     var useBytecodeTransformation by getBooleanProperty(false)
 
+    /**
+     * Type of trace instrumentation.
+     */
+    var traceInstrumentationType by getEnumProperty(TraceInstrumentationType.INSTRUCTION)
+
     /**
      * Limit for number of generated tests per method (in each region)
      */
@@ -751,3 +756,16 @@ enum class ExploreThrowableDepth {
      */
     EXPLORE_ALL_STATEMENTS
 }
+
+enum class TraceInstrumentationType {
+
+    /**
+     * Insert probes before each bytecode instruction.
+     */
+    INSTRUCTION,
+
+    /**
+     * Insert probes into the [control flow graph](https://www.eclemma.org/jacoco/trunk/doc/flow.html).
+     */
+    BRANCH
+}
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/agent/DynamicClassTransformer.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/agent/DynamicClassTransformer.kt
index c3e83d4694..ab0a824f89 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/agent/DynamicClassTransformer.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/agent/DynamicClassTransformer.kt
@@ -4,6 +4,7 @@ import com.jetbrains.rd.util.error
 import com.jetbrains.rd.util.getLogger
 import com.jetbrains.rd.util.info
 import org.utbot.common.asPathToFile
+import org.utbot.framework.TraceInstrumentationType
 import org.utbot.framework.plugin.api.util.UtContext
 import java.lang.instrument.ClassFileTransformer
 import java.nio.file.Paths
@@ -21,6 +22,8 @@ class DynamicClassTransformer : ClassFileTransformer {
     lateinit var transformer: ClassFileTransformer
 
     var useBytecodeTransformation by Delegates.notNull<Boolean>()
+    lateinit var traceInstrumentationType: TraceInstrumentationType
+
     private val pathsToUserClasses = mutableSetOf<String>()
 
     fun addUserPaths(paths: Iterable<String>) {
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/SimpleUtExecutionInstrumentation.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/SimpleUtExecutionInstrumentation.kt
index d888c9cdbe..2fc5867d1d 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/SimpleUtExecutionInstrumentation.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/instrumentation/execution/SimpleUtExecutionInstrumentation.kt
@@ -1,5 +1,7 @@
 package org.utbot.instrumentation.instrumentation.execution
 
+import org.jacoco.core.internal.instr.createClassVisitorForTracingBranchInstructions
+import org.utbot.framework.TraceInstrumentationType
 import org.utbot.framework.plugin.api.EnvironmentModels
 import org.utbot.framework.plugin.api.FieldId
 import org.utbot.framework.plugin.api.UtModel
@@ -153,7 +155,15 @@ class SimpleUtExecutionInstrumentation(
         }
 
         traceHandler.registerClass(className)
-        instrumenter.visitInstructions(traceHandler.computeInstructionVisitor(className))
+        when (Agent.dynamicClassTransformer.traceInstrumentationType) {
+            TraceInstrumentationType.INSTRUCTION -> instrumenter.visitInstructions(
+                traceHandler.computeInstructionVisitor(className)
+            )
+
+            TraceInstrumentationType.BRANCH -> instrumenter.visitClass { writer ->
+                createClassVisitorForTracingBranchInstructions(className, traceHandler.processingStorage, writer)
+            }
+        }
 
         instrumenter.visitClass { writer ->
             NonDeterministicClassVisitor(writer, ndDetector)
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/process/InstrumentedProcessMain.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/process/InstrumentedProcessMain.kt
index 34bb5300c7..c4df71f0a8 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/process/InstrumentedProcessMain.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/process/InstrumentedProcessMain.kt
@@ -147,6 +147,7 @@ private fun InstrumentedProcessModel.setup(kryoHelper: KryoHelper, watchdog: Idl
         instrumentation = instrumentationFactory.create()
         logger.debug { "instrumentation - ${instrumentation.javaClass.name} " }
         Agent.dynamicClassTransformer.useBytecodeTransformation = params.useBytecodeTransformation
+        Agent.dynamicClassTransformer.traceInstrumentationType = kryoHelper.readObject(params.traceInstrumentationType)
         Agent.dynamicClassTransformer.transformer = instrumentation
         Agent.dynamicClassTransformer.addUserPaths(pathsToUserClasses)
     }
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/rd/InstrumentedProcess.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/rd/InstrumentedProcess.kt
index 0bfc2ac99d..57a66439fe 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/rd/InstrumentedProcess.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/rd/InstrumentedProcess.kt
@@ -133,7 +133,8 @@ class InstrumentedProcess private constructor(
             proc.instrumentedProcessModel.setInstrumentation.startSuspending(
                 proc.lifetime, SetInstrumentationParams(
                     proc.kryoHelper.writeObject(instrumentationFactory),
-                    UtSettings.useBytecodeTransformation
+                    UtSettings.useBytecodeTransformation,
+                    proc.kryoHelper.writeObject(UtSettings.traceInstrumentationType)
                 )
             )
             logger.trace("start commands sent")
diff --git a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/rd/generated/InstrumentedProcessModel.Generated.kt b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/rd/generated/InstrumentedProcessModel.Generated.kt
index 5500b9121b..373aa07ac3 100644
--- a/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/rd/generated/InstrumentedProcessModel.Generated.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/utbot/instrumentation/rd/generated/InstrumentedProcessModel.Generated.kt
@@ -67,7 +67,7 @@ class InstrumentedProcessModel private constructor(
         }
         
         
-        const val serializationHash = 2667021258656776274L
+        const val serializationHash = -1318319258783684404L
         
     }
     override val serializersOwner: ISerializersOwner get() = InstrumentedProcessModel
@@ -260,7 +260,7 @@ data class AddPathsParams (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:29]
+ * #### Generated from [InstrumentedProcessModel.kt:30]
  */
 data class CollectCoverageParams (
     val clazz: ByteArray
@@ -317,7 +317,7 @@ data class CollectCoverageParams (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:33]
+ * #### Generated from [InstrumentedProcessModel.kt:34]
  */
 data class CollectCoverageResult (
     val coverageInfo: ByteArray
@@ -374,7 +374,7 @@ data class CollectCoverageResult (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:37]
+ * #### Generated from [InstrumentedProcessModel.kt:38]
  */
 data class ComputeStaticFieldParams (
     val fieldId: ByteArray
@@ -431,7 +431,7 @@ data class ComputeStaticFieldParams (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:41]
+ * #### Generated from [InstrumentedProcessModel.kt:42]
  */
 data class ComputeStaticFieldResult (
     val result: ByteArray
@@ -488,7 +488,7 @@ data class ComputeStaticFieldResult (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:45]
+ * #### Generated from [InstrumentedProcessModel.kt:46]
  */
 data class GetSpringBeanParams (
     val beanName: String
@@ -545,7 +545,7 @@ data class GetSpringBeanParams (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:49]
+ * #### Generated from [InstrumentedProcessModel.kt:50]
  */
 data class GetSpringBeanResult (
     val beanModel: ByteArray
@@ -602,7 +602,7 @@ data class GetSpringBeanResult (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:53]
+ * #### Generated from [InstrumentedProcessModel.kt:54]
  */
 data class GetSpringRepositoriesParams (
     val classId: ByteArray
@@ -659,7 +659,7 @@ data class GetSpringRepositoriesParams (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:57]
+ * #### Generated from [InstrumentedProcessModel.kt:58]
  */
 data class GetSpringRepositoriesResult (
     val springRepositoryIds: ByteArray
@@ -716,7 +716,7 @@ data class GetSpringRepositoriesResult (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:18]
+ * #### Generated from [InstrumentedProcessModel.kt:19]
  */
 data class InvokeMethodCommandParams (
     val classname: String,
@@ -791,7 +791,7 @@ data class InvokeMethodCommandParams (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:25]
+ * #### Generated from [InstrumentedProcessModel.kt:26]
  */
 data class InvokeMethodCommandResult (
     val result: ByteArray
@@ -852,7 +852,8 @@ data class InvokeMethodCommandResult (
  */
 data class SetInstrumentationParams (
     val instrumentation: ByteArray,
-    val useBytecodeTransformation: Boolean
+    val useBytecodeTransformation: Boolean,
+    val traceInstrumentationType: ByteArray
 ) : IPrintable {
     //companion
     
@@ -863,12 +864,14 @@ data class SetInstrumentationParams (
         override fun read(ctx: SerializationCtx, buffer: AbstractBuffer): SetInstrumentationParams  {
             val instrumentation = buffer.readByteArray()
             val useBytecodeTransformation = buffer.readBool()
-            return SetInstrumentationParams(instrumentation, useBytecodeTransformation)
+            val traceInstrumentationType = buffer.readByteArray()
+            return SetInstrumentationParams(instrumentation, useBytecodeTransformation, traceInstrumentationType)
         }
         
         override fun write(ctx: SerializationCtx, buffer: AbstractBuffer, value: SetInstrumentationParams)  {
             buffer.writeByteArray(value.instrumentation)
             buffer.writeBool(value.useBytecodeTransformation)
+            buffer.writeByteArray(value.traceInstrumentationType)
         }
         
         
@@ -886,6 +889,7 @@ data class SetInstrumentationParams (
         
         if (!(instrumentation contentEquals other.instrumentation)) return false
         if (useBytecodeTransformation != other.useBytecodeTransformation) return false
+        if (!(traceInstrumentationType contentEquals other.traceInstrumentationType)) return false
         
         return true
     }
@@ -894,6 +898,7 @@ data class SetInstrumentationParams (
         var __r = 0
         __r = __r*31 + instrumentation.contentHashCode()
         __r = __r*31 + useBytecodeTransformation.hashCode()
+        __r = __r*31 + traceInstrumentationType.contentHashCode()
         return __r
     }
     //pretty print
@@ -902,6 +907,7 @@ data class SetInstrumentationParams (
         printer.indent {
             print("instrumentation = "); instrumentation.print(printer); println()
             print("useBytecodeTransformation = "); useBytecodeTransformation.print(printer); println()
+            print("traceInstrumentationType = "); traceInstrumentationType.print(printer); println()
         }
         printer.print(")")
     }
@@ -911,7 +917,7 @@ data class SetInstrumentationParams (
 
 
 /**
- * #### Generated from [InstrumentedProcessModel.kt:61]
+ * #### Generated from [InstrumentedProcessModel.kt:62]
  */
 data class TryLoadingSpringContextResult (
     val springContextLoadingResult: ByteArray
diff --git a/utbot-rd/src/main/rdgen/org/utbot/rd/models/InstrumentedProcessModel.kt b/utbot-rd/src/main/rdgen/org/utbot/rd/models/InstrumentedProcessModel.kt
index db0d827687..c0b91ef01e 100644
--- a/utbot-rd/src/main/rdgen/org/utbot/rd/models/InstrumentedProcessModel.kt
+++ b/utbot-rd/src/main/rdgen/org/utbot/rd/models/InstrumentedProcessModel.kt
@@ -13,6 +13,7 @@ object InstrumentedProcessModel : Ext(InstrumentedProcessRoot) {
     val SetInstrumentationParams = structdef {
         field("instrumentation", array(PredefinedType.byte))
         field("useBytecodeTransformation", PredefinedType.bool)
+        field("traceInstrumentationType", array(PredefinedType.byte))
     }
 
     val InvokeMethodCommandParams = structdef {

From 0065fc34ba5ad419d602ec2545100e71c390675c Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Fri, 11 Aug 2023 09:33:15 +0300
Subject: [PATCH 11/13] Add inserting probes before method instruction

---
 .../core/internal/instr/TraceMethodInstrumenter.kt    | 11 +++++++++++
 .../jacoco/core/internal/instr/TraceProbeInserter.kt  | 11 +++++++++++
 2 files changed, 22 insertions(+)

diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
index 284c57536c..9fb631c9f1 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
@@ -27,6 +27,17 @@ internal class TraceMethodInstrumenter(
         super.visitCode()
     }
 
+    override fun visitMethodInsn(
+        opcode: Int,
+        owner: String?,
+        name: String?,
+        descriptor: String?,
+        isInterface: Boolean
+    ) {
+        storeInstruction(InvokeInstruction(currentLineNumber, currentMethodSignature))
+        super.visitMethodInsn(opcode, owner, name, descriptor, isInterface)
+    }
+
     override fun visitLineNumber(line: Int, start: Label?) {
         currentLineNumber = line
         super.visitLineNumber(line, start)
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
index f6824752c6..ee5e54aef4 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceProbeInserter.kt
@@ -39,4 +39,15 @@ internal class TraceProbeInserter(
         insertProbe(0)
     }
 
+    override fun visitMethodInsn(
+        opcode: Int,
+        owner: String?,
+        name: String?,
+        descriptor: String?,
+        isInterface: Boolean
+    ) {
+        insertProbe(0)
+        super.visitMethodInsn(opcode, owner, name, descriptor, isInterface)
+    }
+
 }
\ No newline at end of file

From fc6e5ad551ff7dbd3694665961c93268a327fc4a Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Fri, 11 Aug 2023 14:52:37 +0300
Subject: [PATCH 12/13] Add tests for branch coverage instrumentation and
 branch trace instrumentation

---
 .../TestBranchCoverageInstrumentation.kt      |  93 +++++++++++++
 ...TestExecutionBranchTraceInstrumentation.kt | 127 ++++++++++++++++++
 2 files changed, 220 insertions(+)
 create mode 100644 utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestBranchCoverageInstrumentation.kt
 create mode 100644 utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/et/TestExecutionBranchTraceInstrumentation.kt

diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestBranchCoverageInstrumentation.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestBranchCoverageInstrumentation.kt
new file mode 100644
index 0000000000..03b5602f4b
--- /dev/null
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/TestBranchCoverageInstrumentation.kt
@@ -0,0 +1,93 @@
+package org.utbot.examples
+
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import org.utbot.examples.samples.ExampleClass
+import org.utbot.examples.samples.et.ClassSimpleCatch
+import org.utbot.examples.statics.substitution.StaticSubstitution
+import org.utbot.examples.statics.substitution.StaticSubstitutionExamples
+import org.utbot.framework.plugin.api.util.fieldId
+import org.utbot.framework.plugin.api.util.signature
+import org.utbot.instrumentation.execute
+import org.utbot.instrumentation.instrumentation.coverage.BranchCoverageInstrumentation
+import org.utbot.instrumentation.instrumentation.coverage.collectCoverage
+import org.utbot.instrumentation.util.StaticEnvironment
+import org.utbot.instrumentation.withInstrumentation
+
+class TestBranchCoverageInstrumentation {
+    lateinit var utContext: AutoCloseable
+
+    @Test
+    fun testIfBranches() {
+        withInstrumentation(
+            BranchCoverageInstrumentation.Factory,
+            ExampleClass::class.java.protectionDomain.codeSource.location.path
+        ) { executor ->
+            val testObject = ExampleClass()
+
+            executor.execute(ExampleClass::bar, arrayOf(testObject, 2))
+            val coverageInfo1 = executor.collectCoverage(ExampleClass::class.java)
+
+            assertEquals(2, coverageInfo1.visitedInstrs.size)
+            assertEquals(1..3, coverageInfo1.methodToInstrRange[ExampleClass::bar.signature])
+
+            executor.execute(ExampleClass::bar, arrayOf(testObject, 0))
+            val coverageInfo2 = executor.collectCoverage(ExampleClass::class.java)
+
+            assertEquals(2, coverageInfo2.visitedInstrs.size)
+            assertEquals(1..3, coverageInfo2.methodToInstrRange[ExampleClass::bar.signature])
+        }
+    }
+
+    @Test
+    fun testTryCatch() {
+        withInstrumentation(
+            BranchCoverageInstrumentation.Factory,
+            ClassSimpleCatch::class.java.protectionDomain.codeSource.location.path
+        ) { executor ->
+            executor.execute(ClassSimpleCatch::A_catches, emptyArray())
+            val coverageInfo1 = executor.collectCoverage(ClassSimpleCatch::class.java)
+
+            assertEquals(2, coverageInfo1.visitedInstrs.size)
+            assertEquals(3..5, coverageInfo1.methodToInstrRange[ClassSimpleCatch::A_catches.signature])
+
+            executor.execute(ClassSimpleCatch::A_catchesWrongException, emptyArray())
+            val coverageInfo2 = executor.collectCoverage(ClassSimpleCatch::class.java)
+
+            assertEquals(0, coverageInfo2.visitedInstrs.size)
+            assertEquals(7..9, coverageInfo2.methodToInstrRange[ClassSimpleCatch::A_catchesWrongException.signature])
+        }
+    }
+
+    @Test
+    fun testTernaryOperator() {
+        withInstrumentation(
+            BranchCoverageInstrumentation.Factory,
+            StaticSubstitutionExamples::class.java.protectionDomain.codeSource.location.path
+        ) { executor ->
+            val testObject = StaticSubstitutionExamples()
+
+            val emptyStaticEnvironment = StaticEnvironment()
+
+            val res1 = executor.execute(
+                StaticSubstitutionExamples::lessThanZero,
+                arrayOf(testObject),
+                parameters = emptyStaticEnvironment
+            )
+
+            val staticEnvironment = StaticEnvironment(
+                StaticSubstitution::mutableValue.fieldId to -1
+            )
+            val res2 = executor.execute(
+                StaticSubstitutionExamples::lessThanZero,
+                arrayOf(testObject),
+                parameters = staticEnvironment
+            )
+            val coverageInfo = executor.collectCoverage(StaticSubstitutionExamples::class.java)
+
+            assertEquals(res1.getOrNull(), 5)
+            assertEquals(res2.getOrNull(), 0)
+            assertEquals(coverageInfo.visitedInstrs, (1..3).toList())
+        }
+    }
+}
\ No newline at end of file
diff --git a/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/et/TestExecutionBranchTraceInstrumentation.kt b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/et/TestExecutionBranchTraceInstrumentation.kt
new file mode 100644
index 0000000000..923e21fa1d
--- /dev/null
+++ b/utbot-instrumentation-tests/src/test/kotlin/org/utbot/examples/et/TestExecutionBranchTraceInstrumentation.kt
@@ -0,0 +1,127 @@
+package org.utbot.examples.et
+
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.Test
+import org.utbot.examples.samples.et.ClassSimple
+import org.utbot.examples.samples.et.ClassSimpleCatch
+import org.utbot.instrumentation.instrumentation.et.*
+import org.utbot.instrumentation.util.Isolated
+import org.utbot.instrumentation.withInstrumentation
+
+class TestExecutionBranchTraceInstrumentation {
+    lateinit var utContext: AutoCloseable
+
+    val CLASSPATH = ClassSimple::class.java.protectionDomain.codeSource.location.path
+
+    @Test
+    fun testClassSimple() {
+        withInstrumentation(
+            ExecutionBranchTraceInstrumentation.Factory,
+            CLASSPATH
+        ) { executor ->
+            val doesNotThrow = Isolated(ClassSimple::doesNotThrow, executor)
+            val alwaysThrows = Isolated(ClassSimple::alwaysThrows, executor)
+            val maybeThrows = Isolated(ClassSimple::maybeThrows, executor)
+
+            val et1 = doesNotThrow()
+            Assertions.assertEquals(
+                function(doesNotThrow.signature) {
+                    pass()
+                    ret()
+                },
+                convert(et1)
+            )
+
+            val et2 = alwaysThrows()
+            Assertions.assertEquals(
+                function(alwaysThrows.signature) {
+                    pass()
+                    explThr()
+                },
+                convert(et2)
+            )
+
+            val et3 = maybeThrows(-1)
+            Assertions.assertEquals(
+                function(maybeThrows.signature) {
+                    pass()
+                    explThr()
+                },
+                convert(et3)
+            )
+
+            val et4 = maybeThrows(0)
+            Assertions.assertEquals(
+                function(maybeThrows.signature) {
+                    pass()
+                    ret()
+                },
+                convert(et4)
+            )
+        }
+    }
+
+    @Test
+    fun testClasSimpleCatch() {
+        withInstrumentation(
+            ExecutionBranchTraceInstrumentation.Factory,
+            CLASSPATH
+        ) { executor ->
+            val A = Isolated(ClassSimpleCatch::A, executor)
+            val A_catches = Isolated(ClassSimpleCatch::A_catches, executor)
+            val A_catchesWrongException = Isolated(ClassSimpleCatch::A_catchesWrongException, executor)
+            val A_doesNotCatch = Isolated(ClassSimpleCatch::A_doesNotCatch, executor)
+
+            val B = Isolated(ClassSimpleCatch::B, executor)
+            val B_throws = Isolated(ClassSimpleCatch::B_throws, executor)
+
+            val et1 = A()
+            Assertions.assertEquals(
+                function(A.signature) {
+                    pass()
+                    invoke(B.signature) {
+                        pass()
+                        ret()
+                    }
+                    ret()
+                },
+                convert(et1)
+            )
+
+            val et2 = A_catches()
+            Assertions.assertEquals(
+                function(A_catches.signature) {
+                    pass()
+                    invoke(B_throws.signature) {
+                        implThr()
+                    }
+                    pass()
+                    ret()
+                },
+                convert(et2)
+            )
+
+            val et3 = A_catchesWrongException()
+            Assertions.assertEquals(
+                function(A_catchesWrongException.signature) {
+                    pass()
+                    invoke(B_throws.signature) {
+                        implThr()
+                    }
+                },
+                convert(et3)
+            )
+
+            val et4 = A_doesNotCatch()
+            Assertions.assertEquals(
+                function(A_doesNotCatch.signature) {
+                    pass()
+                    invoke(B_throws.signature) {
+                        implThr()
+                    }
+                },
+                convert(et4)
+            )
+        }
+    }
+}

From 47758fefcc14f6760c811681e445d2438dd3c556 Mon Sep 17 00:00:00 2001
From: Pavel Egipti <kinqstone585@gmail.com>
Date: Thu, 17 Aug 2023 17:49:23 +0300
Subject: [PATCH 13/13] Fix bug

---
 .../src/main/kotlin/org/jacoco/core/internal/instr/Api.kt   | 2 +-
 .../jacoco/core/internal/instr/TraceClassInstrumenter.kt    | 6 +++++-
 .../jacoco/core/internal/instr/TraceMethodInstrumenter.kt   | 4 +---
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
index f4e7741fd5..e8ba862153 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/Api.kt
@@ -56,7 +56,7 @@ fun createClassVisitorForTracingBranchInstructions(
         storage.computeId(className, localId)
     }
     return ClassProbesAdapter(
-        TraceClassInstrumenter(strategy, writer, storage, probeIdGenerator),
+        TraceClassInstrumenter(strategy, writer, className, storage, probeIdGenerator),
         false
     )
 }
\ No newline at end of file
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
index 259991406e..f996ff8f59 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceClassInstrumenter.kt
@@ -11,6 +11,7 @@ import org.utbot.instrumentation.instrumentation.et.ProcessingStorage
 class TraceClassInstrumenter(
     private val probeArrayStrategy: IProbeArrayStrategy,
     cv: ClassVisitor,
+    private val className: String,
     private val storage: ProcessingStorage,
     private val probeIdGenerator: ProbeIdGenerator
 ) : ClassInstrumenter(probeArrayStrategy, cv) {
@@ -24,12 +25,15 @@ class TraceClassInstrumenter(
     ): MethodProbesVisitor {
         val mv = cv.visitMethod(access, name, desc, signature, exceptions)
 
+        val currentMethodSignature = name + desc
+        storage.addClassMethod(className, currentMethodSignature)
+
         val frameEliminator = DuplicateFrameEliminator(mv)
         val probeVariableInserter = TraceProbeInserter(
             access, name, desc, frameEliminator, probeArrayStrategy, probeIdGenerator
         )
         return TraceMethodInstrumenter(
-            name, desc, probeVariableInserter, probeVariableInserter, storage, probeIdGenerator
+            currentMethodSignature, probeVariableInserter, probeVariableInserter, storage, probeIdGenerator
         )
     }
 
diff --git a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
index 9fb631c9f1..80f9b83fd9 100644
--- a/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
+++ b/utbot-instrumentation/src/main/kotlin/org/jacoco/core/internal/instr/TraceMethodInstrumenter.kt
@@ -9,8 +9,7 @@ import org.utbot.instrumentation.instrumentation.et.*
 import org.utbot.instrumentation.instrumentation.instrumenter.visitors.util.InstructionVisitorMethodAdapter.Companion.returnInsns
 
 internal class TraceMethodInstrumenter(
-    methodName: String,
-    descriptor: String,
+    private val currentMethodSignature: String,
     mv: MethodVisitor,
     private val probeInserter: IProbeInserter,
     private val storage: ProcessingStorage,
@@ -18,7 +17,6 @@ internal class TraceMethodInstrumenter(
 ) : MethodInstrumenter(mv, probeInserter) {
 
     private var currentLineNumber: Int = 0
-    private val currentMethodSignature: String = methodName + descriptor
 
     // === MethodVisitor ===