Map transformed calls and generated properties to their origin

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<SchemaProperty>) : CallShapeData
 
-    class Scope(val columns: List<SchemaProperty>) : CallShapeData
+    class Scope(val columns: List<SchemaProperty>, val source: KtSourceElement?) : CallShapeData
 
     class RefinedType(val scopes: List<FirRegularClassSymbol>) : CallShapeData
 }

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
@@ -555,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

plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt

@@ -59,15 +59,16 @@ 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),
                             isNullable = false
                         ),
                         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)
                 }

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)
             }

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
         }