From 42cc181fa02ad72e3d25d05d2b454c5af8590221 Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Fri, 9 Feb 2024 21:15:36 +0200 Subject: [PATCH] [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 --- .../kotlin/fir/plugin/CallShapeAttribute.kt | 29 ++ .../DataFrameLikeCallsRefinementExtension.kt | 253 ++++++++++++++++++ .../plugin/DataFrameLikeReturnTypeInjector.kt | 42 +++ .../FirPluginPrototypeExtensionRegistrar.kt | 6 + .../DataFrameLikeTypeMembersGenerator.kt | 125 +++++++++ .../receivers/callShapeBasedInjector.fir.txt | 42 +++ .../receivers/callShapeBasedInjector.kt | 13 + .../FirPsiPluginDiagnosticTestGenerated.java | 6 + 8 files changed, 516 insertions(+) create mode 100644 plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/CallShapeAttribute.kt create mode 100644 plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/DataFrameLikeCallsRefinementExtension.kt create mode 100644 plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/DataFrameLikeReturnTypeInjector.kt create mode 100644 plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/DataFrameLikeTypeMembersGenerator.kt create mode 100644 plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.fir.txt create mode 100644 plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.kt diff --git a/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/CallShapeAttribute.kt b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/CallShapeAttribute.kt new file mode 100644 index 00000000000..205ff25b41b --- /dev/null +++ b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/CallShapeAttribute.kt @@ -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) : CallShapeData + + class Scope(val token: FirClassSymbol, val columns: List) : CallShapeData + + class RefinedType(val scopes: List) : CallShapeData +} + + +object CallShapeAttribute : FirDeclarationDataKey() + +var FirClass.callShapeData: CallShapeData? by FirDeclarationDataRegistry.data(CallShapeAttribute) diff --git a/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/DataFrameLikeCallsRefinementExtension.kt b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/DataFrameLikeCallsRefinementExtension.kt new file mode 100644 index 00000000000..53126f0b8f4 --- /dev/null +++ b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/DataFrameLikeCallsRefinementExtension.kt @@ -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 = 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 { symbol, context -> + context ?: error("context not provided for $symbol") + } +} + +val FirSession.callDataStorage: CallDataStorage by FirSession.sessionComponentAccessor() + +class GeneratedCallData(val type: FirRegularClass, val schema: FirRegularClass) \ No newline at end of file diff --git a/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/DataFrameLikeReturnTypeInjector.kt b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/DataFrameLikeReturnTypeInjector.kt new file mode 100644 index 00000000000..16d96e32a79 --- /dev/null +++ b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/DataFrameLikeReturnTypeInjector.kt @@ -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 { + 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() + .filter { + val data = it.resolvedReturnType.toRegularClassSymbol(session)?.fir?.callShapeData ?: return@filter false + data is CallShapeData.Scope + } + .map { it.resolvedReturnType } + } +} \ No newline at end of file diff --git a/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/FirPluginPrototypeExtensionRegistrar.kt b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/FirPluginPrototypeExtensionRegistrar.kt index fc0fa2f4916..f6e60a3d625 100644 --- a/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/FirPluginPrototypeExtensionRegistrar.kt +++ b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/FirPluginPrototypeExtensionRegistrar.kt @@ -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 } } diff --git a/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/DataFrameLikeTypeMembersGenerator.kt b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/DataFrameLikeTypeMembersGenerator.kt new file mode 100644 index 00000000000..ab264a5c629 --- /dev/null +++ b/plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/DataFrameLikeTypeMembersGenerator.kt @@ -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, Map>?, 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 { + val properties = propertiesCache.getValue(classSymbol) + return properties?.flatMapTo(mutableSetOf(SpecialNames.INIT)) { it.value.map { it.name } } ?: emptySet() + } + + + override fun generateConstructors(context: MemberGenerationContext): List { + 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 { + 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 + } + } +} \ No newline at end of file diff --git a/plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.fir.txt b/plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.fir.txt new file mode 100644 index 00000000000..3e3d891ceff --- /dev/null +++ b/plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.fir.txt @@ -0,0 +1,42 @@ +FILE: callShapeBasedInjector.kt + public abstract interface DataFrame : R|kotlin/Any| { + } + public final annotation class Refine : R|kotlin/Annotation| { + public constructor(): R|Refine| { + super() + } + + } + @R|Refine|() public final fun R|DataFrame|.add(columnName: R|kotlin/String|, expression: R|() -> R|): R|DataFrame| { + ^add R|kotlin/TODO|() + } + public final fun test_1(df: R|DataFrame<*>|): R|kotlin/Unit| { + lval df1: R|DataFrame</DataFrameType>| = R|/df|.R|kotlin/let||, R|DataFrame</DataFrameType>|>( = fun (it: R|DataFrame<*>|): R|DataFrame</DataFrameType>| { + local abstract class Schema1 : R|kotlin/Any| { + local final val column: R|kotlin/Int| + local get(): R|kotlin/Int| + + local constructor(): R|/Schema1| + + } + + local abstract class Scope1 : R|kotlin/Any| { + local final val R|DataFrame</Schema1>|.column: R|kotlin/Int| + local get(): R|kotlin/Int| + + local constructor(): R|/Scope1| + + } + + local abstract class DataFrameType : R|/Schema1| { + local final val scope1: R|/Scope1| + local get(): R|/Scope1| + + local constructor(): R|/DataFrameType| + + } + + } + ) + lval col: R|kotlin/Int| = (this@R|/test_1|, R|/df1|).R|/Scope1.column| + } diff --git a/plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.kt b/plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.kt new file mode 100644 index 00000000000..adf28031cb2 --- /dev/null +++ b/plugins/fir-plugin-prototype/testData/diagnostics/receivers/callShapeBasedInjector.kt @@ -0,0 +1,13 @@ +import kotlin.reflect.KClass + +interface DataFrame + +annotation class Refine + +@Refine +fun DataFrame.add(columnName: String, expression: () -> R): DataFrame = TODO() + +fun test_1(df: DataFrame<*>) { + val df1 = df.add("column") { 1 } + val col = df1.column +} \ No newline at end of file diff --git a/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPsiPluginDiagnosticTestGenerated.java b/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPsiPluginDiagnosticTestGenerated.java index 1782d83b524..1ffb8d9315c 100644 --- a/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPsiPluginDiagnosticTestGenerated.java +++ b/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPsiPluginDiagnosticTestGenerated.java @@ -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() {