An object's closure field keeps its captures through an interface

tslang/include/TypeScript/Defines.h

@@ -170,6 +170,8 @@
 // (`fn?: () => int`, called through a thunk) has its flag at `this` + the offset, so a read of the
 // interface's optional field (`fn?: () => int`) is absent while the instance's flag is clear; -1
 // is always present (a method, or a field that is not optional).
+// An object field holding a function the interface can't call with the object as its `this` (a
+// closure) is filled the same way, with a thunk calling the field; its cell is always present.
 #define INTERFACE_METHOD_CELL_TAG 1
 // on an InterfaceSymbolRefOp of a `v?: T` field that is stored to: true sets the flag of a tagged slot's
 // field, false (a store of undefined) clears it

tslang/include/TypeScript/MLIRLogic/MLIRTypeHelper.h

@@ -1141,6 +1141,26 @@ class MLIRTypeHelper
         return isa<mlir_ts::BoundFunctionType>(inFuncType) || isa<mlir_ts::ExtensionFunctionType>(inFuncType);
     }
 
+    // An object's field holding a function an interface slot can't call with the object as its `this`:
+    // a bound function (a capturing closure, its `this` is the captures), a hybrid one (a type-literal
+    // annotated field) or a plain function without a `this` parameter (`{ add: plusOne }`). Only an
+    // object-literal method (`fn() {...}`, `fn: () => ...` written in the literal) takes the object.
+    static bool isObjectFieldFunctionWithoutThis(mlir::Type fieldType)
+    {
+        if (isa<mlir_ts::BoundFunctionType>(fieldType) || isa<mlir_ts::HybridFunctionType>(fieldType))
+        {
+            return true;
+        }
+
+        if (auto funcType = dyn_cast<mlir_ts::FunctionType>(fieldType))
+        {
+            return funcType.getNumInputs() == 0
+                || !isa<mlir_ts::OpaqueType>(funcType.getInput(0)) && !isa<mlir_ts::ObjectType>(funcType.getInput(0));
+        }
+
+        return false;
+    }
+
     MatchResult TestFunctionTypesMatchWithObjectMethods(mlir::Location location, mlir::Type inFuncType, mlir::Type resFuncType, unsigned startParamIn = 0,
                                                         unsigned startParamRes = 0)
     {
@@ -1424,6 +1444,19 @@ class MLIRTypeHelper
                 if (foundIndex >= 0)
                 {
                     auto foundField = tupleStorageType.getFieldInfo(foundIndex);
+
+                    // a function the slot can't hand out as one taking the object as its `this` is called
+                    // through a thunk reading the field (isObjectFieldFunctionWithoutThis): the slot has the
+                    // interface's type
+                    if (isa<mlir_ts::FunctionType>(fieldType) && fieldType != foundField.type
+                        && isObjectFieldFunctionWithoutThis(foundField.type))
+                    {
+                        if (TestFunctionTypesMatchWithObjectMethods(location, foundField.type, fieldType).result == MatchResultType::Match)
+                        {
+                            return {mlir_ts::FieldInfo{id, fieldType, false, mlir_ts::AccessLevel::Public}, mlir::success()};
+                        }
+                    }
+
                     auto test = isa<mlir_ts::FunctionType>(foundField.type) && isa<mlir_ts::FunctionType>(fieldType)
                                     ? TestFunctionTypesMatchWithObjectMethods(location, foundField.type, fieldType).result == MatchResultType::Match
                                     : stripLiteralType(fieldType) == stripLiteralType(foundField.type);

tslang/lib/TypeScript/LowerToLLVM.cpp

@@ -6476,7 +6476,8 @@ struct InterfaceSymbolRefOpLowering : public TsLlvmPattern<mlir_ts::InterfaceSym
             };
 
             // a function-typed field's slot tagged INTERFACE_METHOD_CELL_TAG: a class fills the field with a
-            // method, and the slot is the address of the cell holding it. Tested ahead of
+            // method, or a class or object with a thunk calling its field, and the slot is the address of the
+            // cell holding it. Tested ahead of
             // INTERFACE_OPTIONAL_STORAGE_TAG, which a cell's address can carry on a 32-bit target.
             // With checkPresence (a read of an optional field), the cell's flag offset decides: -1 is
             // present, any other is the offset from `this` of the instance's flag, and a clear flag reads

tslang/lib/TypeScript/MLIRGenCast.cpp

@@ -21,7 +21,10 @@ namespace mlirgen
     // member order (methods-first, per InterfaceInfo::getTupleTypeFields)
     // silently reads/writes the wrong slots at runtime. Only the TYPES are taken
     // from the interface; interface-only members are appended after the source
-    // fields.
+    // fields. A function the interface can't call with the object as its `this`
+    // (a closure, see MLIRTypeHelper::isObjectFieldFunctionWithoutThis) keeps its
+    // own type: cast to the interface's it would lose its captures. Its slot calls a
+    // thunk reading the field instead (mlirGenObjectVirtualTableDefinitionForInterface).
     static mlir::LogicalResult getInterfaceCloneFields(mlir::ArrayRef<mlir_ts::FieldInfo> srcFields,
                                                        InterfaceInfo::TypePtr interfaceInfo, mlir::MLIRContext *context,
                                                        SmallVector<mlir_ts::FieldInfo> &fields)
@@ -37,7 +40,10 @@ namespace mlirgen
             auto interfaceField =
                 std::find_if(interfaceFields.begin(), interfaceFields.end(),
                              [&](auto &item) { return item.id == origField.id; });
-            fields.push_back(interfaceField != interfaceFields.end() ? *interfaceField : origField);
+            auto keepsOwnType = interfaceField != interfaceFields.end()
+                && isa<mlir_ts::FunctionType>(interfaceField->type)
+                && MLIRTypeHelper::isObjectFieldFunctionWithoutThis(origField.type);
+            fields.push_back(interfaceField != interfaceFields.end() && !keepsOwnType ? *interfaceField : origField);
         }
 
         for (auto &interfaceField : interfaceFields)

tslang/lib/TypeScript/MLIRGenClasses.cpp

@@ -2490,7 +2490,7 @@ genContext);
                         auto funcType = mlir::cast<mlir_ts::FunctionType>(fieldType);
                         auto thunkName = concat(interfaceVTableNameForClass(newClassPtr, newInterfacePtr),
                                                 fieldName.getValue(), "thunk");
-                        if (mlir::failed(generateSynthMethodToCallField(location, newClassPtr, fieldName.getValue(),
+                        if (mlir::failed(generateSynthMethodToCallField(location, newClassPtr->classType, fieldName.getValue(),
                                                                          thunkName, funcType, genContext)))
                         {
                             return {emptyFieldInfo, mlir::failure()};

tslang/lib/TypeScript/MLIRGenImpl.h

@@ -10670,6 +10670,19 @@ class MLIRGenImpl
 
         type = mth.wideStorageType(type);
 
+        // the receiver's slot takes the object as the function's `this` (an interface's function-typed member),
+        // the value does not - a closure, whose `this` is its captures (MLIRTypeHelper::isObjectFieldFunctionWithoutThis):
+        // cast, it would lose them. It keeps its own type, and the object's cast to the interface calls it through a
+        // thunk (mlirGenObjectVirtualTableDefinitionForInterface)
+        if (auto receiverFuncType = dyn_cast_or_null<mlir_ts::FunctionType>(receiverElementType))
+        {
+            if (!mth.isObjectFieldFunctionWithoutThis(receiverFuncType) && mth.isObjectFieldFunctionWithoutThis(type)
+                && mth.TestFunctionTypesMatchWithObjectMethods(location, type, receiverFuncType).result == MatchResultType::Match)
+            {
+                receiverElementType = mlir::Type();
+            }
+        }
+
         if (receiverElementType)
         {
             if (type != receiverElementType)
@@ -12162,9 +12175,10 @@ class MLIRGenImpl
     }
 
     // `<thunkName>(this, p0..pn) { return this.<fieldName>(p0..pn); }`, typed as the interface's
-    // function-typed field (`this` opaque): the method an interface slot calls for a class whose
-    // field of that name holds the function (see mlirGenClassVirtualTableDefinitionForInterface)
-    mlir::LogicalResult generateSynthMethodToCallField(mlir::Location location, ClassInfo::TypePtr newClassPtr,
+    // function-typed field (`this` opaque): the method an interface slot calls for a class or an
+    // object (`thisType`) whose field of that name holds the function (see
+    // mlirGenClassVirtualTableDefinitionForInterface, mlirGenObjectVirtualTableDefinitionForInterface)
+    mlir::LogicalResult generateSynthMethodToCallField(mlir::Location location, mlir::Type thisType,
                                                        StringRef fieldName, StringRef thunkName,
                                                        mlir_ts::FunctionType funcType, const GenContext &genContext)
     {
@@ -12175,8 +12189,8 @@ class MLIRGenImpl
 
         GenContext funcGenContext(genContext);
         funcGenContext.clearScopeVars();
-        funcGenContext.thisType = newClassPtr->classType;
-        funcGenContext.thisClassType = newClassPtr->classType;
+        funcGenContext.thisType = thisType;
+        funcGenContext.thisClassType = dyn_cast<mlir_ts::ClassType>(thisType);
         funcGenContext.disableSpreadParams = true;
 
         return mlirGenFunctionBody(
@@ -12185,8 +12199,8 @@ class MLIRGenImpl
                 // the interface passes the object as an opaque `this`
                 auto funcOp = genContext.funcOp;
                 auto opaqueThis = funcOp.getBody().front().getArgument(0);
-                auto thisValue = builder.create<mlir_ts::CastOp>(location, newClassPtr->classType, opaqueThis);
-                auto thisDecl = std::make_shared<VariableDeclarationDOM>(THIS_NAME, newClassPtr->classType, location);
+                auto thisValue = builder.create<mlir_ts::CastOp>(location, thisType, opaqueThis);
+                auto thisDecl = std::make_shared<VariableDeclarationDOM>(THIS_NAME, thisType, location);
                 DECLARE(thisDecl, thisValue);
 
                 NodeFactory nf(NodeFactoryFlags::None);

tslang/lib/TypeScript/MLIRGenInterfaces.cpp

@@ -254,8 +254,10 @@ namespace mlirgen
                         return mlir::failure();
                     }
 
+                    // a field the interface calls through a thunk has the thunk in its slot already
                     auto fieldInfo = mth.getFieldInfoByIndex(storeType, index);
-                    if (lookupObjectLiteralMethodSymbol(fieldInfo.type, fieldId).empty())
+                    if (lookupObjectLiteralMethodSymbol(fieldInfo.type, fieldId).empty()
+                        && !mth.isObjectFieldFunctionWithoutThis(fieldInfo.type))
                     {
                         methodsNeedingPatch.push_back({methodPtr, vtableOffset});
                     }
@@ -500,8 +502,97 @@ namespace mlirgen
 
         virtualTable.insert(virtualTable.begin(), VirtualMethodOrFieldInfo(noInstanceOf));
 
-        // register global
         auto fullClassInterfaceVTableFieldName = interfaceVTableNameForObject(objectType, newInterfacePtr);
+
+        // a field holding a function the interface can't call with the object as its `this` - a closure, whose
+        // `this` is its captures (MLIRTypeHelper::isObjectFieldFunctionWithoutThis) - is called through a thunk
+        // reading the field, `<vtable>.<field>.thunk(this, ...) { return this.<field>(...) }`. A method's slot holds
+        // the thunk; a function-typed field's slot the address of a cell holding it, tagged with
+        // INTERFACE_METHOD_CELL_TAG, as for a class (mlirGenClassVirtualTableDefinitionForInterface)
+        auto methodCellName = [&](mlir::Attribute id) {
+            return concat(fullClassInterfaceVTableFieldName, mlir::cast<mlir::StringAttr>(id).getValue(), "method");
+        };
+
+        llvm::SmallVector<std::pair<mlir::Attribute, StringRef>> thunksForFields;
+        for (auto &methodOrField : virtualTable)
+        {
+            if (!methodOrField.isField || methodOrField.isMissing)
+            {
+                continue;
+            }
+
+            auto id = methodOrField.fieldInfo.id;
+            auto funcType = dyn_cast<mlir_ts::FunctionType>(methodOrField.fieldInfo.type);
+            auto fieldName = dyn_cast<mlir::StringAttr>(id);
+            auto storageIndex = tupleStorageType.getIndex(id);
+            if (!funcType || !fieldName || storageIndex < 0)
+            {
+                continue;
+            }
+
+            auto storageFieldType = tupleStorageType.getFieldInfo(storageIndex).type;
+            if (storageFieldType == funcType || !mth.isObjectFieldFunctionWithoutThis(storageFieldType))
+            {
+                continue;
+            }
+
+            // a field reached twice through a diamond `extends` has two slots and one thunk
+            if (llvm::any_of(thunksForFields, [&](auto &thunkForField) { return thunkForField.first == id; }))
+            {
+                continue;
+            }
+
+            auto thunkName = concat(fullClassInterfaceVTableFieldName, fieldName.getValue(), "thunk");
+            if (mlir::failed(generateSynthMethodToCallField(location, objectType, fieldName.getValue(), thunkName,
+                                                             funcType, genContext)))
+            {
+                return mlir::failure();
+            }
+
+            thunksForFields.push_back({id, thunkName});
+
+            if (newInterfacePtr->findMethod(fieldName.getValue()))
+            {
+                continue;
+            }
+
+            auto cellType = registerVariable(
+                location, methodCellName(id), true, VariableType::Var,
+                [&](mlir::Location location, const GenContext &genContext) {
+                    auto thunkValue = builder.create<mlir_ts::SymbolRefOp>(
+                        location, funcType, mlir::FlatSymbolRefAttr::get(builder.getContext(), thunkName));
+
+                    // the cell is {method, flag offset}, as a class's: the field holds a function, so the flag
+                    // offset is -1, always present
+                    auto ptrIntType = builder.getIntegerType(compileOptions.sizeBits());
+                    auto flagOffsetValue =
+                        builder.create<mlir_ts::ConstantOp>(location, ptrIntType, builder.getIntegerAttr(ptrIntType, -1));
+                    mlir::SmallVector<mlir_ts::FieldInfo> cellFields{
+                        {MLIRHelper::TupleFieldName("method", builder.getContext()), funcType, false,
+                         mlir_ts::AccessLevel::Public},
+                        {MLIRHelper::TupleFieldName("flagOffset", builder.getContext()), ptrIntType, false,
+                         mlir_ts::AccessLevel::Public}};
+                    auto cellTupleType = getTupleType(cellFields);
+                    mlir::Value cellValue = builder.create<mlir_ts::UndefOp>(location, cellTupleType);
+                    cellValue = builder.create<mlir_ts::InsertPropertyOp>(location, cellTupleType, thunkValue, cellValue,
+                                                                          MLIRHelper::getStructIndex(builder, 0));
+                    cellValue = builder.create<mlir_ts::InsertPropertyOp>(location, cellTupleType, flagOffsetValue,
+                                                                          cellValue, MLIRHelper::getStructIndex(builder, 1));
+                    return TypeValueInitType{cellTupleType, cellValue, TypeProvided::Yes};
+                },
+                genContext);
+            if (!cellType)
+            {
+                return mlir::failure();
+            }
+        }
+
+        auto thunkFor = [&](mlir::Attribute id) -> StringRef {
+            auto found = llvm::find_if(thunksForFields, [&](auto &thunkForField) { return thunkForField.first == id; });
+            return found != thunksForFields.end() ? found->second : StringRef();
+        };
+
+        // register global
         registerVariable(
             location, fullClassInterfaceVTableFieldName, true, VariableType::Var,
             [&](mlir::Location location, const GenContext &genContext) {
@@ -540,16 +631,45 @@ namespace mlirgen
                         // PropertySignature-with-function-type's slot corrupts that
                         // computation (crashes on the resulting garbage address).
                         std::string methodFuncName;
+                        auto thunkName = methodOrField.isMissing ? StringRef() : thunkFor(methodOrField.fieldInfo.id);
+                        auto isMethod = false;
                         if (!methodOrField.isMissing)
                         {
                             auto fieldNameAttr = dyn_cast<mlir::StringAttr>(methodOrField.fieldInfo.id);
-                            if (fieldNameAttr && newInterfacePtr->findMethod(fieldNameAttr.getValue()))
+                            isMethod = fieldNameAttr && newInterfacePtr->findMethod(fieldNameAttr.getValue());
+                            if (isMethod)
                             {
-                                methodFuncName = lookupObjectLiteralMethodSymbol(methodOrField.fieldInfo.type, methodOrField.fieldInfo.id);
+                                methodFuncName = !thunkName.empty()
+                                    ? thunkName.str()
+                                    : lookupObjectLiteralMethodSymbol(methodOrField.fieldInfo.type, methodOrField.fieldInfo.id);
                             }
                         }
 
-                        if (!methodFuncName.empty())
+                        if (!thunkName.empty() && !isMethod)
+                        {
+                            // the slot is the address of the thunk's cell plus INTERFACE_METHOD_CELL_TAG, so an access
+                            // reads the cell instead of `this` + offset (InterfaceSymbolRefOpLowering)
+                            auto cellRef = resolveFullNameIdentifier(location, methodCellName(methodOrField.fieldInfo.id), true, genContext);
+                            if (!cellRef)
+                            {
+                                return TypeValueInitType{mlir::Type(), mlir::Value(), TypeProvided::Yes};
+                            }
+
+                            auto ptrIntType = builder.getIntegerType(compileOptions.sizeBits());
+                            auto cellAddress = builder.create<mlir_ts::CastOp>(location, ptrIntType, cellRef);
+                            auto tagValue = builder.create<mlir_ts::ConstantOp>(
+                                location, ptrIntType, builder.getIntegerAttr(ptrIntType, INTERFACE_METHOD_CELL_TAG));
+                            auto taggedValue = builder.create<mlir_ts::ArithmeticBinaryOp>(
+                                location, ptrIntType, builder.getI32IntegerAttr((int)SyntaxKind::PlusToken), cellAddress,
+                                tagValue);
+                            auto slotValue = builder.create<mlir_ts::CastOp>(
+                                location, mlir_ts::RefType::get(methodOrField.fieldInfo.type), taggedValue);
+
+                            vtableValue = builder.create<mlir_ts::InsertPropertyOp>(
+                                location, virtTuple, slotValue, vtableValue,
+                                MLIRHelper::getStructIndex(builder, fieldIndex));
+                        }
+                        else if (!methodFuncName.empty())
                         {
                             auto methodConstName = builder.create<mlir_ts::SymbolRefOp>(
                                 location, mlir::cast<mlir_ts::FunctionType>(methodOrField.fieldInfo.type),