[FIR Plugin prototype] Implement DataFrame-like extension

Proof of concept that FirFunctionCallRefinementExtension (new) and existing extension points can be used together to update the return type of the specific calls and generate members based on call arguments. Important implementation details: - FirDeclarationGenerationExtension must be used to generate members of generated local classes. - FirExtensionSessionComponent together with firCachesFactory to pass information between `intercept` and `transform` Actual plugin is developed as a part of Kotlin dataframe repository KT-65859
2024-02-09 21:15:36 +02:00
parent c49edfef04
commit 42cc181fa0
8 changed files with 516 additions and 0 deletions
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.plugin
+
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataRegistry
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.name.Name
+
+class Column(val name: Name, val type: FirTypeRef)
+
+sealed interface CallShapeData {
+    class Schema(val columns: List<Column>) : CallShapeData
+
+    class Scope(val token: FirClassSymbol<out FirClass>, val columns: List<Column>) : CallShapeData
+
+    class RefinedType(val scopes: List<FirRegularClassSymbol>) : CallShapeData
+}
+
+
+object CallShapeAttribute : FirDeclarationDataKey()
+
+var FirClass.callShapeData: CallShapeData? by FirDeclarationDataRegistry.data(CallShapeAttribute)
@@ -0,0 +1,253 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.plugin
+
+import org.jetbrains.kotlin.GeneratedDeclarationKey
+import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.descriptors.EffectiveVisibility
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.getValue
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.builder.*
+import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.buildResolvedArgumentList
+import org.jetbrains.kotlin.fir.expressions.builder.buildAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildBlock
+import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
+import org.jetbrains.kotlin.fir.expressions.builder.buildLambdaArgumentExpression
+import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
+import org.jetbrains.kotlin.fir.extensions.FirExtensionSessionComponent
+import org.jetbrains.kotlin.fir.extensions.FirFunctionCallRefinementExtension
+import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.references.resolved
+import org.jetbrains.kotlin.fir.references.toResolvedNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.resolve.calls.CallInfo
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildTypeProjectionWithVariance
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.types.impl.FirImplicitAnyTypeRef
+import org.jetbrains.kotlin.fir.types.resolvedType
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.types.Variance
+
+@OptIn(FirExtensionApiInternals::class)
+class DataFrameLikeCallsRefinementExtension(session: FirSession) : FirFunctionCallRefinementExtension(session) {
+    companion object {
+        val REFINE = ClassId(FqName.ROOT, Name.identifier("Refine"))
+        val DATAFRAME = ClassId(FqName.ROOT, Name.identifier("DataFrame"))
+
+        object KEY : GeneratedDeclarationKey()
+    }
+
+    override fun intercept(callInfo: CallInfo, symbol: FirNamedFunctionSymbol): CallReturnType? {
+        if (!symbol.hasAnnotation(REFINE, session)) return null
+        val lookupTag = ConeClassLikeLookupTagImpl(DATAFRAME)
+        val refinedTypeId = localClassId(Name.identifier("DataFrameType"))
+        val schemaId = localClassId(Name.identifier("Schema1"))
+        val schemaSymbol = FirRegularClassSymbol(schemaId)
+        val schemaClass = buildRegularClass {
+            resolvePhase = FirResolvePhase.BODY_RESOLVE
+            moduleData = session.moduleData
+            origin = FirDeclarationOrigin.Source
+            status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.ABSTRACT, EffectiveVisibility.Local)
+            deprecationsProvider = EmptyDeprecationsProvider
+            classKind = ClassKind.CLASS
+            scopeProvider = FirKotlinScopeProvider()
+            superTypeRefs += FirImplicitAnyTypeRef(null)
+
+            name = schemaId.shortClassName
+            this.symbol = schemaSymbol
+        }
+
+
+        val refinedTypeSymbol = FirRegularClassSymbol(refinedTypeId)
+        val refinedTypeDeclaration = buildRegularClass {
+            resolvePhase = FirResolvePhase.BODY_RESOLVE
+            moduleData = session.moduleData
+            origin = FirDeclarationOrigin.Source
+            status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.ABSTRACT, EffectiveVisibility.Local)
+            deprecationsProvider = EmptyDeprecationsProvider
+            classKind = ClassKind.CLASS
+            scopeProvider = FirKotlinScopeProvider()
+
+            name = refinedTypeId.shortClassName
+            this.symbol = refinedTypeSymbol
+            superTypeRefs += buildResolvedTypeRef {
+                type = ConeClassLikeTypeImpl(
+                    ConeClassLookupTagWithFixedSymbol(schemaId, schemaSymbol),
+                    emptyArray(),
+                    isNullable = false
+                )
+            }
+        }
+
+        val typeRef = buildResolvedTypeRef {
+            type = ConeClassLikeTypeImpl(
+                lookupTag,
+                arrayOf(
+                    ConeClassLikeTypeImpl(
+                        ConeClassLookupTagWithFixedSymbol(refinedTypeId, refinedTypeSymbol),
+                        emptyArray(),
+                        isNullable = false
+                    )
+                ),
+                isNullable = false
+            )
+        }
+        return CallReturnType(typeRef) { functionSymbol ->
+            session.callDataStorage.generatedCallData.getValue(functionSymbol, GeneratedCallData(refinedTypeDeclaration, schemaClass))
+        }
+    }
+
+    @OptIn(SymbolInternals::class)
+    override fun transform(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall {
+        val resolvedLet = findLet()
+        val parameter = resolvedLet.valueParameterSymbols[0]
+
+        val explicitReceiver = call.explicitReceiver ?: return call
+        val receiverType = explicitReceiver.resolvedType
+        val returnType = call.resolvedType
+
+        val symbol = call.calleeReference.resolved?.toResolvedNamedFunctionSymbol() ?: return call
+        val callData = session.callDataStorage.generatedCallData.getValue(symbol)
+        val refinedType = callData.type
+        val schemaClass = callData.schema
+
+        val scope = localClassId(Name.identifier("Scope1"))
+        val scopeSymbol = FirRegularClassSymbol(scope)
+        val columns: List<Column> = listOf(Column(Name.identifier("column"), session.builtinTypes.intType))
+
+        schemaClass.callShapeData = CallShapeData.Schema(columns)
+
+        val scopeClass = buildRegularClass {
+            resolvePhase = FirResolvePhase.BODY_RESOLVE
+            moduleData = session.moduleData
+            origin = FirDeclarationOrigin.Source
+            status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.ABSTRACT, EffectiveVisibility.Local)
+            deprecationsProvider = EmptyDeprecationsProvider
+            classKind = ClassKind.CLASS
+            scopeProvider = FirKotlinScopeProvider()
+            superTypeRefs += FirImplicitAnyTypeRef(null)
+            name = scope.shortClassName
+            this.symbol = scopeSymbol
+        }
+
+        scopeClass.callShapeData = CallShapeData.Scope(schemaClass.symbol, columns)
+
+        refinedType.callShapeData = CallShapeData.RefinedType(listOf(scopeSymbol))
+
+        val argument = buildLambdaArgumentExpression {
+            expression = buildAnonymousFunctionExpression {
+                val fSymbol = FirAnonymousFunctionSymbol()
+                anonymousFunction = buildAnonymousFunction {
+                    resolvePhase = FirResolvePhase.BODY_RESOLVE
+                    moduleData = session.moduleData
+                    origin = FirDeclarationOrigin.Plugin(KEY)
+                    status = FirResolvedDeclarationStatusImpl(Visibilities.Local, Modality.FINAL, EffectiveVisibility.Local)
+                    deprecationsProvider = EmptyDeprecationsProvider
+                    returnTypeRef = buildResolvedTypeRef {
+                        type = returnType
+                    }
+                    val itName = Name.identifier("it")
+                    val parameterSymbol = FirValueParameterSymbol(itName)
+                    valueParameters += buildValueParameter {
+                        moduleData = session.moduleData
+                        origin = FirDeclarationOrigin.Plugin(KEY)
+                        returnTypeRef = buildResolvedTypeRef {
+                            type = receiverType
+                        }
+                        name = itName
+                        this.symbol = parameterSymbol
+                        containingFunctionSymbol = fSymbol
+                        isCrossinline = false
+                        isNoinline = false
+                        isVararg = false
+                    }.also { parameterSymbol.bind(it) }
+                    body = buildBlock {
+                        this.coneTypeOrNull = returnType
+
+                        // Schema is required for static extensions resolve and holds information for subsequent call modifications
+                        statements += schemaClass
+
+                        // Scope (provides extensions API)
+                        statements += scopeClass
+
+                        // Return type - dataframe schema
+                        statements += refinedType
+                    }
+                    this.symbol = fSymbol
+                    isLambda = true
+                    hasExplicitParameterList = false
+                    typeRef = buildResolvedTypeRef {
+                        type = ConeClassLikeTypeImpl(
+                            ConeClassLikeLookupTagImpl(ClassId(FqName("kotlin"), Name.identifier("Function1"))),
+                            typeArguments = arrayOf(receiverType, returnType),
+                            isNullable = false
+                        )
+                    }
+                    invocationKind = EventOccurrencesRange.EXACTLY_ONCE
+                    inlineStatus = InlineStatus.Inline
+                }.also { fSymbol.bind(it) }
+            }
+        }
+
+        val newCall = buildFunctionCall {
+            this.coneTypeOrNull = returnType
+            typeArguments += buildTypeProjectionWithVariance {
+                typeRef = buildResolvedTypeRef {
+                    type = receiverType
+                }
+                variance = Variance.INVARIANT
+            }
+
+            typeArguments += buildTypeProjectionWithVariance {
+                typeRef = buildResolvedTypeRef {
+                    type = returnType
+                }
+                variance = Variance.INVARIANT
+            }
+            dispatchReceiver = call.dispatchReceiver
+            this.explicitReceiver = call.explicitReceiver
+            extensionReceiver = call.extensionReceiver
+            argumentList = buildResolvedArgumentList(linkedMapOf(argument to parameter.fir))
+            calleeReference = buildResolvedNamedReference {
+                name = Name.identifier("let")
+                resolvedSymbol = resolvedLet
+            }
+        }
+        return newCall
+    }
+
+    private fun localClassId(name: Name) = ClassId(CallableId.PACKAGE_FQ_NAME_FOR_LOCAL, FqName.ROOT.child(name), isLocal = true)
+
+    private fun findLet(): FirFunctionSymbol<*> {
+        return session.symbolProvider.getTopLevelFunctionSymbols(FqName("kotlin"), Name.identifier("let")).single()
+    }
+}
+
+class CallDataStorage(session: FirSession) : FirExtensionSessionComponent(session) {
+    val generatedCallData =
+        session.firCachesFactory.createCache<FirNamedFunctionSymbol, GeneratedCallData, GeneratedCallData?> { symbol, context ->
+            context ?: error("context not provided for $symbol")
+        }
+}
+
+val FirSession.callDataStorage: CallDataStorage by FirSession.sessionComponentAccessor()
+
+class GeneratedCallData(val type: FirRegularClass, val schema: FirRegularClass)
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.plugin
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.extensions.FirExpressionResolutionExtension
+import org.jetbrains.kotlin.fir.scopes.collectAllProperties
+import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+
+class DataFrameLikeReturnTypeInjector(session: FirSession) : FirExpressionResolutionExtension(session) {
+    companion object {
+        val DF_CLASS_ID: ClassId = ClassId.topLevel(FqName.fromSegments(listOf("DataFrame")))
+    }
+
+    @OptIn(SymbolInternals::class)
+    override fun addNewImplicitReceivers(functionCall: FirFunctionCall): List<ConeKotlinType> {
+        val callReturnType = functionCall.resolvedType
+        if (callReturnType.classId != DF_CLASS_ID) return emptyList()
+        val rootMarker = callReturnType.typeArguments[0]
+        if (rootMarker !is ConeClassLikeType) {
+            return emptyList()
+        }
+        val symbol = rootMarker.toRegularClassSymbol(session) ?: return emptyList()
+        return symbol.declaredMemberScope(session, FirResolvePhase.DECLARATIONS).collectAllProperties()
+            .filterIsInstance<FirPropertySymbol>()
+            .filter {
+                val data = it.resolvedReturnType.toRegularClassSymbol(session)?.fir?.callShapeData ?: return@filter false
+                data is CallShapeData.Scope
+            }
+            .map { it.resolvedReturnType }
+    }
+}
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.plugin
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 import org.jetbrains.kotlin.fir.plugin.generators.*
@@ -27,6 +28,10 @@ class FirPluginPrototypeExtensionRegistrar : FirExtensionRegistrar() {
        +::FirNumberSignAttributeExtension
        +::AlgebraReceiverInjector
        +::ComposableLikeFunctionTypeKindExtension
+        @OptIn(FirExtensionApiInternals::class)
+        +::DataFrameLikeCallsRefinementExtension
+        +::DataFrameLikeReturnTypeInjector
+        +::CallDataStorage

        // Declaration generators
        +::TopLevelDeclarationsGenerator
@@ -35,6 +40,7 @@ class FirPluginPrototypeExtensionRegistrar : FirExtensionRegistrar() {
        +::AdditionalMembersGenerator
        +::CompanionGenerator
        +::MembersOfSerializerGenerator
+        +::DataFrameLikeTypeMembersGenerator
    }
 }

@@ -0,0 +1,125 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.plugin.generators
+
+import org.jetbrains.kotlin.descriptors.Visibilities
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.createCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.getValue
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
+import org.jetbrains.kotlin.fir.plugin.*
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.SpecialNames
+
+class DataFrameLikeTypeMembersGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
+
+    @OptIn(SymbolInternals::class)
+    private val propertiesCache: FirCache<FirClassSymbol<*>, Map<Name, List<FirProperty>>?, Nothing?> =
+        session.firCachesFactory.createCache { k ->
+            val callShapeData = k.fir.callShapeData ?: return@createCache null
+            when (callShapeData) {
+                is CallShapeData.RefinedType -> callShapeData.scopes.associate {
+                    val propertyName = Name.identifier(it.name.asString().replaceFirstChar { it.lowercaseChar() })
+                    propertyName to listOf(buildScopeReferenceProperty(it.classId, it, propertyName))
+                }
+                is CallShapeData.Scope -> callShapeData.columns.associate {
+                    it.name to listOf(buildScopeApiProperty(callShapeData.token, scopeSymbol = k, it.name))
+                }
+                is CallShapeData.Schema -> callShapeData.columns.associate {
+                    it.name to listOf(buildTokenProperty(k, it.name))
+                }
+            }
+        }
+
+
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
+        val properties = propertiesCache.getValue(classSymbol)
+        return properties?.flatMapTo(mutableSetOf(SpecialNames.INIT)) { it.value.map { it.name } } ?: emptySet()
+    }
+
+
+    override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
+        val constructor = createConstructor(context.owner, DataFrameLikeCallsRefinementExtension.Companion.KEY, isPrimary = true) {
+            visibility = Visibilities.Local
+        }
+        return listOf(constructor.symbol)
+    }
+
+    override fun generateProperties(callableId: CallableId, context: MemberGenerationContext?): List<FirPropertySymbol> {
+        val owner = context?.owner ?: return emptyList()
+        return propertiesCache.getValue(owner)?.flatMap { it.value.map { it.symbol } } ?: emptyList()
+    }
+
+    private fun buildScopeApiProperty(
+        tokenSymbol: FirClassSymbol<*>,
+        scopeSymbol: FirClassSymbol<*>,
+        propName: Name,
+    ): FirProperty {
+        return createMemberProperty(
+            scopeSymbol,
+            DataFrameLikeCallsRefinementExtension.Companion.KEY,
+            propName,
+            session.builtinTypes.intType.type
+        ) {
+            visibility = Visibilities.Local
+            extensionReceiverType {
+                ConeClassLikeTypeImpl(
+                    ConeClassLikeLookupTagImpl(DataFrameLikeCallsRefinementExtension.DATAFRAME),
+                    arrayOf(
+                        ConeClassLikeTypeImpl(
+                            ConeClassLookupTagWithFixedSymbol(tokenSymbol.classId, tokenSymbol),
+                            emptyArray(),
+                            isNullable = false
+                        )
+                    ),
+                    isNullable = false
+                )
+            }
+        }
+    }
+
+    private fun buildTokenProperty(
+        tokenSymbol: FirClassSymbol<*>,
+        propName: Name,
+    ): FirProperty {
+        return createMemberProperty(
+            tokenSymbol,
+            DataFrameLikeCallsRefinementExtension.Companion.KEY,
+            propName,
+            session.builtinTypes.intType.type
+        ) {
+            visibility = Visibilities.Local
+        }
+    }
+
+    private fun buildScopeReferenceProperty(
+        scope: ClassId,
+        scopeSymbol: FirRegularClassSymbol,
+        name: Name
+    ): FirProperty {
+        return createMemberProperty(
+            scopeSymbol,
+            DataFrameLikeCallsRefinementExtension.Companion.KEY,
+            name,
+            ConeClassLikeTypeImpl(
+                ConeClassLookupTagWithFixedSymbol(scope, scopeSymbol),
+                emptyArray(),
+                isNullable = false
+            )
+        ) {
+            visibility = Visibilities.Local
+        }
+    }
+}
@@ -0,0 +1,42 @@
+FILE: callShapeBasedInjector.kt
+    public abstract interface DataFrame<out T> : R|kotlin/Any| {
+    }
+    public final annotation class Refine : R|kotlin/Annotation| {
+        public constructor(): R|Refine| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    @R|Refine|() public final fun <T, R> R|DataFrame<T>|.add(columnName: R|kotlin/String|, expression: R|() -> R|): R|DataFrame<kotlin/Any?>| {
+        ^add R|kotlin/TODO|()
+    }
+    public final fun test_1(df: R|DataFrame<*>|): R|kotlin/Unit| {
+        lval df1: R|DataFrame<<local>/DataFrameType>| = R|<local>/df|.R|kotlin/let|<R|DataFrame<*>|, R|DataFrame<<local>/DataFrameType>|>(<L> = fun <anonymous>(it: R|DataFrame<*>|): R|DataFrame<<local>/DataFrameType>| <inline=Inline, kind=EXACTLY_ONCE>  {
+            local abstract class Schema1 : R|kotlin/Any| {
+                local final val column: R|kotlin/Int|
+                    local get(): R|kotlin/Int|
+
+                local constructor(): R|<local>/Schema1|
+
+            }
+
+            local abstract class Scope1 : R|kotlin/Any| {
+                local final val R|DataFrame<<local>/Schema1>|.column: R|kotlin/Int|
+                    local get(): R|kotlin/Int|
+
+                local constructor(): R|<local>/Scope1|
+
+            }
+
+            local abstract class DataFrameType : R|<local>/Schema1| {
+                local final val scope1: R|<local>/Scope1|
+                    local get(): R|<local>/Scope1|
+
+                local constructor(): R|<local>/DataFrameType|
+
+            }
+
+        }
+        )
+        lval col: R|kotlin/Int| = (this@R|/test_1|, R|<local>/df1|).R|<local>/Scope1.column|
+    }
@@ -0,0 +1,13 @@
+import kotlin.reflect.KClass
+
+interface DataFrame<out T>
+
+annotation class Refine
+
+@Refine
+fun <T, R> DataFrame<T>.add(columnName: String, expression: () -> R): DataFrame<Any?> = TODO()
+
+fun test_1(df: DataFrame<*>) {
+    val df1 = df.add("column") { 1 }
+    val col = df1.column
+}
@@ -147,6 +147,12 @@ public class FirPsiPluginDiagnosticTestGenerated extends AbstractFirPsiPluginDia
      KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("plugins/fir-plugin-prototype/testData/diagnostics/receivers"), Pattern.compile("^(.+)\\.kt$"), null, true);
    }

+    @Test
+    @TestMetadata("callShapeBasedInjector.kt")
+    public void testCallShapeBasedInjector() {
+      runTest("plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.kt");
+    }
+
    @Test
    @TestMetadata("receiverInjection.kt")
    public void testReceiverInjection() {