A class field's function type inferred from its initializer works

tslang/include/TypeScript/MLIRLogic/MLIRGenStore.h

@@ -773,6 +773,10 @@ struct ClassInfo
     NodeArray<ClassElement> extraMembers;
     NodeArray<ClassElement> extraMembersPost;
 
+    // fields whose initializer has no type yet in a discovery pass (`fn = this.m`: the methods are
+    // declared after the fields), typed again once the members are (mlirGenClassRetypeFields)
+    llvm::SmallVector<PropertyDeclaration> fieldsTypedAfterMembers;
+
     llvm::StringMap<std::pair<TypeParameterDOM::TypePtr, mlir::Type>> typeParamsWithArgs;
 
     bool isDeclaration;

tslang/lib/TypeScript/MLIRGenClasses.cpp

@@ -179,6 +179,7 @@ namespace mlirgen
         resetScope();   
 
         setProcessingState(newClassPtr, ProcessingStages::ProcessingStorageClass, genContext);
+        newClassPtr->fieldsTypedAfterMembers.clear();
         if (mlir::failed(mlirGenClassStorageType(location, classDeclarationAST, newClassPtr, classGenContext)))
         {
             setProcessingState(newClassPtr, ProcessingStages::ErrorInStorageClass, genContext);
@@ -270,6 +271,12 @@ namespace mlirgen
             return {mlir::failure(), ""};
         }
 
+        if (mlir::failed(mlirGenClassRetypeFields(newClassPtr, classGenContext)))
+        {
+            setProcessingState(newClassPtr, ProcessingStages::ErrorInMembers, genContext);
+            return {mlir::failure(), ""};
+        }
+
         // the slots of a @cpp class's methods, each where its @slot(N) puts it: two in one slot is
         // an error (getVirtualTable finds it, here it's reported)
         if (newClassPtr->isCppVTable)
@@ -1419,6 +1426,69 @@ namespace mlirgen
         return mlir::success();
     }    
 
+    std::pair<mlir::Type, bool> MLIRGenImpl::mlirGenClassDataFieldType(mlir::Location location, PropertyDeclaration propertyDeclaration,
+                                                                     const GenContext &genContext)
+    {
+        auto [type, init, typeProvided] = evaluateTypeAndInit(propertyDeclaration, genContext);
+        if (init)
+        {
+            type = mth.wideStorageType(type);
+        }
+
+        // a field's type inferred from `fn = () => 0` is `() => number`, which a capturing closure
+        // can be assigned to later: store it as a hybrid function, as an annotated field and a `let`
+        // (adjustLocalVariableType) are, not as a plain one with no room for the closure's context
+        if (!typeProvided)
+        {
+            if (auto funcType = dyn_cast_or_null<mlir_ts::FunctionType>(type); funcType && !funcType.isCVarArg())
+            {
+                type = mlir_ts::HybridFunctionType::get(builder.getContext(), funcType);
+            }
+        }
+
+        // `v?: T` is `v: T | undefined`, as in a type literal: unassigned it reads as undefined, not as T's zero.
+        // An interface's `v?: T` member reads it by its flag (see mlirGenClassVirtualTableDefinitionForInterface).
+        if (propertyDeclaration->questionToken && type && !mth.isNoneType(type))
+        {
+            type = getUnionType(location, type, getUndefinedType());
+        }
+
+        return {type, init};
+    }
+
+    // the fields mlirGenClassDataFieldMember left untyped in a discovery pass get their types now that
+    // the members are declared; code made before against the placeholder is thrown away with the pass
+    mlir::LogicalResult MLIRGenImpl::mlirGenClassRetypeFields(ClassInfo::TypePtr newClassPtr, const GenContext &genContext)
+    {
+        if (newClassPtr->fieldsTypedAfterMembers.empty())
+        {
+            return mlir::success();
+        }
+
+        auto storageType = mlir::cast<mlir_ts::ClassStorageType>(newClassPtr->classType.getStorageType());
+        SmallVector<mlir_ts::FieldInfo> fieldInfos(storageType.getFields().begin(), storageType.getFields().end());
+        for (auto propertyDeclaration : newClassPtr->fieldsTypedAfterMembers)
+        {
+            auto [type, init] = mlirGenClassDataFieldType(loc(propertyDeclaration), propertyDeclaration, genContext);
+            if (mth.isNoneType(type))
+            {
+                continue;
+            }
+
+            auto fieldId = TupleFieldName(propertyDeclaration->name, genContext);
+            for (auto &fieldInfo : fieldInfos)
+            {
+                if (fieldInfo.id == fieldId)
+                {
+                    fieldInfo.type = type;
+                }
+            }
+        }
+
+        newClassPtr->fieldsTypedAfterMembers.clear();
+        return storageType.setFields(fieldInfos);
+    }
+
     mlir::LogicalResult MLIRGenImpl::mlirGenClassDataFieldMember(mlir::Location location, ClassInfo::TypePtr newClassPtr, SmallVector<mlir_ts::FieldInfo> &fieldInfos, 
                                                     PropertyDeclaration propertyDeclaration, const GenContext &genContext)
     {
@@ -1450,22 +1520,14 @@ namespace mlirgen
             }
         }
 
-        auto [type, init, typeProvided] = evaluateTypeAndInit(propertyDeclaration, genContext);
+        auto [type, init] = mlirGenClassDataFieldType(location, propertyDeclaration, genContext);
         if (init)
         {
             newClassPtr->hasInitializers = true;
-            type = mth.wideStorageType(type);
         }
 
         auto hasType = !!propertyDeclaration->type;
 
-        // `v?: T` is `v: T | undefined`, as in a type literal: unassigned it reads as undefined, not as T's zero.
-        // An interface's `v?: T` member reads it by its flag (see mlirGenClassVirtualTableDefinitionForInterface).
-        if (propertyDeclaration->questionToken && type && !mth.isNoneType(type))
-        {
-            type = getUnionType(location, type, getUndefinedType());
-        }
-
         LLVM_DEBUG(dbgs() << "\n!! class field: " << fieldId << " type: " << type << " access level: " << accessLevel);
         if (mth.isNoneType(type))
         {
@@ -1474,6 +1536,15 @@ namespace mlirgen
                 return mlir::failure();
             }
 
+            // `fn = this.m`: a discovery pass lays out the fields before it declares the methods, so
+            // the initializer has no type yet. Hold the place and type the field once they are.
+            if (init && !isAccessor && genContext.allowPartialResolve)
+            {
+                newClassPtr->fieldsTypedAfterMembers.push_back(propertyDeclaration);
+                fieldInfos.push_back({fieldId, getAnyType(), false, accessLevel});
+                return mlir::success();
+            }
+
 #ifndef ANY_AS_DEFAULT
             emitError(location)
                 << "type for field '" << fieldId << "' is not provided, field must have type or initializer";

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -11856,7 +11856,12 @@ class MLIRGenImpl
                     : mlir_ts::AccessLevel::Public;
     }
 
-    mlir::LogicalResult mlirGenClassDataFieldMember(mlir::Location location, ClassInfo::TypePtr newClassPtr, SmallVector<mlir_ts::FieldInfo> &fieldInfos, 
+    std::pair<mlir::Type, bool> mlirGenClassDataFieldType(mlir::Location location, PropertyDeclaration propertyDeclaration,
+                                                          const GenContext &genContext);
+
+    mlir::LogicalResult mlirGenClassRetypeFields(ClassInfo::TypePtr newClassPtr, const GenContext &genContext);
+
+    mlir::LogicalResult mlirGenClassDataFieldMember(mlir::Location location, ClassInfo::TypePtr newClassPtr, SmallVector<mlir_ts::FieldInfo> &fieldInfos,
                                                     PropertyDeclaration propertyDeclaration, const GenContext &genContext);
 
     mlir::LogicalResult mlirGenClassStaticFieldMember(mlir::Location location, ClassInfo::TypePtr newClassPtr, PropertyDeclaration propertyDeclaration, const GenContext &genContext);

tslang/test/tester/CMakeLists.txt

@@ -613,6 +613,7 @@ tslang_add_test(NAME test-compile-00-interface-optional-cast-order COMMAND test-
 tslang_add_test(NAME test-compile-00-interface-function-typed-field COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_function_typed_field.ts")
 tslang_add_test(NAME test-compile-00-interface-function-typed-field-class COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_function_typed_field_class.ts")
 tslang_add_test(NAME test-compile-00-interface-function-typed-field-class-field COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_function_typed_field_class_field.ts")
+tslang_add_test(NAME test-compile-00-class-function-typed-field-inferred COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_function_typed_field_inferred.ts")
 tslang_add_test(NAME test-compile-00-interface-optional-function-typed-field COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_optional_function_typed_field.ts")
 tslang_add_test(NAME test-compile-00-interface-captures COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_captures.ts")
 tslang_add_test(NAME test-compile-00-object-annotated-method COMMAND test-runner "${PROJECT_SOURCE_DIR}/test/tester/tests/00object_annotated_method.ts")
@@ -1169,6 +1170,7 @@ tslang_add_test(NAME test-jit-00-interface-optional-cast-order COMMAND test-runn
 tslang_add_test(NAME test-jit-00-interface-function-typed-field COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_function_typed_field.ts")
 tslang_add_test(NAME test-jit-00-interface-function-typed-field-class COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_function_typed_field_class.ts")
 tslang_add_test(NAME test-jit-00-interface-function-typed-field-class-field COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_function_typed_field_class_field.ts")
+tslang_add_test(NAME test-jit-00-class-function-typed-field-inferred COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00class_function_typed_field_inferred.ts")
 tslang_add_test(NAME test-jit-00-interface-optional-function-typed-field COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_optional_function_typed_field.ts")
 tslang_add_test(NAME test-jit-00-interface-captures COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00interface_captures.ts")
 tslang_add_test(NAME test-jit-00-object-annotated-method COMMAND test-runner -jit "${PROJECT_SOURCE_DIR}/test/tester/tests/00object_annotated_method.ts")
@@ -2033,6 +2035,7 @@ set(TSLANG_CORPUS
     00interface_function_typed_field.ts
     00interface_function_typed_field_class.ts
     00interface_function_typed_field_class_field.ts
+    00class_function_typed_field_inferred.ts
     00interface_generic.ts
     00interface_global_method.ts
     00interface_indexer.ts

tslang/test/tester/tests/00class_function_typed_field_inferred.ts

@@ -0,0 +1,83 @@
+// regression test: a class field whose function type is inferred from its
+// initializer.
+// - `fn = () => 0` was stored as a plain function, which has no room for a
+//   closure's context: assigning a capturing closure to it warned "losing this
+//   reference" and the call crashed. It is stored as a hybrid function now, as
+//   an annotated `fn: () => int` field and a `let` are.
+// - `fn = this.m` had no type in the first pass over the program (the class's
+//   methods aren't known while its fields are), so the field was `any` there
+//   and `w.fn()` failed with "not a function to call".
+
+interface HasFn {
+    fn: () => int;
+}
+
+class Capturing {
+    fn = () => 0;
+    constructor(k: int) {
+        this.fn = () => k;
+    }
+}
+
+class WithParams {
+    add = (a: int, b: int) => a + b;
+}
+
+class BoundMethod {
+    v = 4;
+    m(): int {
+        return this.v;
+    }
+    fn = this.m;
+}
+
+class BoundBeforeMethod {
+    fn = this.twice;
+    constructor(public v: int) {}
+    twice(): int {
+        return this.v * 2;
+    }
+}
+
+function main() {
+    const c = new Capturing(9);
+    print(c.fn());
+    assert(c.fn() == 9);
+
+    let k2 = 5;
+    c.fn = () => k2 + 1;
+    assert(c.fn() == 6);
+    k2 = 10;
+    assert(c.fn() == 11);
+
+    const p = new WithParams();
+    assert(p.add(2, 3) == 5);
+    const base = 100;
+    p.add = (a: int, b: int) => a * b + base;
+    assert(p.add(2, 3) == 106);
+
+    // under -mm=rc a bound method does not hold its receiver yet: `b` keeps it alive
+    const b = new BoundMethod();
+    print(b.fn());
+    assert(b.fn() == 4);
+    b.v = 7;
+    assert(b.fn() == 7);
+    const g = b.fn;
+    assert(g() == 7);
+
+    const bb = new BoundBeforeMethod(21);
+    assert(bb.fn() == 42);
+
+    // an inferred function-typed field fills an interface's function-typed field
+    const h: HasFn = c;
+    assert(h.fn() == 11);
+    const hb: HasFn = b;
+    assert(hb.fn() == 7);
+
+    // a closure replaces the bound method
+    const k3 = 3;
+    b.fn = () => k3;
+    assert(b.fn() == 3);
+
+    print("done.");
+}