[KDF] Postpone access to @ColumnName argument from too early STATUS phase

generateProperties can happen at Phase.STATUS
arguments are available later at ANNOTATION_ARGUMENTS
Now value is accessed at backend IR lowering phase

Author: koperagen

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

@@ -109,10 +109,10 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
         val call = IrCallImpl(-1, -1, context.irBuiltIns.anyNType, get, 0).also {
             val thisSymbol: IrValueSymbol = getterExtensionReceiver.symbol
             it.arguments[0] = IrGetValueImpl(-1, -1, thisSymbol)
-            val annotation = declaration.annotations.findAnnotation(Names.COLUMN_NAME_ANNOTATION.asSingleFqName())
-            val columnName = (annotation?.arguments?.get(0) as? IrConst)?.value as? String
-            val columName = columnName ?: declaration.name.identifier
-            it.arguments[1] = IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, columName)
+            val columnName = marker.properties.firstOrNull { it.name == declaration.name }
+                ?.getAnnotationArgumentValue<String>(Names.COLUMN_NAME_ANNOTATION.asSingleFqName(), Names.COLUMN_NAME_ARGUMENT.identifier)
+            val arg = columnName ?: declaration.name.identifier
+            it.arguments[1] = IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, arg)
         }
 
         val typeOp = IrTypeOperatorCallImpl(-1, -1, returnType, IrTypeOperator.CAST, returnType, call)

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

@@ -121,7 +121,6 @@ class ImportedSchemasGenerator(
                     it.owner,
                     mode,
                     it.returnType,
-                    null,
                     callableId.callableName
                 )
             }.orEmpty()

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

@@ -65,7 +65,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
                         receiverType = Names.DATA_ROW_CLASS_ID.constructClassLikeType(
                             typeArguments = arrayOf(schemaProperty.marker)
                         ),
-                        propertyName = propertyName,
+                        propertyName = propertyName.identifier,
                         returnType = schemaProperty.dataRowReturnType,
                         symbol = k,
                         effectiveVisibility = EffectiveVisibility.Local,
@@ -77,7 +77,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
                         receiverType = Names.COLUMNS_SCOPE_CLASS_ID.constructClassLikeType(
                             typeArguments = arrayOf(schemaProperty.marker)
                         ),
-                        propertyName = propertyName,
+                        propertyName = propertyName.identifier,
                         returnType = schemaProperty.columnContainerReturnType,
                         symbol = k,
                         effectiveVisibility = EffectiveVisibility.Local,

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

@@ -5,10 +5,8 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.builder.buildTypeParameter
 import org.jetbrains.kotlin.fir.declarations.declaredProperties
-import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
 import org.jetbrains.kotlin.fir.declarations.hasAnnotation
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
-import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
 import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
 import org.jetbrains.kotlin.fir.moduleData
@@ -27,7 +25,6 @@ import org.jetbrains.kotlin.name.packageName
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
 import org.jetbrains.kotlinx.dataframe.plugin.DataFramePlugin
-import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
 import org.jetbrains.kotlinx.dataframe.plugin.utils.CallableIdOrSymbol
 import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
@@ -89,9 +86,6 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio
                 owner,
                 mode,
                 property.resolvedReturnType,
-                property.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, this@TopLevelExtensionsGenerator.session)?.let { annotation ->
-                    (annotation.argumentMapping.mapping[Names.COLUMN_NAME_ARGUMENT] as? FirLiteralExpression)?.value as? String?
-                },
                 property.name
             )
         }
@@ -111,12 +105,9 @@ fun FirDeclarationGenerationExtension.buildExtensionPropertiesApi(
     owner: FirRegularClassSymbol,
     mode: TopLevelExtensionsGenerator.Receiver,
     resolvedReturnType: ConeKotlinType,
-    columnName: String?,
     name: Name
 ): FirPropertySymbol {
     var resolvedReturnType = resolvedReturnType
-    val columnName = columnName
-
     val firPropertySymbol = FirRegularPropertySymbol(callableId)
 
     val typeParameters = owner.typeParameterSymbols.map {
@@ -182,7 +173,7 @@ fun FirDeclarationGenerationExtension.buildExtensionPropertiesApi(
                 typeArguments = arrayOf(marker),
                 isMarkedNullable = false
             ),
-            propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
+            propertyName = name,
             returnType = resolvedReturnType,
             source = owner.source,
             typeParameters = typeParameters,
@@ -193,7 +184,7 @@ fun FirDeclarationGenerationExtension.buildExtensionPropertiesApi(
                 typeArguments = arrayOf(marker),
                 isMarkedNullable = false
             ),
-            propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
+            propertyName = name,
             returnType = columnReturnType,
             source = owner.source,
             typeParameters = typeParameters,

plugins/kotlin-dataframe/kotlin-dataframe.k2/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/PropertyName.kt

@@ -44,9 +44,5 @@ data class PropertyName(val identifier: Name, val columnNameAnnotation: FirAnnot
                 }
             }
         }
-
-        fun of(identifier: Name, columnNameAnnotation: FirAnnotation?): PropertyName {
-            return PropertyName(identifier, columnNameAnnotation)
-        }
     }
 }

plugins/kotlin-dataframe/kotlin-dataframe.k2/src/org/jetbrains/kotlinx/dataframe/plugin/utils/firFactories.kt

@@ -20,13 +20,13 @@ import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.constructClassType
 import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlinx.dataframe.plugin.DataFramePlugin
-import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName
 
 internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
     callableIdOrSymbol: CallableIdOrSymbol,
     receiverType: ConeClassLikeType,
-    propertyName: PropertyName,
+    propertyName: Name,
     returnType: ConeKotlinType,
     symbol: FirClassSymbol<*>? = null,
     effectiveVisibility: EffectiveVisibility = EffectiveVisibility.Public,
@@ -40,9 +40,6 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
 
     return buildProperty {
         this.source = source
-        propertyName.columnNameAnnotation?.let {
-            annotations += it
-        }
         moduleData = session.moduleData
         resolvePhase = FirResolvePhase.BODY_RESOLVE
         origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
@@ -84,7 +81,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
                 effectiveVisibility
             )
         }
-        name = propertyName.identifier
+        name = propertyName
         this.symbol = firPropertySymbol
         isVar = false
     }

plugins/kotlin-dataframe/testData/box/resolveGeneratedPropertyOnStatusPhase.kt

@@ -0,0 +1,18 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+data class MySchema(
+    @ColumnName("Col")
+    val col: Int
+)
+
+fun box(): String {
+    dataFrameOf("Col" to listOf(1, 2, 3)).cast<MySchema>().col
+    return "OK"
+}
+
+// triggers property generation before annotation arguments are resolved
+private fun DataFrame<MySchema>.convert() = col