From 52108637770bffdc3b133b1d54bf064b29afdecd Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Wed, 29 Jan 2025 20:49:12 +0200 Subject: [PATCH 1/2] Map transformed call to its original source --- .../dataframe/plugin/extensions/FunctionCallTransformer.kt | 1 + 1 file changed, 1 insertion(+) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt index e882cfcfd2..2330d18496 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt @@ -328,6 +328,7 @@ class FunctionCallTransformer( return if (element is FirResolvedNamedReference) { @Suppress("UNCHECKED_CAST") buildResolvedNamedReference { + source = call.calleeReference.source this.name = element.name resolvedSymbol = originalSymbol } as E From 48dab7a911f2c1d2c16fda86c84f692fb5fcd65a Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Wed, 29 Jan 2025 21:16:39 +0200 Subject: [PATCH 2/2] Point plugin-generated properties to their origin --- .../dataframe/plugin/extensions/CallShapeAttribute.kt | 3 ++- .../plugin/extensions/FunctionCallTransformer.kt | 2 +- .../kotlinx/dataframe/plugin/extensions/TokenGenerator.kt | 8 +++++--- .../plugin/extensions/TopLevelExtensionsGenerator.kt | 6 ++++-- .../kotlinx/dataframe/plugin/utils/firFactories.kt | 5 ++++- 5 files changed, 16 insertions(+), 8 deletions(-) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/CallShapeAttribute.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/CallShapeAttribute.kt index 0058198b44..6ece6eeaf7 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/CallShapeAttribute.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/CallShapeAttribute.kt @@ -1,5 +1,6 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions +import org.jetbrains.kotlin.KtSourceElement import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty import org.jetbrains.kotlin.fir.declarations.FirClass import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey @@ -9,7 +10,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol sealed interface CallShapeData { class Schema(val columns: List) : CallShapeData - class Scope(val columns: List) : CallShapeData + class Scope(val columns: List, val source: KtSourceElement?) : CallShapeData class RefinedType(val scopes: List) : CallShapeData } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt index 2330d18496..7668541839 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt @@ -556,7 +556,7 @@ class FunctionCallTransformer( } } schema.callShapeData = CallShapeData.Schema(properties) - scope.callShapeData = CallShapeData.Scope(properties) + scope.callShapeData = CallShapeData.Scope(properties, call.calleeReference.source) val schemaApi = DataSchemaApi(schema, scope) dataSchemaApis.add(schemaApi) return schemaApi diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt index 03e8009690..12ed210fc1 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt @@ -59,7 +59,6 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se val callableId = CallableId(k.classId, propertyName.identifier) val dataRowExtension = generateExtensionProperty( callableId = callableId, - symbol = k, receiverType = ConeClassLikeTypeImpl( ConeClassLikeLookupTagImpl(Names.DATA_ROW_CLASS_ID), typeArguments = arrayOf(schemaProperty.marker), @@ -67,7 +66,9 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se ), propertyName = propertyName, returnTypeRef = schemaProperty.dataRowReturnType.toFirResolvedTypeRef(), - effectiveVisibility = EffectiveVisibility.Local + symbol = k, + effectiveVisibility = EffectiveVisibility.Local, + source = callShapeData.source ) val columnContainerExtension = generateExtensionProperty( @@ -80,7 +81,8 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se propertyName = propertyName, returnTypeRef = schemaProperty.columnContainerReturnType.toFirResolvedTypeRef(), symbol = k, - effectiveVisibility = EffectiveVisibility.Local + effectiveVisibility = EffectiveVisibility.Local, + source = callShapeData.source ) propertyName.identifier to listOf(dataRowExtension, columnContainerExtension) } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt index ff2435c971..f6e3e0fdf3 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt @@ -120,7 +120,8 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio isNullable = false ), propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }), - returnTypeRef = resolvedReturnTypeRef + returnTypeRef = resolvedReturnTypeRef, + source = owner.source ) val columnReturnType = when { @@ -142,7 +143,8 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio isNullable = false ), propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }), - returnTypeRef = columnReturnType + returnTypeRef = columnReturnType, + source = owner.source ) listOf(rowExtension.symbol, columnsContainerExtension.symbol) } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/firFactories.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/firFactories.kt index e6fa15743b..acaffd56c2 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/firFactories.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/firFactories.kt @@ -1,5 +1,6 @@ package org.jetbrains.kotlinx.dataframe.plugin.utils +import org.jetbrains.kotlin.KtSourceElement import org.jetbrains.kotlin.descriptors.EffectiveVisibility import org.jetbrains.kotlin.descriptors.Modality import org.jetbrains.kotlin.descriptors.Visibilities @@ -30,10 +31,12 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty( propertyName: PropertyName, returnTypeRef: FirResolvedTypeRef, symbol: FirClassSymbol<*>? = null, - effectiveVisibility: EffectiveVisibility = EffectiveVisibility.Public + effectiveVisibility: EffectiveVisibility = EffectiveVisibility.Public, + source: KtSourceElement? ): FirProperty { val firPropertySymbol = FirPropertySymbol(callableId) return buildProperty { + this.source = source propertyName.columnNameAnnotation?.let { annotations += it }