Introduce common AbstractDeterministicRandomTransformer, add more tes…

…ts for deterministic randomness and fix found bugs

bootstrap/src/sun/nio/ch/lincheck/Injections.java

@@ -275,20 +275,6 @@ public static int nextInt() {
         return getEventTracker().randomNextInt();
     }
 
-    /**
-     * Called from the instrumented code to replace `ThreadLocalRandom.nextInt(origin, bound)` with a deterministic random value.
-     */
-    public static int nextInt2(int origin, int bound) {
-        boolean enteredIgnoredSection = enterIgnoredSection();
-        try {
-            return deterministicRandom().nextInt(bound);
-        } finally {
-            if (enteredIgnoredSection) {
-                leaveIgnoredSection();
-            }
-        }
-    }
-
     /**
      * Called from the instrumented code to get a random instance that is deterministic and controlled by Lincheck.
      */

src/jvm/main/org/jetbrains/kotlinx/lincheck/strategy/managed/NativeCallsUtils.kt

@@ -133,7 +133,7 @@ fun GeneratorAdapter.customGetOrThrow(successType: Type) {
     dup()
     ifStatement(
         condition = { instanceOf(throwableWrapperType) },
-        ifClause = {
+        thenClause = {
             invokeStatic(ThrowableWrapper::throwable)
             throwException()
         },

src/jvm/main/org/jetbrains/kotlinx/lincheck/transformation/LincheckClassVisitor.kt

@@ -180,7 +180,7 @@ internal class LincheckClassVisitor(
         } else {
             mv = FakeDeterministicTimeTransformer(mv.newAdapter())
         }
-        mv = DeterministicRandomTransformer(fileName, className, methodName, mv.newAdapter())
+        mv = FakeDeterministicRandomTransformer(fileName, className, methodName, mv.newAdapter())
         // `SharedMemoryAccessTransformer` goes first because it relies on `AnalyzerAdapter`,
         // which should be put in front of the byte-code transformer chain,
         // so that it can correctly analyze the byte-code and compute required type-information
@@ -254,7 +254,7 @@ private class WrapMethodInIgnoredSectionTransformer(
             ARETURN, DRETURN, FRETURN, IRETURN, LRETURN, RETURN -> {
                 ifStatement(
                     condition = { loadLocal(enteredInIgnoredSectionLocal) },
-                    ifClause = { invokeStatic(Injections::leaveIgnoredSection) },
+                    thenClause = { invokeStatic(Injections::leaveIgnoredSection) },
                     elseClause = {}
                 )
             }

src/jvm/main/org/jetbrains/kotlinx/lincheck/transformation/TraceDebuggerClassVisitor.kt

@@ -86,7 +86,7 @@ private class TraceDebuggerRunMethodTransformer(
             condition = {
                 invokeStatic(TraceDebuggerInjections::isFirstRun)
             },
-            ifClause = {
+            thenClause = {
                 push(classUnderTraceDebugging)
                 push(methodUnderTraceDebugging)
                 // STACK: testClassName, testMethodName

src/jvm/main/org/jetbrains/kotlinx/lincheck/transformation/TransformationUtils.kt

@@ -273,7 +273,7 @@ internal fun GeneratorAdapter.invokeBeforeEvent(debugMessage: String, setMethodE
         condition = {
             invokeStatic(Injections::shouldInvokeBeforeEvent)
         },
-        ifClause = {
+        thenClause = {
             if (setMethodEventId) {
                 invokeStatic(Injections::setLastMethodCallEventId)
             }
@@ -284,7 +284,7 @@ internal fun GeneratorAdapter.invokeBeforeEvent(debugMessage: String, setMethodE
                 condition = {
                     invokeStatic(Injections::isBeforeEventRequested)
                 },
-                ifClause = {
+                thenClause = {
                     push(debugMessage)
                     invokeStatic(Injections::beforeEvent)
                 },
@@ -316,13 +316,13 @@ internal fun GeneratorAdapter.invokeStatic(function: KFunction<*>) {
  * Generates an if-statement in bytecode.
  *
  * @param condition the condition code.
- * @param ifClause the if-clause code.
+ * @param thenClause the then-clause code.
  * @param elseClause the else-clause code.
  */
 internal inline fun GeneratorAdapter.ifStatement(
     condition: GeneratorAdapter.() -> Unit,
-    ifClause: GeneratorAdapter.() -> Unit,
-    elseClause: GeneratorAdapter.() -> Unit
+    thenClause: GeneratorAdapter.() -> Unit,
+    elseClause: GeneratorAdapter.() -> Unit = { },
 ) {
     val ifClauseStart = newLabel()
     val end = newLabel()
@@ -331,7 +331,7 @@ internal inline fun GeneratorAdapter.ifStatement(
     elseClause()
     goTo(end)
     visitLabel(ifClauseStart)
-    ifClause()
+    thenClause()
     visitLabel(end)
 }
 
@@ -351,7 +351,7 @@ internal inline fun GeneratorAdapter.invokeIfInTestingCode(
 ) {
     ifStatement(
         condition = { invokeStatic(Injections::inIgnoredSection) },
-        ifClause = original,
+        thenClause = original,
         elseClause = code
     )
 }
@@ -373,7 +373,7 @@ internal inline fun GeneratorAdapter.invokeInIgnoredSection(
         condition = {
             loadLocal(enteredIgnoredSection)
         },
-        ifClause = {
+        thenClause = {
             invokeStatic(Injections::leaveIgnoredSection)
         },
         elseClause = {}

src/jvm/main/org/jetbrains/kotlinx/lincheck/transformation/transformers/DeterminismTransformers.kt

@@ -17,7 +17,9 @@ import org.objectweb.asm.Type.getType
 import org.objectweb.asm.commons.Method
 import org.objectweb.asm.commons.GeneratorAdapter
 import org.jetbrains.kotlinx.lincheck.transformation.*
+import org.objectweb.asm.Opcodes
 import sun.nio.ch.lincheck.*
+import java.lang.reflect.Modifier
 import java.util.*
 
 /**
@@ -126,16 +128,25 @@ internal abstract class AbstractDeterministicTimeMethodTransformer(val adapter:
     }
 }
 
-/**
- * [DeterministicRandomTransformer] tracks invocations of various random number generation functions,
- * such as [Random.nextInt], and replaces them with corresponding [Injections] methods to prevent non-determinism.
- */
-internal class DeterministicRandomTransformer(
+internal abstract class AbstractDeterministicRandomTransformer(
     fileName: String,
     className: String,
     methodName: String,
     adapter: GeneratorAdapter,
 ) : ManagedStrategyMethodVisitor(fileName, className, methodName, adapter) {
+    protected abstract fun GeneratorAdapter.generateInstrumentedCodeForNextSecondarySeedAndGetProbe(
+        opcode: Int, owner: String, name: String, desc: String, itf: Boolean,
+    )
+
+    protected abstract fun GeneratorAdapter.generateInstrumentedCodeForAdvanceProbe(
+        opcode: Int, owner: String, name: String, desc: String, itf: Boolean,
+    )
+
+    protected abstract fun GeneratorAdapter.generateInstrumentedCodeForRegularRandomMethod(
+        opcode: Int, owner: String, name: String, desc: String, itf: Boolean,
+        receiverLocal: Int,
+        argumentsLocals: IntArray,
+    )
 
     override fun visitMethodInsn(opcode: Int, owner: String, name: String, desc: String, itf: Boolean) = adapter.run {
         if (owner == "java/util/concurrent/ThreadLocalRandom" ||
@@ -151,7 +162,7 @@ internal class DeterministicRandomTransformer(
                         visitMethodInsn(opcode, owner, name, desc, itf)
                     },
                     code = {
-                        invokeStatic(Injections::nextInt)
+                        generateInstrumentedCodeForNextSecondarySeedAndGetProbe(opcode, owner, name, desc, itf)
                     }
                 )
                 return
@@ -162,28 +173,13 @@ internal class DeterministicRandomTransformer(
                         visitMethodInsn(opcode, owner, name, desc, itf)
                     },
                     code = {
-                        pop()
-                        invokeStatic(Injections::nextInt)
-                    }
-                )
-                return
-            }
-            if (name == "nextInt" && desc == "(II)I") {
-                invokeIfInTestingCode(
-                    original = {
-                        visitMethodInsn(opcode, owner, name, desc, itf)
-                    },
-                    code = {
-                        val arguments = storeArguments(desc)
-                        pop()
-                        loadLocals(arguments)
-                        invokeStatic(Injections::nextInt2)
+                        generateInstrumentedCodeForAdvanceProbe(opcode, owner, name, desc, itf)
                     }
                 )
                 return
             }
         }
-        if (isRandomMethod(name, desc)) {
+        if (isRandomMethod(opcode, name, desc)) {
             invokeIfInTestingCode(
                 original = {
                     visitMethodInsn(opcode, owner, name, desc, itf)
@@ -197,17 +193,8 @@ internal class DeterministicRandomTransformer(
                             loadLocal(ownerLocal)
                             invokeStatic(Injections::isRandom)
                         },
-                        ifClause = {
-                            invokeInIgnoredSection {
-                                invokeStatic(Injections::deterministicRandom)
-                                loadLocals(arguments)
-                                /*
-                                 * In Java 21 RandomGenerator interface was introduced,
-                                 * so sometimes data structures interact with java.util.Random through this interface.
-                                 */
-                                val randomOwner = if (owner.endsWith("RandomGenerator")) "java/util/random/RandomGenerator" else "java/util/Random"
-                                visitMethodInsn(opcode, randomOwner, name, desc, itf)
-                            }
+                        thenClause = {
+                            generateInstrumentedCodeForRegularRandomMethod(opcode, owner, name, desc, itf, ownerLocal, arguments)
                         },
                         elseClause = {
                             loadLocal(ownerLocal)
@@ -222,12 +209,75 @@ internal class DeterministicRandomTransformer(
         visitMethodInsn(opcode, owner, name, desc, itf)
     }
 
-    private fun isRandomMethod(methodName: String, desc: String): Boolean = RANDOM_METHODS.any {
-        it.name == methodName && it.descriptor == desc
+    private fun isRandomMethod(opcode: Int, methodName: String, desc: String): Boolean = allRandomMethods.any {
+        opcode.isInstanceMethodOpcode && it.name == methodName && it.descriptor == desc
     }
 
     private companion object {
-        private val RANDOM_METHODS = Random::class.java.declaredMethods.map { Method.getMethod(it) }
+        private val randomGeneratorClass = runCatching { Class.forName("java.util.random.RandomGenerator") }.getOrNull()
+        private fun Class<*>.getMethodsToReplace() = declaredMethods
+            .filter { Modifier.isPublic(it.modifiers) || Modifier.isProtected(it.modifiers) }
+            .map { Method.getMethod(it) }
+        private val randomGeneratorMethods = randomGeneratorClass?.getMethodsToReplace() ?: emptyList()
+        private val randomClassMethods = Random::class.java.getMethodsToReplace()
+        private val allRandomMethods = randomGeneratorMethods + randomClassMethods
+        private val Int.isInstanceMethodOpcode
+            get() = when (this) {
+                Opcodes.INVOKEVIRTUAL, Opcodes.INVOKEINTERFACE, Opcodes.INVOKESPECIAL -> true
+                else -> false
+            }
     }
 }
 
+/**
+ * [FakeDeterministicRandomTransformer] tracks invocations of various random number generation functions,
+ * such as [Random.nextInt], and replaces them with corresponding [Injections] methods to prevent non-determinism.
+ */
+internal class FakeDeterministicRandomTransformer(
+    fileName: String,
+    className: String,
+    methodName: String,
+    adapter: GeneratorAdapter,
+) : AbstractDeterministicRandomTransformer(fileName, className, methodName, adapter) {
+    override fun GeneratorAdapter.generateInstrumentedCodeForNextSecondarySeedAndGetProbe(
+        opcode: Int,
+        owner: String,
+        name: String,
+        desc: String,
+        itf: Boolean
+    ) {
+        invokeStatic(Injections::nextInt)
+    }
+
+    override fun GeneratorAdapter.generateInstrumentedCodeForAdvanceProbe(
+        opcode: Int,
+        owner: String,
+        name: String,
+        desc: String,
+        itf: Boolean
+    ) {
+        pop()
+        invokeStatic(Injections::nextInt)
+    }
+
+    override fun GeneratorAdapter.generateInstrumentedCodeForRegularRandomMethod(
+        opcode: Int,
+        owner: String,
+        name: String,
+        desc: String,
+        itf: Boolean,
+        receiverLocal: Int,
+        argumentsLocals: IntArray,
+    ) {
+        invokeInIgnoredSection {
+            invokeStatic(Injections::deterministicRandom)
+            loadLocals(argumentsLocals)
+            /*
+             * In Java 21 RandomGenerator interface was introduced,
+             * so sometimes data structures interact with java.util.Random through this interface.
+             */
+            val randomOwner = if (owner.endsWith("RandomGenerator")) "java/util/random/RandomGenerator" else "java/util/Random"
+            visitMethodInsn(opcode, randomOwner, name, desc, itf)
+        }
+    }
+}

src/jvm/main/org/jetbrains/kotlinx/lincheck/transformation/transformers/SharedMemoryAccessTransformer.kt

@@ -63,7 +63,7 @@ internal class SharedMemoryAccessTransformer(
                         // STACK: isTracePointCreated
                         ifStatement(
                             condition = { /* already on stack */ },
-                            ifClause = {
+                            thenClause = {
                                 invokeBeforeEventIfPluginEnabled("read static field")
                             },
                             elseClause = {})
@@ -96,7 +96,7 @@ internal class SharedMemoryAccessTransformer(
                         // STACK: obj, isTracePointCreated
                         ifStatement(
                             condition = { /* already on stack */ },
-                            ifClause = {
+                            thenClause = {
                                 invokeBeforeEventIfPluginEnabled("read field")
                             },
                             elseClause = {}
@@ -134,7 +134,7 @@ internal class SharedMemoryAccessTransformer(
                         // STACK: isTracePointCreated
                         ifStatement(
                             condition = { /* already on stack */ },
-                            ifClause = {
+                            thenClause = {
                                 invokeBeforeEventIfPluginEnabled("write static field")
                             },
                             elseClause = {}
@@ -172,7 +172,7 @@ internal class SharedMemoryAccessTransformer(
                         // STACK: isTracePointCreated
                         ifStatement(
                             condition = { /* already on stack */ },
-                            ifClause = {
+                            thenClause = {
                                 invokeBeforeEventIfPluginEnabled("write field")
                             },
                             elseClause = {}
@@ -211,7 +211,7 @@ internal class SharedMemoryAccessTransformer(
                         invokeStatic(Injections::beforeReadArray)
                         ifStatement(
                             condition = { /* already on stack */ },
-                            ifClause = {
+                            thenClause = {
                                 invokeBeforeEventIfPluginEnabled("read array")
                             },
                             elseClause = {}
@@ -245,7 +245,7 @@ internal class SharedMemoryAccessTransformer(
                         invokeStatic(Injections::beforeWriteArray)
                         ifStatement(
                             condition = { /* already on stack */ },
-                            ifClause = {
+                            thenClause = {
                                 invokeBeforeEventIfPluginEnabled("write array")
                             },
                             elseClause = {}

src/jvm/test/org/jetbrains/kotlinx/lincheck_test/trace_debugger/AbstractDeterministicTest.kt

@@ -27,8 +27,12 @@ import java.io.PrintStream
  * Checks for absence of non-determinism and absence (or existence) of exceptions.
  */
 abstract class AbstractDeterministicTest {
+    open val alsoRunInLincheckMode: Boolean get() = false
+
     private fun testTraceDebugger() {
-        assumeTrue(isInTraceDebuggerMode)
+        if (!alsoRunInLincheckMode) {
+            assumeTrue(isInTraceDebuggerMode)
+        }
         val oldStdOut = System.out
         val oldErr = System.err
         val stdOutOutputCollector = ByteArrayOutputStream()
@@ -51,9 +55,9 @@ abstract class AbstractDeterministicTest {
                     val results = lincheckFailure?.results?.parallelResults?.flatten()?.takeIf { it.isNotEmpty() }
                     require(results != null) { lincheckFailure.toString() }
                     if (shouldFail()) {
-                        require(results.all { it is ExceptionResult })// { lincheckFailure.toString() }
+                        require(results.all { it is ExceptionResult }) { lincheckFailure.toString() }
                     } else {
-                        require(results.none { it is ExceptionResult })// { lincheckFailure.toString() }
+                        require(results.none { it is ExceptionResult }) { lincheckFailure.toString() }
                     }
                 }
         } finally {