[JVM] Adapt invocation results to call-site return types

usvm-core/src/main/kotlin/org/usvm/memory/RegistersStack.kt

@@ -45,8 +45,13 @@ class URegistersStack(
     fun push(arguments: Array<UExpr<out USort>>, localsCount: Int) =
         frames.add(arguments.copyOf(arguments.size + localsCount))
 
-    private fun <Sort : USort> Array<UExpr<out USort>?>?.read(index: Int, sort: Sort): UExpr<Sort> =
-        this?.get(index)?.asExpr(sort) ?: sort.uctx.mkRegisterReading(index, sort)
+    private fun <Sort : USort> Array<UExpr<out USort>?>?.read(index: Int, sort: Sort): UExpr<Sort> {
+        val value = this?.get(index) ?: return sort.uctx.mkRegisterReading(index, sort)
+        check(value.sort == sort) {
+            "Register sort mismatch: index=$index requested=$sort stored=${value.sort} frames=${frames.size}"
+        }
+        return value.asExpr(sort)
+    }
 
     override fun <Sort : USort> readRegister(index: Int, sort: Sort): UExpr<Sort> =
         frames.lastOrNull().read(index, sort)

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

@@ -77,6 +77,7 @@ import org.jacodb.api.jvm.ext.byte
 import org.jacodb.api.jvm.ext.char
 import org.jacodb.api.jvm.ext.double
 import org.jacodb.api.jvm.ext.enumValues
+import org.jacodb.api.jvm.ext.findType
 import org.jacodb.api.jvm.ext.float
 import org.jacodb.api.jvm.ext.ifArrayGetElementType
 import org.jacodb.api.jvm.ext.int
@@ -120,6 +121,7 @@ import org.usvm.machine.state.JcState
 import org.usvm.machine.state.addConcreteMethodCallStmt
 import org.usvm.machine.state.addDynamicCall
 import org.usvm.machine.state.addVirtualMethodCallStmt
+import org.usvm.machine.state.lastStmt
 import org.usvm.machine.state.throwExceptionWithoutStackFrameDrop
 import org.usvm.memory.ULValue
 import org.usvm.memory.URegisterStackLValue
@@ -479,17 +481,20 @@ class JcExprResolver(
             arguments += resolveJcExpr(expr, type) ?: return null
         }
 
-        return resolveInvokeNoStaticInitializationCheck { onNoCallPresent(arguments) }
+        return resolveInvokeNoStaticInitializationCheck(expectedReturnType = method.returnType) {
+            onNoCallPresent(arguments)
+        }
     }
 
     private inline fun resolveInvokeNoStaticInitializationCheck(
+        expectedReturnType: JcType,
         onNoCallPresent: JcStepScope.() -> Unit,
     ): UExpr<out USort>? {
         val result = scope.calcOnState { methodResult }
         return when (result) {
             is JcMethodResult.Success -> {
                 scope.doWithState { methodResult = JcMethodResult.NoCall }
-                result.value
+                adaptInvocationResult(result, expectedReturnType)
             }
 
             is JcMethodResult.NoCall -> {
@@ -501,6 +506,22 @@ class JcExprResolver(
         }
     }
 
+    /** Lambda implementation and functional-interface results may have different JVM representations. */
+    private fun adaptInvocationResult(result: JcMethodResult.Success, expectedType: JcType): UExpr<out USort>? {
+        if (expectedType == ctx.cp.void) return ctx.voidValue
+        val actualType = ctx.cp.findType(result.method.returnType.typeName)
+        val requiresBoxing = actualType is JcPrimitiveType && expectedType !is JcPrimitiveType
+        if (!requiresBoxing) return resolveLambdaArgument(result.value, actualType, expectedType)
+
+        val wrapper = actualType.autoboxIfNeeded() as JcClassType
+        val valueOf = wrapper.declaredMethods.single {
+            it.name == "valueOf" && it.isStatic && it.parameters.singleOrNull()?.type == actualType
+        }
+        // Execute the ordinary boxing method so wrapper caching and its models keep their semantics.
+        scope.doWithState { addConcreteMethodCallStmt(valueOf.method, listOf(result.value)) }
+        return null
+    }
+
     // endregion
 
     // region jc locals
@@ -1222,7 +1243,15 @@ class JcSimpleValueResolver(
 
     override fun visitJcLocalVar(value: JcLocalVar): UExpr<out USort> = with(ctx) {
         val ref = resolveLocal(value)
-        scope.calcOnState { memory.read(ref) }
+        scope.calcOnState {
+            try {
+                memory.read(ref)
+            } catch (failure: IllegalStateException) {
+                val context = "entry=$entrypoint method=$lastEnteredMethod statement=$lastStmt " +
+                    "local=${value.name} type=${value.type} register=${ref.idx} requested=${ref.sort}"
+                throw IllegalStateException(context, failure)
+            }
+        }
     }
 
     override fun visitJcThis(value: JcThis): UExpr<out USort> = with(ctx) {

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

@@ -0,0 +1,40 @@
+package org.usvm.samples.lambda;
+
+import java.util.function.Function;
+import java.util.function.IntFunction;
+import java.util.function.IntToLongFunction;
+import java.util.function.ToIntFunction;
+
+public class LambdaReturnAdaptation {
+    public static boolean boxedBoolean(int value) {
+        Function<Integer, Boolean> positive = LambdaReturnAdaptation::positive;
+        return positive.apply(value);
+    }
+
+    public static int boxedInteger(int value) {
+        IntFunction<Integer> identity = LambdaReturnAdaptation::identity;
+        return identity.apply(value);
+    }
+
+    public static int unboxedInteger(Integer value) {
+        ToIntFunction<Integer> identity = LambdaReturnAdaptation::boxedIdentity;
+        return identity.applyAsInt(value);
+    }
+
+    public static long widenedInteger(int value) {
+        IntToLongFunction identity = LambdaReturnAdaptation::identity;
+        return identity.applyAsLong(value);
+    }
+
+    private static boolean positive(int value) {
+        return value > 0;
+    }
+
+    private static int identity(int value) {
+        return value;
+    }
+
+    private static Integer boxedIdentity(Integer value) {
+        return value;
+    }
+}

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

@@ -0,0 +1,51 @@
+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 LambdaReturnAdaptationTest : ApproximationsTestRunner() {
+    init {
+        options = options.copy(throwExceptionOnStepFailure = true)
+    }
+
+    @Test
+    fun `primitive boolean result is boxed for erased Function return`() {
+        checkDiscoveredProperties(
+            LambdaReturnAdaptation::boxedBoolean,
+            eq(count = 2),
+            { value, result -> value > 0 && result == true },
+            { value, result -> value <= 0 && result == false },
+        )
+    }
+
+    @Test
+    fun `primitive integer result is boxed for IntFunction return`() {
+        checkDiscoveredProperties(
+            LambdaReturnAdaptation::boxedInteger,
+            eq(count = 1),
+            { value, result -> result == value },
+        )
+    }
+
+    @Test
+    fun `boxed result is unboxed and null preserves NullPointerException`() {
+        checkDiscoveredPropertiesWithExceptions(
+            LambdaReturnAdaptation::unboxedInteger,
+            eq(count = 2),
+            { value, result -> value == null && result.isException<NullPointerException>() },
+            { value, result -> value != null && result.getOrNull() == value },
+        )
+    }
+
+    @Test
+    fun `primitive result widens to functional interface return type`() {
+        checkDiscoveredProperties(
+            LambdaReturnAdaptation::widenedInteger,
+            eq(count = 1),
+            { value, result -> result == value.toLong() },
+        )
+    }
+
+}