From 3bca7599c79e5ddb6d98f4216401912fa78c971f Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Wed, 9 Oct 2024 13:43:10 +0300 Subject: [PATCH 1/3] [Compiler plugin] Rename ExtensionsGenerator --- .../dataframe/plugin/FirDataFrameComponentRegistrar.kt | 4 ++-- ...xtensionsGenerator.kt => TopLevelExtensionsGenerator.kt} | 6 +++++- 2 files changed, 7 insertions(+), 3 deletions(-) rename plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/{ExtensionsGenerator.kt => TopLevelExtensionsGenerator.kt} (96%) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/FirDataFrameComponentRegistrar.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/FirDataFrameComponentRegistrar.kt index 3d05e53fbc..68f4accc73 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/FirDataFrameComponentRegistrar.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/FirDataFrameComponentRegistrar.kt @@ -19,7 +19,7 @@ import org.jetbrains.kotlin.fir.caches.FirCache import org.jetbrains.kotlin.fir.caches.firCachesFactory import org.jetbrains.kotlinx.dataframe.plugin.extensions.DataRowSchemaSupertype import org.jetbrains.kotlinx.dataframe.plugin.extensions.ExpressionAnalysisAdditionalChecker -import org.jetbrains.kotlinx.dataframe.plugin.extensions.ExtensionsGenerator +import org.jetbrains.kotlinx.dataframe.plugin.extensions.TopLevelExtensionsGenerator import org.jetbrains.kotlinx.dataframe.plugin.extensions.FunctionCallTransformer import org.jetbrains.kotlinx.dataframe.plugin.extensions.IrBodyFiller import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade @@ -68,7 +68,7 @@ class FirDataFrameExtensionRegistrar( ) : FirExtensionRegistrar() { @OptIn(FirExtensionApiInternals::class) override fun ExtensionRegistrarContext.configurePlugin() { - +::ExtensionsGenerator + +::TopLevelExtensionsGenerator +::ReturnTypeBasedReceiverInjector +{ it: FirSession -> FunctionCallTransformer(path, it, jsonCache(it), schemasDirectory, isTest) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExtensionsGenerator.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt similarity index 96% rename from plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExtensionsGenerator.kt rename to plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt index a5b26bc6b5..7c634dedfb 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ExtensionsGenerator.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TopLevelExtensionsGenerator.kt @@ -33,7 +33,11 @@ import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.types.Variance import org.jetbrains.kotlinx.dataframe.annotations.DataSchema -class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) { +/** + * extensions inside scope classes are generated here: + * @see org.jetbrains.kotlinx.dataframe.plugin.extensions.TokenGenerator + */ +class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) { private companion object { val dataSchema = FqName(DataSchema::class.qualifiedName!!) } From 883d45146d5711815175b09f585e79ef308166e0 Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Wed, 9 Oct 2024 13:50:17 +0300 Subject: [PATCH 2/3] [Compiler plugin] Add empty test for debugging purposes --- plugins/kotlin-dataframe/testData/box/playground.kt | 9 +++++++++ .../dataframe/DataFrameBlackBoxCodegenTestGenerated.java | 6 ++++++ 2 files changed, 15 insertions(+) create mode 100644 plugins/kotlin-dataframe/testData/box/playground.kt diff --git a/plugins/kotlin-dataframe/testData/box/playground.kt b/plugins/kotlin-dataframe/testData/box/playground.kt new file mode 100644 index 0000000000..fb4913d5c9 --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/playground.kt @@ -0,0 +1,9 @@ +import org.jetbrains.kotlinx.dataframe.* +import org.jetbrains.kotlinx.dataframe.annotations.* +import org.jetbrains.kotlinx.dataframe.api.* +import org.jetbrains.kotlinx.dataframe.io.* + +fun box(): String { + + return "OK" +} diff --git a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java index 605e49f492..0f12c6ec7a 100644 --- a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java +++ b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java @@ -304,6 +304,12 @@ public void testPlatformType() { runTest("testData/box/platformType.kt"); } + @Test + @TestMetadata("playground.kt") + public void testPlayground() { + runTest("testData/box/playground.kt"); + } + @Test @TestMetadata("propertiesOrder.kt") public void testPropertiesOrder() { From bd33f628f4d1afc7c7adb0263cfa84cdafa666bb Mon Sep 17 00:00:00 2001 From: Nikita Klimenko Date: Wed, 9 Oct 2024 13:56:18 +0300 Subject: [PATCH 3/3] [Compiler plugin] Generate ColumnName annotations on frontend for all names that contain illegal characters The way ColumnName is used: 1. When extracting schemas. PluginDataFrameSchema always contains actual names, and it's responsibility of the plugin to generate valid property name 2. When generating getter of extension property. --- .../plugin/analyzeRefinedCallShape.kt | 10 ---- .../plugin/extensions/CallShapeAttribute.kt | 2 +- .../plugin/extensions/DataFramePlugin.kt | 2 +- .../extensions/FunctionCallTransformer.kt | 9 ++-- .../plugin/extensions/IrBodyFiller.kt | 6 ++- .../plugin/extensions/TokenGenerator.kt | 18 ++++--- .../extensions/TopLevelExtensionsGenerator.kt | 10 ++-- .../plugin/extensions/impl/PropertyName.kt | 52 +++++++++++++++++++ .../plugin/extensions/impl/SchemaProperty.kt | 12 +++++ .../dataframe/plugin/utils/firFactories.kt | 14 ++--- .../testData/box/columnName_invalidSymbol.kt | 10 ++++ ...DataFrameBlackBoxCodegenTestGenerated.java | 6 +++ 12 files changed, 115 insertions(+), 36 deletions(-) create mode 100644 plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/PropertyName.kt create mode 100644 plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/SchemaProperty.kt create mode 100644 plugins/kotlin-dataframe/testData/box/columnName_invalidSymbol.kt diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/analyzeRefinedCallShape.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/analyzeRefinedCallShape.kt index 7c16ee968e..00f678811a 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/analyzeRefinedCallShape.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/analyzeRefinedCallShape.kt @@ -8,8 +8,6 @@ package org.jetbrains.kotlinx.dataframe.plugin import org.jetbrains.kotlin.fir.expressions.FirExpression import org.jetbrains.kotlin.fir.expressions.FirFunctionCall import org.jetbrains.kotlin.fir.types.ConeClassLikeType -import org.jetbrains.kotlin.fir.types.ConeKotlinType -import org.jetbrains.kotlin.fir.types.ConeTypeProjection import org.jetbrains.kotlin.fir.types.classId import org.jetbrains.kotlin.fir.types.resolvedType import org.jetbrains.kotlin.name.ClassId @@ -63,11 +61,3 @@ data class RefinedArgument(val name: Name, val expression: FirExpression) { return "RefinedArgument(name=$name, expression=${expression})" } } - -data class SchemaProperty( - val marker: ConeTypeProjection, - val name: String, - val dataRowReturnType: ConeKotlinType, - val columnContainerReturnType: ConeKotlinType, - val override: Boolean = false -) 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 a74d07d172..0058198b44 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,6 +1,6 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions -import org.jetbrains.kotlinx.dataframe.plugin.SchemaProperty +import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty import org.jetbrains.kotlin.fir.declarations.FirClass import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataRegistry diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/DataFramePlugin.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/DataFramePlugin.kt index 695b42ed06..b158414ff0 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/DataFramePlugin.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/DataFramePlugin.kt @@ -2,4 +2,4 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions import org.jetbrains.kotlin.GeneratedDeclarationKey -class DataFramePlugin(val columnName: String?) : GeneratedDeclarationKey() +data object DataFramePlugin : GeneratedDeclarationKey() 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 c741de7fbf..c7b6af6766 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 @@ -13,7 +13,7 @@ import org.jetbrains.kotlin.fir.FirSession import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId import org.jetbrains.kotlin.fir.caches.FirCache import org.jetbrains.kotlinx.dataframe.plugin.InterpretationErrorReporter -import org.jetbrains.kotlinx.dataframe.plugin.SchemaProperty +import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.SchemaProperty import org.jetbrains.kotlinx.dataframe.plugin.analyzeRefinedCallShape import org.jetbrains.kotlinx.dataframe.plugin.utils.Names import org.jetbrains.kotlinx.dataframe.plugin.utils.projectOverDataColumnType @@ -74,6 +74,7 @@ import org.jetbrains.kotlin.name.FqName import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.text import org.jetbrains.kotlin.types.Variance +import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn @@ -520,7 +521,7 @@ class FunctionCallTransformer( isNullable = false ) - SchemaProperty(schema.defaultType(), it.name, dataRowReturnType, columnsContainerReturnType) + SchemaProperty(schema.defaultType(), PropertyName.of(it.name), dataRowReturnType, columnsContainerReturnType) } is SimpleFrameColumn -> { @@ -534,7 +535,7 @@ class FunctionCallTransformer( SchemaProperty( marker = schema.defaultType(), - name = it.name, + propertyName = PropertyName.of(it.name), dataRowReturnType = frameColumnReturnType, columnContainerReturnType = frameColumnReturnType.toFirResolvedTypeRef() .projectOverDataColumnType() @@ -543,7 +544,7 @@ class FunctionCallTransformer( is SimpleDataColumn -> SchemaProperty( marker = schema.defaultType(), - name = it.name, + propertyName = PropertyName.of(it.name), dataRowReturnType = it.type.type(), columnContainerReturnType = it.type.type().toFirResolvedTypeRef().projectOverDataColumnType() ) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt index 4e81d8764d..fa395d4c9b 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/IrBodyFiller.kt @@ -48,6 +48,7 @@ import org.jetbrains.kotlin.ir.types.classOrFail import org.jetbrains.kotlin.ir.types.classifierOrNull import org.jetbrains.kotlin.ir.types.getClass import org.jetbrains.kotlin.ir.util.constructors +import org.jetbrains.kotlin.ir.util.findAnnotation import org.jetbrains.kotlin.ir.util.parentAsClass import org.jetbrains.kotlin.ir.util.primaryConstructor import org.jetbrains.kotlin.ir.util.superTypes @@ -63,6 +64,7 @@ import org.jetbrains.kotlinx.dataframe.api.schema import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup import org.jetbrains.kotlinx.dataframe.plugin.impl.data.IoSchema import org.jetbrains.kotlinx.dataframe.plugin.impl.data.serialize +import org.jetbrains.kotlinx.dataframe.plugin.utils.Names import java.io.File class IrBodyFiller( @@ -209,7 +211,9 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering val call = IrCallImpl(-1, -1, context.irBuiltIns.anyNType, get, 0, 1).also { val thisSymbol: IrValueSymbol = getter.extensionReceiverParameter?.symbol!! it.dispatchReceiver = IrGetValueImpl(-1, -1, thisSymbol) - val columName = pluginKey.columnName ?: declaration.name.identifier + val annotation = declaration.annotations.findAnnotation(Names.COLUMN_NAME_ANNOTATION.asSingleFqName()) + val columnName = (annotation?.valueArguments?.get(0) as? IrConst<*>)?.value as? String + val columName = columnName ?: declaration.name.identifier it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, columName)) } 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 739b0ab9c8..0fed46fcd6 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 @@ -47,16 +47,16 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se val resolvedTypeRef = buildResolvedTypeRef { type = property.dataRowReturnType } - val propertyName = Name.identifier(property.name) - propertyName to listOf(buildProperty(resolvedTypeRef, propertyName, k, order = index)) + val identifier = property.propertyName.identifier + identifier to listOf(buildProperty(resolvedTypeRef, identifier, k, property.propertyName.columnNameAnnotation, order = index)) } is CallShapeData.RefinedType -> callShapeData.scopes.associate { - val propertyName = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() }) - propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isScopeProperty = true)) + val identifier = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() }) + identifier to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), identifier, k, isScopeProperty = true)) } is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty -> - val propertyName = Name.identifier(schemaProperty.name) - val callableId = CallableId(k.classId, propertyName) + val propertyName = schemaProperty.propertyName + val callableId = CallableId(k.classId, propertyName.identifier) val dataRowExtension = generateExtensionProperty( callableId = callableId, symbol = k, @@ -82,7 +82,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se symbol = k, effectiveVisibility = EffectiveVisibility.Local ) - propertyName to listOf(dataRowExtension, columnContainerExtension) + propertyName.identifier to listOf(dataRowExtension, columnContainerExtension) } } } @@ -109,6 +109,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se resolvedTypeRef: FirResolvedTypeRef, propertyName: Name, k: FirClassSymbol<*>, + columnNameAnnotation: FirAnnotation? = null, isScopeProperty: Boolean = false, order: Int? = null, ): FirProperty { @@ -135,6 +136,9 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se argumentMapping = buildAnnotationArgumentMapping() } } + columnNameAnnotation?.let { + annotations += it + } replaceAnnotations(annotations) } } 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 7c634dedfb..ff2435c971 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 @@ -29,9 +29,9 @@ import org.jetbrains.kotlin.fir.types.toSymbol import org.jetbrains.kotlin.fir.types.toTypeProjection import org.jetbrains.kotlin.name.CallableId import org.jetbrains.kotlin.name.FqName -import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.types.Variance import org.jetbrains.kotlinx.dataframe.annotations.DataSchema +import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName /** * extensions inside scope classes are generated here: @@ -87,7 +87,7 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio val columnName = property.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, session)?.let { annotation -> (annotation.argumentMapping.mapping[Names.COLUMN_NAME_ARGUMENT] as? FirLiteralExpression)?.value as? String? } - val propertyName = property.name + val name = property.name val marker = owner.constructType(arrayOf(), isNullable = false).toTypeProjection(Variance.INVARIANT) val columnGroupProjection: ConeTypeProjection? = if (resolvedReturnTypeRef.coneType.classId?.equals( @@ -119,8 +119,7 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio typeArguments = arrayOf(marker), isNullable = false ), - propertyName = propertyName, - columnName = columnName, + propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }), returnTypeRef = resolvedReturnTypeRef ) @@ -142,8 +141,7 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio typeArguments = arrayOf(marker), isNullable = false ), - propertyName = propertyName, - columnName = columnName, + propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }), returnTypeRef = columnReturnType ) listOf(rowExtension.symbol, columnsContainerExtension.symbol) diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/PropertyName.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/PropertyName.kt new file mode 100644 index 0000000000..8a64de99d1 --- /dev/null +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/PropertyName.kt @@ -0,0 +1,52 @@ +package org.jetbrains.kotlinx.dataframe.plugin.extensions.impl + +import org.jetbrains.kotlin.fir.expressions.FirAnnotation +import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping +import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression +import org.jetbrains.kotlin.fir.resolve.defaultType +import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef +import org.jetbrains.kotlin.name.Name +import org.jetbrains.kotlin.types.ConstantValueKind +import org.jetbrains.kotlinx.dataframe.codeGen.ValidFieldName +import org.jetbrains.kotlinx.dataframe.plugin.utils.Names + +data class PropertyName(val identifier: Name, val columnNameAnnotation: FirAnnotation?) { + companion object { + fun of(name: String): PropertyName { + val valid = ValidFieldName.of(name) + var columnName = false + val identifier = if (valid.unquoted != name) { + columnName = true + Name.identifier(valid.unquoted) + } else { + Name.identifier(name) + } + val columnNameAnnotation: FirAnnotation? = if (columnName) { + buildAnnotation(name) + } else { + null + } + return PropertyName(identifier, columnNameAnnotation) + } + + fun buildAnnotation(name: String): FirAnnotation { + return org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation { + annotationTypeRef = buildResolvedTypeRef { + type = Names.COLUMN_NAME_ANNOTATION.defaultType(emptyList()) + } + argumentMapping = buildAnnotationArgumentMapping { + mapping[Names.COLUMN_NAME_ARGUMENT] = buildLiteralExpression( + source = null, + kind = ConstantValueKind.String, + value = name, + setType = true + ) + } + } + } + + fun of(identifier: Name, columnNameAnnotation: FirAnnotation?): PropertyName { + return PropertyName(identifier, columnNameAnnotation) + } + } +} diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/SchemaProperty.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/SchemaProperty.kt new file mode 100644 index 0000000000..07198b39a6 --- /dev/null +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/impl/SchemaProperty.kt @@ -0,0 +1,12 @@ +package org.jetbrains.kotlinx.dataframe.plugin.extensions.impl + +import org.jetbrains.kotlin.fir.types.ConeKotlinType +import org.jetbrains.kotlin.fir.types.ConeTypeProjection + +data class SchemaProperty( + val marker: ConeTypeProjection, + val propertyName: PropertyName, + val dataRowReturnType: ConeKotlinType, + val columnContainerReturnType: ConeKotlinType, + val override: Boolean = false +) 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 32c1291550..e6fa15743b 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 @@ -21,23 +21,25 @@ import org.jetbrains.kotlin.fir.toFirResolvedTypeRef import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl import org.jetbrains.kotlin.name.CallableId -import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlinx.dataframe.plugin.extensions.DataFramePlugin +import org.jetbrains.kotlinx.dataframe.plugin.extensions.impl.PropertyName internal fun FirDeclarationGenerationExtension.generateExtensionProperty( callableId: CallableId, receiverType: ConeClassLikeTypeImpl, - propertyName: Name, + propertyName: PropertyName, returnTypeRef: FirResolvedTypeRef, - columnName: String? = null, symbol: FirClassSymbol<*>? = null, effectiveVisibility: EffectiveVisibility = EffectiveVisibility.Public ): FirProperty { val firPropertySymbol = FirPropertySymbol(callableId) return buildProperty { + propertyName.columnNameAnnotation?.let { + annotations += it + } moduleData = session.moduleData resolvePhase = FirResolvePhase.BODY_RESOLVE - origin = FirDeclarationOrigin.Plugin(DataFramePlugin(columnName)) + origin = FirDeclarationOrigin.Plugin(DataFramePlugin) status = FirResolvedDeclarationStatusImpl( Visibilities.Public, Modality.FINAL, @@ -67,7 +69,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty( getter = buildPropertyAccessor { moduleData = session.moduleData resolvePhase = FirResolvePhase.BODY_RESOLVE - origin = FirDeclarationOrigin.Plugin(DataFramePlugin(columnName)) + origin = FirDeclarationOrigin.Plugin(DataFramePlugin) this.returnTypeRef = returnTypeRef dispatchReceiverType = receiverType this.symbol = firPropertyAccessorSymbol @@ -79,7 +81,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty( effectiveVisibility ) } - name = propertyName + name = propertyName.identifier this.symbol = firPropertySymbol isVar = false isLocal = false diff --git a/plugins/kotlin-dataframe/testData/box/columnName_invalidSymbol.kt b/plugins/kotlin-dataframe/testData/box/columnName_invalidSymbol.kt new file mode 100644 index 0000000000..4301b94899 --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/columnName_invalidSymbol.kt @@ -0,0 +1,10 @@ +import org.jetbrains.kotlinx.dataframe.* +import org.jetbrains.kotlinx.dataframe.annotations.* +import org.jetbrains.kotlinx.dataframe.api.* +import org.jetbrains.kotlinx.dataframe.io.* + +fun box(): String { + val df = dataFrameOf("a.b")(1) + df.`a b` + return "OK" +} diff --git a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java index 0f12c6ec7a..1a0a06b08a 100644 --- a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java +++ b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java @@ -52,6 +52,12 @@ public void testColumnName() { runTest("testData/box/columnName.kt"); } + @Test + @TestMetadata("columnName_invalidSymbol.kt") + public void testColumnName_invalidSymbol() { + runTest("testData/box/columnName_invalidSymbol.kt"); + } + @Test @TestMetadata("columnWithStarProjection.kt") public void testColumnWithStarProjection() {