[TS] Capture lexical receivers for stored arrows (#391)

jacodb-ets/src/main/kotlin/org/jacodb/ets/dto/Convert.kt

@@ -713,6 +713,7 @@ fun TypeDto.toEtsType(): EtsType = when (this) {
     is FunctionTypeDto -> EtsFunctionType(
         signature = signature.toEtsMethodSignature(),
         typeParameters = typeParameters.map { it.toEtsType() },
+        isArrow = isArrow,
     )
 
     is GenericTypeDto -> EtsGenericType(

jacodb-ets/src/main/kotlin/org/jacodb/ets/dto/Types.kt

@@ -205,4 +205,5 @@ data class TupleTypeDto(
 data class FunctionTypeDto(
     val signature: MethodSignatureDto,
     val typeParameters: List<TypeDto> = emptyList(),
+    val isArrow: Boolean = false,
 ) : TypeDto

jacodb-ets/src/main/kotlin/org/jacodb/ets/model/Type.kt

@@ -443,6 +443,7 @@ data class EtsTupleType(
 data class EtsFunctionType(
     val signature: EtsMethodSignature,
     val typeParameters: List<EtsType> = emptyList(),
+    val isArrow: Boolean = false,
 ) : EtsRefType {
     override val typeName: String
         get() = if (typeParameters.isNotEmpty()) {

jacodb-ets/src/test/kotlin/org/jacodb/ets/test/EtsTsLexicalThisTest.kt

@@ -0,0 +1,97 @@
+/*
+ *  Copyright 2022 UnitTestBot contributors (utbot.org)
+ * <p>
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ * <p>
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package org.jacodb.ets.test
+
+import kotlin.io.path.createTempDirectory
+import kotlin.io.path.exists
+import kotlin.io.path.readText
+import kotlin.io.path.writeText
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+import kotlinx.serialization.encodeToString
+import kotlinx.serialization.json.Json
+import org.jacodb.ets.dto.EtsFileDto
+import org.jacodb.ets.dto.dtoModule
+import org.jacodb.ets.dto.toEtsFile
+import org.jacodb.ets.utils.DEFAULT_ARK_CLASS_NAME
+import org.jacodb.ets.utils.EtsIrProvider
+import org.jacodb.ets.utils.generateEtsIR
+import org.junit.jupiter.api.Test
+import org.jacodb.ets.model.EtsScene
+import org.jacodb.ets.model.EtsAssignStmt
+import org.jacodb.ets.model.EtsLocal
+import org.jacodb.ets.model.EtsFunctionType
+import org.jacodb.ets.model.EtsClosureFieldRef
+import org.jacodb.ets.model.EtsThis
+
+class EtsTsLexicalThisTest {
+    companion object {
+        /** Run the PRODUCTION integration path: generateEtsIR + EtsFileDto.loadFromJson. */
+        private fun runFrontend(source: String, fileName: String = "test.ts"): EtsFileDto {
+            val dir = createTempDirectory("ts-frontend-test")
+            val inputPath = dir.resolve(fileName)
+            inputPath.writeText(source)
+
+            val outputPath = generateEtsIR(
+                inputPath,
+                isProject = false,
+                timeout = null,
+                provider = EtsIrProvider.TS_FRONTEND,
+            )
+            assertTrue(outputPath.exists(), "ts-frontend did not produce output: $outputPath")
+
+            return EtsFileDto.loadFromJson(outputPath.readText())
+        }
+    }
+
+    @Test
+    fun `stored arrow lexical receiver survives JSON and model conversion`() {
+        val frontendDto = runFrontend(
+            source = """
+                class Box {
+                    offset = 11;
+                    installOn(target: Box) {
+                        target.callback = (value: number) => this.offset + value;
+                    }
+                    installOrdinaryOn(target: Box) {
+                        target.callback = function(value: number) { return this.offset + value; };
+                    }
+                }
+            """.trimIndent(),
+        )
+        val roundTripped = EtsFileDto.loadFromJson(
+            Json { serializersModule = dtoModule }.encodeToString(frontendDto),
+        )
+        val box = roundTripped.toEtsFile().classes.single { it.name == "Box" }
+        val install = box.methods.single { it.name == "installOn" }
+        val arrowValue = install.cfg.stmts.filterIsInstance<EtsAssignStmt>()
+            .map { it.rhv }
+            .filterIsInstance<EtsLocal>()
+            .single { it.name.startsWith("%AM") }
+        val arrow = box.methods.single { it.name == arrowValue.name }
+        val ordinary = box.methods.single { it.name.endsWith("\$installOrdinaryOn") }
+
+        assertTrue((arrowValue.type as EtsFunctionType).isArrow)
+        val receiver = arrow.cfg.stmts.filterIsInstance<EtsAssignStmt>()
+            .single { (it.lhv as? EtsLocal)?.name == "this" }.rhv
+        assertTrue(receiver is EtsClosureFieldRef)
+        assertEquals("this", receiver.fieldName)
+        assertTrue(ordinary.cfg.stmts.filterIsInstance<EtsAssignStmt>()
+            .single { (it.lhv as? EtsLocal)?.name == "this" }.rhv is org.jacodb.ets.model.EtsThis)
+    }
+
+}

jacodb-ets/ts-frontend/src/dto/types.ts

@@ -179,6 +179,8 @@ export interface FunctionTypeDto {
     readonly _: "FunctionType";
     signature: MethodSignatureDto;
     typeParameters?: TypeDto[]; // Kotlin default: empty list
+    /** Arrow values use the receiver saved in their lexical environment. */
+    isArrow?: boolean; // Kotlin default: false
 }
 
 // Singletons for field-less types.

jacodb-ets/ts-frontend/src/lowering/exprLowering.ts

@@ -1574,18 +1574,28 @@ export class ExprLowerer {
             .filter((identifier) => this.m.moduleFieldForIdentifier(identifier) === undefined)
             .map((identifier) => this.m.captureForIdentifier(identifier, this.safeTypeOf(identifier)));
 
+        const isArrow = ts.isArrowFunction(node);
+        const lexicalThis = isArrow && usesLexicalThis(node)
+            ? this.m.getOrCreateLocal("this", this.m.thisType())
+            : undefined;
+
         let signature = baseSignature;
         let environment:
             | { name: string; type: Extract<TypeDto, { _: "LexicalEnvType" }> }
             | undefined;
-        if (captures.length > 0) {
+        if (captures.length > 0 || lexicalThis !== undefined) {
+            const capturedLocals = captures.map((capture) => ({
+                name: capture.outerLocal.name,
+                type: capture.outerLocal.type,
+            }));
+            if (lexicalThis !== undefined) {
+                capturedLocals.push({ name: lexicalThis.name, type: lexicalThis.type });
+            }
+
             const environmentType: Extract<TypeDto, { _: "LexicalEnvType" }> = {
                 _: "LexicalEnvType",
                 method: baseSignature,
-                closures: [...new Map(captures.map((capture) => [
-                    capture.outerLocal.name,
-                    { name: capture.outerLocal.name, type: capture.outerLocal.type },
-                ])).values()],
+                closures: [...new Map(capturedLocals.map((local) => [local.name, local])).values()],
             };
             const environmentLocal = this.m.newClosureEnvironment(environmentType);
             environment = { name: environmentLocal.name, type: environmentType };
@@ -1608,7 +1618,7 @@ export class ExprLowerer {
         if (environment === undefined) {
             closureContext.emitPrologue(prologueParams);
         } else {
-            closureContext.emitClosurePrologue(environment.name, environment.type, captures, prologueParams);
+            closureContext.emitClosurePrologue(environment.name, environment.type, captures, prologueParams, lexicalThis);
         }
         this.lowerFunctionBody(closureContext, node.body, prologueParams);
         registry.methods.push({
@@ -1618,7 +1628,7 @@ export class ExprLowerer {
             body: closureContext.build(),
         });
 
-        return this.m.getOrCreateLocal(name, { _: "FunctionType", signature });
+        return this.m.getOrCreateLocal(name, { _: "FunctionType", signature, ...(isArrow ? { isArrow } : {}) });
     }
 
     /**
@@ -1957,6 +1967,24 @@ function numericConstantText(node: ts.NumericLiteral): string {
     return Number.isFinite(Number(node.text)) ? node.text : node.getText();
 }
 
+/** Nested arrows share the receiver; ordinary functions and class members introduce their own. */
+function usesLexicalThis(arrow: ts.ArrowFunction): boolean {
+    let found = false;
+    const visit = (node: ts.Node): void => {
+        if (found) return;
+        if (node.kind === ts.SyntaxKind.ThisKeyword) {
+            found = true;
+            return;
+        }
+        if (!ts.isArrowFunction(node) && (ts.isFunctionLike(node) || ts.isClassLike(node))) return;
+        ts.forEachChild(node, visit);
+    };
+
+    for (const parameter of arrow.parameters) visit(parameter);
+    visit(arrow.body);
+    return found;
+}
+
 function bigIntConstantText(node: ts.BigIntLiteral): string {
     return BigInt(node.text.replace(/_/g, "").replace(/n$/, "")).toString();
 }
@@ -2136,7 +2164,7 @@ function collectCapturedIdentifiers(
         if (isCapturableDeclaration(node as ts.Declaration)) {
             ownDeclarations.add(node as ts.Declaration);
         }
-        if (ts.isIdentifier(node)) {
+        if (ts.isIdentifier(node) && !(ts.isPropertyAccessExpression(node.parent) && node.parent.name === node)) {
             let symbol: ts.Symbol | undefined;
             try {
                 symbol = checker.getSymbolAtLocation(node);

jacodb-ets/ts-frontend/src/lowering/methodBuilder.ts

@@ -313,6 +313,7 @@ export class MethodContext {
         environmentType: LexicalEnvTypeDto,
         captures: ClosureCapture[],
         parameters: { name: string; type: TypeDto; identifier?: ts.Identifier }[],
+        lexicalThis?: LocalDto,
     ): void {
         const environment = this.getOrCreateLocal(environmentName, environmentType);
         this.cfg.emit({
@@ -366,16 +367,22 @@ export class MethodContext {
                 });
             }
         }
-        this.emitThisAssignment();
+        const receiver: ClosureFieldRefDto | undefined = lexicalThis === undefined ? undefined : {
+            _: "ClosureFieldRef",
+            base: { name: environment.name, type: environment.type },
+            fieldName: lexicalThis.name,
+            type: lexicalThis.type,
+        };
+        this.emitThisAssignment(receiver);
     }
 
-    private emitThisAssignment(): void {
+    private emitThisAssignment(receiver?: ClosureFieldRefDto): void {
         const thisType = this.thisType();
         const thisLocal = this.getOrCreateLocal("this", thisType);
         this.cfg.emit({
             _: "AssignStmt",
             left: thisLocal,
-            right: { _: "ThisRef", type: thisType },
+            right: receiver ?? { _: "ThisRef", type: thisType },
         });
     }
 

jacodb-ets/ts-frontend/test/lexical-this.spec.ts

@@ -0,0 +1,95 @@
+import { describe, expect, it } from "vitest";
+import * as ts from "typescript";
+import { serializeEtsFile } from "../src/serialize";
+import { executeObjectIr } from "./object-runtime";
+import { lower, methodByName } from "./util";
+
+describe("lexical arrow receivers", () => {
+    it("uses the creation receiver when a stored arrow body reads this", () => {
+        const { file, diagnostics } = lower(`
+            export function result() {
+                const source = {
+                    offset: 11,
+                    install: function(target: any) { target.callback = () => this.offset + 1; },
+                };
+                const target: any = { offset: 100 };
+                source.install(target);
+                return target.callback();
+            }
+        `);
+        const roundTripped = JSON.parse(serializeEtsFile(file));
+
+        expect(diagnostics.messages).toEqual([]);
+        expect(executeObjectIr(roundTripped, "result")).toBe(12);
+    });
+
+    it("preserves the creation receiver when an arrow is stored on another object", () => {
+        const source = `
+            class Box {
+                callback!: (value: number) => number;
+                constructor(public offset: number) {}
+                installOn(target: Box): void {
+                    target.callback = (value: number) => this.offset + value;
+                }
+                installOrdinaryOn(target: Box): void {
+                    target.callback = function(value: number) { return this.offset + value; };
+                }
+                call(value: number): number { return this.callback(value); }
+            }
+            export function results(): number[] {
+                const source = new Box(11);
+                const target = new Box(100);
+                source.installOn(target);
+                const arrowResult = target.call(1);
+                source.installOrdinaryOn(target);
+                return [arrowResult, target.call(1)];
+            }
+        `;
+        const { file, diagnostics } = lower(source);
+        const roundTripped = JSON.parse(serializeEtsFile(file));
+        const box = roundTripped.classes.find((classDto) => classDto.signature.name === "Box");
+        const install = box.methods.find((method) => method.signature.name === "installOn");
+        const arrowValue = install.body.locals.find((local) => local.name.startsWith("%AM"));
+        const arrow = methodByName(roundTripped, arrowValue.name);
+        const ordinary = box.methods.find((method) => method.signature.name.includes("$installOrdinaryOn"));
+
+        expect(diagnostics.messages).toEqual([]);
+        expect(arrowValue.type.isArrow).toBe(true);
+        expect(arrow.signature.parameters[0].type).toMatchObject({
+            _: "LexicalEnvType", closures: [{ name: "this", type: { _: "ClassType" } }],
+        });
+        expect(arrow.body.cfg.blocks.flatMap((block) => block.stmts)).toContainEqual(expect.objectContaining({
+            _: "AssignStmt", left: expect.objectContaining({ name: "this" }),
+            right: expect.objectContaining({ _: "ClosureFieldRef", fieldName: "this" }),
+        }));
+        expect(ordinary.body.cfg.blocks.flatMap((block) => block.stmts)).toContainEqual(expect.objectContaining({
+            _: "AssignStmt", left: expect.objectContaining({ name: "this" }),
+            right: expect.objectContaining({ _: "ThisRef" }),
+        }));
+
+        const javascript = ts.transpileModule(source, {
+            compilerOptions: { target: ts.ScriptTarget.ES2020, module: ts.ModuleKind.CommonJS },
+        }).outputText;
+        const concrete = new Function("exports", `${javascript}\nreturn exports.results();`)({});
+        expect(concrete).toEqual([12, 101]);
+    });
+
+    it("forwards lexical this through nested arrows and stops at ordinary functions", () => {
+        const { file, diagnostics } = lower(`
+            class Box {
+                offset = 11;
+                nested() { return () => () => this.offset; }
+                ordinary() { return () => function() { return this.offset; }; }
+            }
+        `);
+        const methods = file.classes.flatMap((classDto) => classDto.methods)
+            .filter((method) => method.signature.name.startsWith("%AM"));
+        const nested = methods.filter((method) => method.signature.name.includes("$nested"));
+        const ordinary = methods.filter((method) => method.signature.name.includes("$ordinary"));
+
+        expect(diagnostics.messages).toEqual([]);
+        expect(nested).toHaveLength(2);
+        expect(nested.every((method) => method.signature.parameters[0].type._ === "LexicalEnvType")).toBe(true);
+        expect(ordinary.every((method) => method.signature.parameters.length === 0)).toBe(true);
+    });
+});