[JVM] Unbox lambda invocation arguments before implementation dispatch

buildSrc/src/main/kotlin/Dependencies.kt

@@ -6,7 +6,7 @@ object Versions {
     const val clikt = "5.0.0"
     const val detekt = "1.23.7"
     const val ini4j = "0.5.4"
-    const val jacodb = "4251d2023d"
+    const val jacodb = "1822dea851"
     const val juliet = "1.3.2"
     const val junit = "5.9.3"
     const val kotlin = "2.1.0"

usvm-jvm/src/main/kotlin/org/usvm/machine/JcMethodCallInst.kt

@@ -11,6 +11,7 @@ import org.jacodb.impl.cfg.JcInstLocationImpl
 import org.usvm.UExpr
 import org.usvm.UHeapRef
 import org.usvm.USort
+import org.usvm.machine.interpreter.JcLambdaCallSite
 
 /**
  * An interface for instructions that replace or surround some [originalInst].
@@ -102,3 +103,13 @@ data class JcDynamicMethodCallInst(
     override val method: JcMethod = dynamicCall.method.method
     override val originalInst: JcInst = returnSite
 }
+
+/** Adapts a lambda's invocation arguments in their own step, after receiver dispatch has forked. */
+data class JcLambdaMethodCallInst(
+    val virtualCall: JcVirtualMethodCallInst,
+    val callSite: JcLambdaCallSite,
+) : JcMethodCallBaseInst {
+    override val method: JcMethod = callSite.lambda.actualMethod.method.method
+    override val location: JcInstLocation = virtualCall.location
+    override val originalInst: JcInst = virtualCall.originalInst
+}

usvm-jvm/src/main/kotlin/org/usvm/machine/interpreter/JcExprResolver.kt

@@ -71,6 +71,7 @@ import org.jacodb.api.jvm.cfg.JcValue
 import org.jacodb.api.jvm.cfg.JcValueVisitor
 import org.jacodb.api.jvm.cfg.JcVirtualCallExpr
 import org.jacodb.api.jvm.cfg.JcXorExpr
+import org.jacodb.api.jvm.ext.autoboxIfNeeded
 import org.jacodb.api.jvm.ext.boolean
 import org.jacodb.api.jvm.ext.byte
 import org.jacodb.api.jvm.ext.char
@@ -980,6 +981,30 @@ class JcExprResolver(
         return expr
     }
 
+    internal fun resolveLambdaArgument(
+        value: UExpr<out USort>,
+        argumentType: JcType,
+        parameterType: JcType,
+    ): UExpr<out USort>? {
+        if (parameterType !is JcPrimitiveType) return value
+        if (argumentType is JcPrimitiveType) return resolvePrimitiveCast(value, argumentType, parameterType)
+
+        val unboxedType = ctx.primitiveTypes.firstOrNull { primitive ->
+            primitive != ctx.cp.void && primitive.autoboxIfNeeded().typeName == argumentType.typeName
+        } ?: error("Lambda unboxing requires a primitive wrapper, got ${argumentType.typeName}")
+        val wrapperType = unboxedType.autoboxIfNeeded() as JcClassType
+        val reference = value.asExpr(ctx.addressSort)
+        checkNullPointer(reference) ?: return null
+        checkClassCast(reference, wrapperType) ?: return null
+
+        // The primitive wrappers are final, and their unboxing methods read this final field.
+        val valueField = wrapperType.declaredFields.single { it.name == "value" }
+        val valueSlot = UFieldLValue(ctx.typeToSort(unboxedType), reference, valueField.field)
+        val unboxed = scope.calcOnState { memory.read(valueSlot) }
+
+        return resolvePrimitiveCast(unboxed, unboxedType, parameterType)
+    }
+
     private fun resolvePrimitiveCast(
         expr: JcExpr,
         type: JcPrimitiveType

usvm-jvm/src/main/kotlin/org/usvm/machine/interpreter/JcInterpreter.kt

@@ -48,6 +48,7 @@ import org.usvm.jvm.util.name
 import org.usvm.jvm.util.outerClassInstanceField
 import org.usvm.machine.JcApplicationGraph
 import org.usvm.machine.JcConcreteMethodCallInst
+import org.usvm.machine.JcLambdaMethodCallInst
 import org.usvm.machine.JcContext
 import org.usvm.machine.JcDynamicMethodCallInst
 import org.usvm.machine.JcInterpreterObserver
@@ -265,6 +266,11 @@ class JcInterpreter(
                 }
             }
 
+            is JcLambdaMethodCallInst -> {
+                val concreteCall = resolveLambdaInvoke(stmt, scope, exprResolver) ?: return
+                scope.doWithState { newStmt(concreteCall) }
+            }
+
             is JcConcreteMethodCallInst -> {
                 observer?.onMethodCallWithResolvedArguments(simpleValueResolver, stmt, scope)
                 if (approximateMethod(scope, stmt)) {

usvm-jvm/src/main/kotlin/org/usvm/machine/interpreter/JcVirtualInvokeResolver.kt

@@ -7,6 +7,7 @@ import org.jacodb.api.jvm.JcClassOrInterface
 import org.jacodb.api.jvm.JcRefType
 import org.jacodb.api.jvm.JcType
 import org.jacodb.api.jvm.ext.enumValues
+import org.jacodb.api.jvm.ext.findType
 import org.jacodb.api.jvm.ext.isEnum
 import org.jacodb.api.jvm.ext.toType
 import org.usvm.UBoolExpr
@@ -19,6 +20,7 @@ import org.usvm.isAllocatedConcreteHeapRef
 import org.usvm.isStaticHeapRef
 import org.usvm.machine.JcConcreteMethodCallInst
 import org.usvm.machine.JcContext
+import org.usvm.machine.JcLambdaMethodCallInst
 import org.usvm.machine.JcVirtualMethodCallInst
 import org.usvm.machine.interpreter.statics.JcStaticFieldLValue
 import org.usvm.machine.state.JcState
@@ -76,7 +78,7 @@ private fun resolveVirtualInvokeWithModel(
     if (isAllocatedConcreteHeapRef(concreteRef) || isStaticHeapRef(concreteRef)) {
         val callSite = findLambdaCallSite(methodCall, scope, concreteRef)
         if (callSite != null) {
-            val lambdaCall = makeLambdaCallSiteCall(scope, callSite)
+            val lambdaCall = JcLambdaMethodCallInst(virtualCall = methodCall, callSite = callSite)
             scope.doWithState {
                 newStmt(lambdaCall)
             }
@@ -233,7 +235,7 @@ private fun resolveVirtualInvokeWithoutModel(
     }
 
     lambdaCallSitesWithConditions.mapTo(conditionsWithBlocks) { (callSite, condition) ->
-        val concreteCall = makeLambdaCallSiteCall(scope, callSite)
+        val concreteCall = JcLambdaMethodCallInst(virtualCall = methodCall, callSite = callSite)
         condition to { state: JcState -> state.newStmt(concreteCall) }
     }
 
@@ -300,7 +302,7 @@ private fun findLambdaCallSite(
     val callSites = scope.calcOnState { memory.getRegion(ctx.lambdaCallSiteRegionId) as JcLambdaCallSiteMemoryRegion }
     val callSite = callSites.findCallSite(ref) ?: return null
 
-    val lambdaMethodType = callSite.lambda.dynamicMethodType
+    val lambdaMethodType = callSite.lambda.interfaceMethodType
 
     // Match function signature
     when {
@@ -312,22 +314,40 @@ private fun findLambdaCallSite(
     return callSite
 }
 
-private fun JcVirtualMethodCallInst.makeLambdaCallSiteCall(
+internal fun resolveLambdaInvoke(
+    instruction: JcLambdaMethodCallInst,
     scope: JcStepScope,
-    callSite: JcLambdaCallSite,
-): JcConcreteMethodCallInst {
+    exprResolver: JcExprResolver,
+): JcConcreteMethodCallInst? {
+    val call = instruction.virtualCall
+    val callSite = instruction.callSite
     val lambda = callSite.lambda
     val lambdaMethod = lambda.actualMethod.method
-    // Instance was already resolved to the call site
-    var callArgsWithoutInstance = this.arguments.drop(1)
+    val classpath = scope.calcOnState { ctx.cp }
+    val invocationArgumentTypes = lambda.dynamicMethodType.argumentTypes.map { classpath.findType(it.typeName) }
+    var arguments = callSite.callSiteArgs + call.arguments.drop(1)
+    var argumentTypes = lambda.callSiteArgTypes + invocationArgumentTypes
+    var parameterTypes = lambdaMethod.parameters.map { it.type }
+
+    if (!lambdaMethod.method.isStatic) {
+        parameterTypes = listOf(lambdaMethod.enclosingType) + parameterTypes
+    }
+
     if (lambda.isNewInvokeSpecial) {
         check(lambdaMethod.method.isConstructor) {
-            "unexpected JcLambdaExpr: calling NEWINVOKESPECIAL for non-constructor method"
+            "Unexpected constructor reference to ${lambdaMethod.method}"
         }
-        val thisArg = scope.calcOnState { memory.allocConcrete(lambdaMethod.enclosingType) }
-        callArgsWithoutInstance = listOf(thisArg) + callArgsWithoutInstance
+        val receiver = scope.calcOnState { memory.allocConcrete(lambdaMethod.enclosingType) }
+        arguments = listOf(receiver) + arguments
+        argumentTypes = listOf(lambdaMethod.enclosingType) + argumentTypes
+    }
+
+    require(arguments.size == argumentTypes.size && arguments.size == parameterTypes.size) {
+        "Lambda argument count differs from implementation ${lambdaMethod.method}"
+    }
+    val adaptedArguments = arguments.indices.map { index ->
+        exprResolver.resolveLambdaArgument(arguments[index], argumentTypes[index], parameterTypes[index]) ?: return null
     }
-    val lambdaMethodArgs = callSite.callSiteArgs + callArgsWithoutInstance
 
-    return JcConcreteMethodCallInst(location, lambdaMethod.method, lambdaMethodArgs, returnSite)
+    return JcConcreteMethodCallInst(call.location, lambdaMethod.method, adaptedArguments, call.returnSite)
 }

usvm-jvm/src/samples/java/org/usvm/samples/lambda/LambdaArgumentAdaptation.java

@@ -0,0 +1,38 @@
+package org.usvm.samples.lambda;
+
+import java.util.function.IntPredicate;
+import java.util.function.Predicate;
+
+public class LambdaArgumentAdaptation {
+    public static int isWhitespace(int codePoint) {
+        if (codePoint != ' ' && codePoint != 'A') {
+            return -1;
+        }
+
+        Predicate<Integer> isDelimiter = Character::isWhitespace;
+        return isDelimiter.test(codePoint) ? 1 : 0;
+    }
+
+    public static boolean boxedPositive(Integer value) {
+        Predicate<Integer> positive = LambdaArgumentAdaptation::positiveInt;
+        return positive.test(value);
+    }
+
+    public static boolean boxedWideningPositive(Integer value) {
+        Predicate<Integer> positive = LambdaArgumentAdaptation::positiveLong;
+        return positive.test(value);
+    }
+
+    public static boolean primitivePositive(int value) {
+        IntPredicate positive = LambdaArgumentAdaptation::positiveInt;
+        return positive.test(value);
+    }
+
+    private static boolean positiveInt(int value) {
+        return value > 0;
+    }
+
+    private static boolean positiveLong(long value) {
+        return value > 0;
+    }
+}

usvm-jvm/src/test/kotlin/org/usvm/samples/lambda/LambdaArgumentAdaptationTest.kt

@@ -0,0 +1,55 @@
+package org.usvm.samples.lambda
+
+import org.junit.jupiter.api.Test
+import org.usvm.samples.approximations.ApproximationsTestRunner
+import org.usvm.test.util.checkers.eq
+import org.usvm.util.isException
+
+class LambdaArgumentAdaptationTest : ApproximationsTestRunner() {
+    init {
+        options = options.copy(throwExceptionOnStepFailure = true)
+    }
+
+    @Test
+    fun `boxed predicate argument is unboxed for the primitive method reference`() {
+        checkDiscoveredProperties(
+            LambdaArgumentAdaptation::isWhitespace,
+            eq(count = 3),
+            { value, result -> value == ' '.code && result == 1 },
+            { value, result -> value == 'A'.code && result == 0 },
+            { value, result -> value != ' '.code && value != 'A'.code && result == -1 },
+        )
+    }
+
+    @Test
+    fun `unboxing a nullable lambda argument preserves NullPointerException`() {
+        checkDiscoveredPropertiesWithExceptions(
+            LambdaArgumentAdaptation::boxedPositive,
+            eq(count = 3),
+            { value, result -> value == null && result.isException<NullPointerException>() },
+            { value, result -> value != null && value > 0 && result.getOrNull() == true },
+            { value, result -> value != null && value <= 0 && result.getOrNull() == false },
+        )
+    }
+
+    @Test
+    fun `unboxing a lambda argument widens to the implementation primitive type`() {
+        checkDiscoveredPropertiesWithExceptions(
+            LambdaArgumentAdaptation::boxedWideningPositive,
+            eq(count = 3),
+            { value, result -> value == null && result.isException<NullPointerException>() },
+            { value, result -> value != null && value > 0 && result.getOrNull() == true },
+            { value, result -> value != null && value <= 0 && result.getOrNull() == false },
+        )
+    }
+
+    @Test
+    fun `primitive lambda arguments keep their values`() {
+        checkDiscoveredProperties(
+            LambdaArgumentAdaptation::primitivePositive,
+            eq(count = 2),
+            { value, result -> value > 0 && result == true },
+            { value, result -> value <= 0 && result == false },
+        )
+    }
+}