Skip to content

Commit

Permalink
Merge pull request #913 from Kotlin/bugfixes-1
Browse files Browse the repository at this point in the history
[Compiler plugin] Generate ColumnName annotations on frontend for all names that contain illegal characters
  • Loading branch information
koperagen authored Oct 11, 2024
2 parents 61865d8 + bd33f62 commit b3b1f64
Show file tree
Hide file tree
Showing 14 changed files with 137 additions and 39 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
)
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 -> {
Expand All @@ -534,7 +535,7 @@ class FunctionCallTransformer(

SchemaProperty(
marker = schema.defaultType(),
name = it.name,
propertyName = PropertyName.of(it.name),
dataRowReturnType = frameColumnReturnType,
columnContainerReturnType = frameColumnReturnType.toFirResolvedTypeRef()
.projectOverDataColumnType()
Expand All @@ -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()
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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))
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
}
}
}
Expand All @@ -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 {
Expand All @@ -135,6 +136,9 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
argumentMapping = buildAnnotationArgumentMapping()
}
}
columnNameAnnotation?.let {
annotations += it
}
replaceAnnotations(annotations)
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,15 @@ 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

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!!)
}
Expand Down Expand Up @@ -83,7 +87,7 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
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(
Expand Down Expand Up @@ -115,8 +119,7 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
typeArguments = arrayOf(marker),
isNullable = false
),
propertyName = propertyName,
columnName = columnName,
propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
returnTypeRef = resolvedReturnTypeRef
)

Expand All @@ -138,8 +141,7 @@ class ExtensionsGenerator(session: FirSession) : FirDeclarationGenerationExtensi
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)
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
}
}
}
Original file line number Diff line number Diff line change
@@ -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
)
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -79,7 +81,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
effectiveVisibility
)
}
name = propertyName
name = propertyName.identifier
this.symbol = firPropertySymbol
isVar = false
isLocal = false
Expand Down
10 changes: 10 additions & 0 deletions plugins/kotlin-dataframe/testData/box/columnName_invalidSymbol.kt
Original file line number Diff line number Diff line change
@@ -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"
}
9 changes: 9 additions & 0 deletions plugins/kotlin-dataframe/testData/box/playground.kt
Original file line number Diff line number Diff line change
@@ -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"
}
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down Expand Up @@ -304,6 +310,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() {
Expand Down

0 comments on commit b3b1f64

Please sign in to comment.