Preserve extracted native Object.assign identity

usvm-ts/src/main/kotlin/org/usvm/machine/expr/NativeObjectAssignIdentity.kt

@@ -0,0 +1,62 @@
+package org.usvm.machine.expr
+
+import org.jacodb.ets.model.EtsInstanceFieldRef
+import org.jacodb.ets.model.EtsMethodSignature
+import org.jacodb.ets.model.EtsUnknownType
+import org.usvm.UConcreteHeapRef
+import org.usvm.isFalse
+import org.usvm.machine.TsContext
+import org.usvm.machine.runtimeFunctionType
+
+/** Object.assign ignores its this argument, so an extracted native function keeps the same call boundary. */
+internal fun TsExprResolver.tryCaptureNativeObjectAssign(field: EtsInstanceFieldRef): TsExprApproximationResult {
+    if (field.field.name != "assign") return TsExprApproximationResult.NoApproximation
+    val receiver = resolve(field.instance) ?: return TsExprApproximationResult.ResolveFailure
+    if (receiver !is UConcreteHeapRef) return TsExprApproximationResult.NoApproximation
+    if (!ctx.isNativeObjectConstructor(receiver)) return TsExprApproximationResult.NoApproximation
+
+    val changed = scope.calcOnState {
+        val written = memory.read(ctx.writtenPropertyLValue(receiver, field.field.name))
+        val deleted = memory.read(ctx.deletedFieldLValue(receiver, field.field))
+        !written.isFalse || !deleted.isFalse
+    }
+    if (changed) {
+        val stored = ctx.resolveField(scope, field.instance, receiver, field.field, hierarchy)
+            ?: return TsExprApproximationResult.ResolveFailure
+        return TsExprApproximationResult.SuccessfulApproximation(stored)
+    }
+
+    val owner = requireNotNull(ctx.classConstructorSignature(receiver))
+    val prior = scope.calcOnState {
+        capturedInstanceCalls.firstOrNull { (_, captured) ->
+            val native = captured.nativeCallee
+            native != null && native.enclosingClass == owner && native.name == field.field.name
+        }?.key
+    }
+    if (prior != null) return TsExprApproximationResult.SuccessfulApproximation(prior)
+
+    val callee = EtsMethodSignature(
+        enclosingClass = owner,
+        name = field.field.name,
+        parameters = emptyList(),
+        returnType = EtsUnknownType,
+    )
+    val pointer = scope.calcOnState { memory.allocConcrete(ctx.scene.runtimeFunctionType()) }
+    scope.doWithState {
+        capturedInstanceCalls = capturedInstanceCalls.put(
+            pointer,
+            TsCapturedInstanceCall(
+                receiver = receiver,
+                sourceReceiver = field.instance,
+                nativeCallee = callee,
+            ),
+            ownership,
+        )
+    }
+    return TsExprApproximationResult.SuccessfulApproximation(pointer)
+}
+
+private fun TsContext.isNativeObjectConstructor(instance: UConcreteHeapRef): Boolean {
+    val signature = classConstructorSignature(instance) ?: return false
+    return scene.sdkClasses.any { clazz -> clazz.name == "Object" && clazz.signature == signature }
+}

usvm-ts/src/main/kotlin/org/usvm/machine/expr/TsCapturedInstanceCall.kt

@@ -30,11 +30,18 @@ import org.usvm.util.mkFieldLValue
 class TsCapturedInstanceCall(
     val receiver: UExpr<*>,
     val sourceReceiver: EtsLocal,
+    /** Extracted native methods retain their owner/name; each invocation keeps its own argument and result types. */
+    val nativeCallee: EtsMethodSignature? = null,
 )
 
 internal fun TsExprResolver.tryCaptureInstanceCall(
     field: EtsInstanceFieldRef,
 ): TsExprApproximationResult = with(ctx) {
+    when (val native = tryCaptureNativeObjectAssign(field)) {
+        TsExprApproximationResult.NoApproximation -> {}
+        else -> return native
+    }
+
     val method = scope.calcOnState { lastEnteredMethod }
     val assignment = scope.calcOnState { lastStmt } as? EtsAssignStmt
         ?: return TsExprApproximationResult.NoApproximation
@@ -107,10 +114,16 @@ internal fun TsExprResolver.handleCapturedInstanceCall(
     call: EtsPtrCallExpr,
     captured: TsCapturedInstanceCall,
 ): UExpr<out USort>? = with(ctx) {
-    val receiver = requireNotNull(call.receiver)
+    val nativeCallee = captured.nativeCallee
+    val receiver = if (nativeCallee != null) captured.sourceReceiver else requireNotNull(call.receiver)
+    val callee = if (nativeCallee != null) {
+        call.callee.copy(enclosingClass = nativeCallee.enclosingClass, name = nativeCallee.name)
+    } else {
+        call.callee
+    }
     val instanceCall = EtsInstanceCallExpr(
         instance = receiver,
-        callee = call.callee,
+        callee = callee,
         args = call.args,
         type = call.type,
     )
@@ -131,7 +144,7 @@ internal fun TsExprResolver.handleCapturedInstanceCall(
         call = instanceCall,
         instance = captured.receiver,
         args = args,
-        originalCall = call.copy(receiver = captured.sourceReceiver),
+        originalCall = call.copy(callee = callee, receiver = captured.sourceReceiver),
         targetCaptured = true,
     )
 }
@@ -141,7 +154,9 @@ internal fun TsExprResolver.resolveCallPointer(local: EtsLocal): UExpr<out USort
     if (!value.isFakeObject()) return value
 
     val reference = value.extractRef(scope)
-    val knownFunction = reference is UConcreteHeapRef && scope.calcOnState { associatedFunction[reference] != null }
+    val knownFunction = reference is UConcreteHeapRef && scope.calcOnState {
+        associatedFunction[reference] != null || capturedInstanceCalls[reference] != null
+    }
     if (!knownFunction) return value
 
     scope.assert(value.getFakeType(scope).refTypeExpr) ?: return null

usvm-ts/src/main/kotlin/org/usvm/machine/expr/TsExprResolver.kt

@@ -489,6 +489,7 @@ class TsExprResolver(
         val isKnownFunction = scope.calcOnState {
             ref is UConcreteHeapRef && (
                 associatedFunction[ref] != null ||
+                    capturedInstanceCalls[ref] != null ||
                     memory.types.getTypeStream(ref).singleOrNull() is EtsFunctionType
                 )
         }

usvm-ts/src/test/kotlin/org/usvm/machine/expr/NativeObjectAssignIdentityTest.kt

@@ -0,0 +1,119 @@
+package org.usvm.machine.expr
+
+import org.jacodb.ets.model.EtsMethod
+import org.junit.jupiter.api.io.TempDir
+import org.usvm.PathSelectionStrategy
+import org.usvm.StateCollectionStrategy
+import org.usvm.UMachineOptions
+import org.usvm.api.TsTest
+import org.usvm.api.TsTestValue
+import org.usvm.machine.TsMachine
+import org.usvm.machine.TsOptions
+import org.usvm.machine.call.TsResidualCallPolicy
+import org.usvm.machine.call.TsUnknownCallModelSelection
+import org.usvm.util.TsMethodTestRunner
+import org.usvm.util.TsTestResolver
+import org.usvm.util.assertNodeReplay
+import org.usvm.util.getResourcePath
+import java.nio.file.Path
+import kotlin.io.path.readText
+import kotlin.test.Test
+
+class NativeObjectAssignIdentityTest : TsMethodTestRunner() {
+    @TempDir
+    lateinit var directory: Path
+
+    override val scene = loadScene("/models/NativeObjectAssignIdentity.ts")
+    override val tsOptions = TsOptions(
+        unknownCallModelSelection = TsUnknownCallModelSelection.Only(emptySet()),
+        unknownCallFallback = TsResidualCallPolicy.FRESH_SYMBOLIC_RETURN,
+    )
+    override val runner: (EtsMethod, UMachineOptions) -> List<TsTest> = { method, machineOptions ->
+        TsMachine(scene = scene, options = machineOptions, tsOptions = tsOptions).use { machine ->
+            machine.analyze(methods = listOf(method)).map { state -> TsTestResolver().resolve(method, state) }
+        }
+    }
+
+    init {
+        options = options.copy(
+            pathSelectionStrategies = listOf(PathSelectionStrategy.BFS),
+            stateCollectionStrategy = StateCollectionStrategy.ALL,
+            stopOnCoverage = 0,
+            throwExceptionOnStepFailure = true,
+        )
+    }
+
+    @Test
+    fun `extracted native assign has stable function identity`() {
+        discoverProperties<TsTestValue.TsNumber>(
+            method = method(name = "stableIdentity"),
+            { result -> result.number == 1.0 },
+            invariants = arrayOf({ result -> result.number == 1.0 }),
+        )
+    }
+
+    @Test
+    fun `local and captured unbound calls retain the native call boundary`() {
+        val methods = listOf("extractedCall" to 2.0, "capturedCall" to 3.0, "nullishTarget" to 4.0, "boxedCall" to 8.0)
+        for ((name, expected) in methods) {
+            discoverProperties<TsTestValue.TsNumber>(
+                method = method(name = name),
+                { result -> result.number == expected },
+                invariants = arrayOf({ result -> result.number == expected }),
+            )
+        }
+    }
+
+    @Test
+    fun `source object named Object executes its own assign method`() {
+        discoverProperties<TsTestValue.TsNumber>(
+            method = method(name = "localShadow"),
+            { result -> result.number == 5.0 },
+            invariants = arrayOf({ result -> result.number == 5.0 }),
+        )
+    }
+
+    @Test
+    fun `native assign replacement and deletion take precedence over its initial identity`() {
+        val methods = listOf(
+            "writeNative" to 5.0,
+            "deleteNative" to 6.0,
+            "originalAfterReplacement" to 7.0,
+            "originalAfterDeletion" to 9.0,
+        )
+        for ((name, expected) in methods) {
+            discoverProperties<TsTestValue.TsNumber>(
+                method = method(name = name),
+                { result -> result.number == expected },
+                invariants = arrayOf({ result -> result.number == expected }),
+            )
+        }
+    }
+
+    @Test
+    fun `native JavaScript confirms extracted identity calls nullish target and source shadow`() {
+        val source = getResourcePath("/models/NativeObjectAssignIdentity.ts").readText()
+        val checks = """
+            if (NativeObjectAssignIdentity.stableIdentity() !== 1) throw Error("identity");
+            if (NativeObjectAssignIdentity.extractedCall() !== 2) throw Error("local call");
+            if (NativeObjectAssignIdentity.capturedCall() !== 3) throw Error("captured call");
+            if (NativeObjectAssignIdentity.boxedCall() !== 8) throw Error("boxed call");
+            if (NativeObjectAssignIdentity.nullishTarget() !== 4) throw Error("nullish target");
+            if (NativeObjectAssignIdentity.localShadow() !== 5) throw Error("source shadow");
+            if (NativeObjectAssignIdentity.writeNative() !== 5) throw Error("replacement");
+            if (NativeObjectAssignIdentity.originalAfterReplacement() !== 7) throw Error("captured replacement");
+            if (NativeObjectAssignIdentity.deleteNative() !== 6) throw Error("deletion");
+            if (NativeObjectAssignIdentity.originalAfterDeletion() !== 9) throw Error("captured deletion");
+            Object.assign = extractedNativeAssign;
+        """.trimIndent()
+
+        assertNodeReplay(
+            source = source + checks,
+            directory = directory,
+            name = "native-object-assign-identity",
+            timeoutMessage = "Native Object.assign identity replay timed out",
+        )
+    }
+
+    private fun method(name: String) = getMethod(methodName = name, className = "NativeObjectAssignIdentity")
+}

usvm-ts/src/test/resources/models/NativeObjectAssignIdentity.ts

@@ -0,0 +1,88 @@
+const extractedNativeAssign = Object.assign;
+
+class SourceAssignReceiver {
+    assign(left: number, right: number): number {
+        return left + right;
+    }
+}
+
+export class NativeObjectAssignIdentity {
+    static stableIdentity(): number {
+        const same = Object.assign;
+        const native: any = Object;
+        return typeof extractedNativeAssign === "function" && same === extractedNativeAssign &&
+            native.assign === extractedNativeAssign ? 1 : 0;
+    }
+
+    static extractedCall(): number {
+        const assign = Object.assign;
+        assign({}, { value: 1 });
+        return 2;
+    }
+
+    static capturedCall(): number {
+        extractedNativeAssign({}, { value: 1 });
+        return 3;
+    }
+
+    static boxedCall(): number {
+        const values: any[] = [Object.assign];
+        const assign = values[0];
+        try {
+            assign(undefined, { value: 1 });
+            return 0;
+        } catch {
+            return 8;
+        }
+    }
+
+    static nullishTarget(): number {
+        try {
+            extractedNativeAssign(undefined, { value: 1 });
+            return 0;
+        } catch {
+            return 4;
+        }
+    }
+
+    static localShadow(): number {
+        const Object = new SourceAssignReceiver();
+        return Object.assign(2, 3);
+    }
+
+    static writeNative(): number {
+        const native: any = Object;
+        native.assign = (left: number, right: number) => left + right;
+        return native.assign(2, 3);
+    }
+
+    static originalAfterReplacement(): number {
+        const original = extractedNativeAssign;
+        const native: any = Object;
+        native.assign = () => 99;
+        try {
+            original(null, { value: 1 });
+            return 0;
+        } catch {
+            return 7;
+        }
+    }
+
+    static originalAfterDeletion(): number {
+        const original = extractedNativeAssign;
+        const native: any = Object;
+        delete native.assign;
+        try {
+            original(null, { value: 1 });
+            return 0;
+        } catch {
+            return 9;
+        }
+    }
+
+    static deleteNative(): number {
+        const native: any = Object;
+        delete native.assign;
+        return typeof native.assign === "undefined" ? 6 : 0;
+    }
+}