Skip to content

Commit 42ba0b4

Browse files
koperagenkoperagen
committed
[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.
1 parent 883d451 commit 42ba0b4

File tree

10 files changed

+111
-34
lines changed

10 files changed

+111
-34
lines changed

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/SchemaProperty.kt

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package org.jetbrains.kotlinx.dataframe.plugin
2+
3+
import org.jetbrains.kotlin.fir.expressions.FirAnnotation
4+
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
5+
import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
6+
import org.jetbrains.kotlin.fir.resolve.defaultType
7+
import org.jetbrains.kotlin.fir.types.ConeKotlinType
8+
import org.jetbrains.kotlin.fir.types.ConeTypeProjection
9+
import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
10+
import org.jetbrains.kotlin.name.Name
11+
import org.jetbrains.kotlin.types.ConstantValueKind
12+
import org.jetbrains.kotlinx.dataframe.codeGen.ValidFieldName
13+
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
14+
15+
data class SchemaProperty(
16+
val marker: ConeTypeProjection,
17+
val propertyName: PropertyName,
18+
val dataRowReturnType: ConeKotlinType,
19+
val columnContainerReturnType: ConeKotlinType,
20+
val override: Boolean = false
21+
)
22+
23+
data class PropertyName(val identifier: Name, val columnNameAnnotation: FirAnnotation?) {
24+
companion object {
25+
fun of(name: String): PropertyName {
26+
val valid = ValidFieldName.of(name)
27+
var columnName = false
28+
val identifier = if (valid.unquoted != name) {
29+
columnName = true
30+
Name.identifier(valid.unquoted)
31+
} else {
32+
Name.identifier(name)
33+
}
34+
val columnNameAnnotation: FirAnnotation? = if (columnName) {
35+
buildAnnotation(name)
36+
} else {
37+
null
38+
}
39+
return PropertyName(identifier, columnNameAnnotation)
40+
}
41+
42+
fun buildAnnotation(name: String): FirAnnotation {
43+
return org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation {
44+
annotationTypeRef = buildResolvedTypeRef {
45+
type = Names.COLUMN_NAME_ANNOTATION.defaultType(emptyList())
46+
}
47+
argumentMapping = buildAnnotationArgumentMapping {
48+
mapping[Names.COLUMN_NAME_ARGUMENT] = buildLiteralExpression(
49+
source = null,
50+
kind = ConstantValueKind.String,
51+
value = name,
52+
setType = true
53+
)
54+
}
55+
}
56+
}
57+
58+
fun of(identifier: Name, columnNameAnnotation: FirAnnotation?): PropertyName {
59+
return PropertyName(identifier, columnNameAnnotation)
60+
}
61+
}
62+
}

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/analyzeRefinedCallShape.kt

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@ package org.jetbrains.kotlinx.dataframe.plugin
88
import org.jetbrains.kotlin.fir.expressions.FirExpression
99
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
1010
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
11-
import org.jetbrains.kotlin.fir.types.ConeKotlinType
12-
import org.jetbrains.kotlin.fir.types.ConeTypeProjection
1311
import org.jetbrains.kotlin.fir.types.classId
1412
import org.jetbrains.kotlin.fir.types.resolvedType
1513
import org.jetbrains.kotlin.name.ClassId
@@ -63,11 +61,3 @@ data class RefinedArgument(val name: Name, val expression: FirExpression) {
6361
return "RefinedArgument(name=$name, expression=${expression})"
6462
}
6563
}
66-
67-
data class SchemaProperty(
68-
val marker: ConeTypeProjection,
69-
val name: String,
70-
val dataRowReturnType: ConeKotlinType,
71-
val columnContainerReturnType: ConeKotlinType,
72-
val override: Boolean = false
73-
)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions
22

33
import org.jetbrains.kotlin.GeneratedDeclarationKey
44

5-
class DataFramePlugin(val columnName: String?) : GeneratedDeclarationKey()
5+
data object DataFramePlugin : GeneratedDeclarationKey()

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ import org.jetbrains.kotlin.name.FqName
7474
import org.jetbrains.kotlin.name.Name
7575
import org.jetbrains.kotlin.text
7676
import org.jetbrains.kotlin.types.Variance
77+
import org.jetbrains.kotlinx.dataframe.plugin.PropertyName
7778
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
7879
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
7980
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
@@ -520,7 +521,7 @@ class FunctionCallTransformer(
520521
isNullable = false
521522
)
522523

523-
SchemaProperty(schema.defaultType(), it.name, dataRowReturnType, columnsContainerReturnType)
524+
SchemaProperty(schema.defaultType(), PropertyName.of(it.name), dataRowReturnType, columnsContainerReturnType)
524525
}
525526

526527
is SimpleFrameColumn -> {
@@ -534,7 +535,7 @@ class FunctionCallTransformer(
534535

535536
SchemaProperty(
536537
marker = schema.defaultType(),
537-
name = it.name,
538+
propertyName = PropertyName.of(it.name),
538539
dataRowReturnType = frameColumnReturnType,
539540
columnContainerReturnType = frameColumnReturnType.toFirResolvedTypeRef()
540541
.projectOverDataColumnType()
@@ -543,7 +544,7 @@ class FunctionCallTransformer(
543544

544545
is SimpleDataColumn -> SchemaProperty(
545546
marker = schema.defaultType(),
546-
name = it.name,
547+
propertyName = PropertyName.of(it.name),
547548
dataRowReturnType = it.type.type(),
548549
columnContainerReturnType = it.type.type().toFirResolvedTypeRef().projectOverDataColumnType()
549550
)

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ import org.jetbrains.kotlin.ir.types.classOrFail
4848
import org.jetbrains.kotlin.ir.types.classifierOrNull
4949
import org.jetbrains.kotlin.ir.types.getClass
5050
import org.jetbrains.kotlin.ir.util.constructors
51+
import org.jetbrains.kotlin.ir.util.findAnnotation
5152
import org.jetbrains.kotlin.ir.util.parentAsClass
5253
import org.jetbrains.kotlin.ir.util.primaryConstructor
5354
import org.jetbrains.kotlin.ir.util.superTypes
@@ -63,6 +64,7 @@ import org.jetbrains.kotlinx.dataframe.api.schema
6364
import org.jetbrains.kotlinx.dataframe.columns.ColumnGroup
6465
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.IoSchema
6566
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.serialize
67+
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
6668
import java.io.File
6769

6870
class IrBodyFiller(
@@ -209,7 +211,9 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
209211
val call = IrCallImpl(-1, -1, context.irBuiltIns.anyNType, get, 0, 1).also {
210212
val thisSymbol: IrValueSymbol = getter.extensionReceiverParameter?.symbol!!
211213
it.dispatchReceiver = IrGetValueImpl(-1, -1, thisSymbol)
212-
val columName = pluginKey.columnName ?: declaration.name.identifier
214+
val annotation = declaration.annotations.findAnnotation(Names.COLUMN_NAME_ANNOTATION.asSingleFqName())
215+
val columnName = (annotation?.valueArguments?.get(0) as? IrConst<*>)?.value as? String
216+
val columName = columnName ?: declaration.name.identifier
213217
it.putValueArgument(0, IrConstImpl.string(-1, -1, context.irBuiltIns.stringType, columName))
214218
}
215219

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

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -47,16 +47,16 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
4747
val resolvedTypeRef = buildResolvedTypeRef {
4848
type = property.dataRowReturnType
4949
}
50-
val propertyName = Name.identifier(property.name)
51-
propertyName to listOf(buildProperty(resolvedTypeRef, propertyName, k, order = index))
50+
val identifier = property.propertyName.identifier
51+
identifier to listOf(buildProperty(resolvedTypeRef, identifier, k, property.propertyName.columnNameAnnotation, order = index))
5252
}
5353
is CallShapeData.RefinedType -> callShapeData.scopes.associate {
54-
val propertyName = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
55-
propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isScopeProperty = true))
54+
val identifier = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
55+
identifier to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), identifier, k, isScopeProperty = true))
5656
}
5757
is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty ->
58-
val propertyName = Name.identifier(schemaProperty.name)
59-
val callableId = CallableId(k.classId, propertyName)
58+
val propertyName = schemaProperty.propertyName
59+
val callableId = CallableId(k.classId, propertyName.identifier)
6060
val dataRowExtension = generateExtensionProperty(
6161
callableId = callableId,
6262
symbol = k,
@@ -82,7 +82,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
8282
symbol = k,
8383
effectiveVisibility = EffectiveVisibility.Local
8484
)
85-
propertyName to listOf(dataRowExtension, columnContainerExtension)
85+
propertyName.identifier to listOf(dataRowExtension, columnContainerExtension)
8686
}
8787
}
8888
}
@@ -109,6 +109,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
109109
resolvedTypeRef: FirResolvedTypeRef,
110110
propertyName: Name,
111111
k: FirClassSymbol<*>,
112+
columnNameAnnotation: FirAnnotation? = null,
112113
isScopeProperty: Boolean = false,
113114
order: Int? = null,
114115
): FirProperty {
@@ -135,6 +136,9 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
135136
argumentMapping = buildAnnotationArgumentMapping()
136137
}
137138
}
139+
columnNameAnnotation?.let {
140+
annotations += it
141+
}
138142
replaceAnnotations(annotations)
139143
}
140144
}

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,9 @@ import org.jetbrains.kotlin.fir.types.toSymbol
2929
import org.jetbrains.kotlin.fir.types.toTypeProjection
3030
import org.jetbrains.kotlin.name.CallableId
3131
import org.jetbrains.kotlin.name.FqName
32-
import org.jetbrains.kotlin.name.Name
3332
import org.jetbrains.kotlin.types.Variance
3433
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
34+
import org.jetbrains.kotlinx.dataframe.plugin.PropertyName
3535

3636
/**
3737
* extensions inside scope classes are generated here:
@@ -87,7 +87,7 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio
8787
val columnName = property.getAnnotationByClassId(Names.COLUMN_NAME_ANNOTATION, session)?.let { annotation ->
8888
(annotation.argumentMapping.mapping[Names.COLUMN_NAME_ARGUMENT] as? FirLiteralExpression)?.value as? String?
8989
}
90-
val propertyName = property.name
90+
val name = property.name
9191
val marker = owner.constructType(arrayOf(), isNullable = false).toTypeProjection(Variance.INVARIANT)
9292

9393
val columnGroupProjection: ConeTypeProjection? = if (resolvedReturnTypeRef.coneType.classId?.equals(
@@ -119,8 +119,7 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio
119119
typeArguments = arrayOf(marker),
120120
isNullable = false
121121
),
122-
propertyName = propertyName,
123-
columnName = columnName,
122+
propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
124123
returnTypeRef = resolvedReturnTypeRef
125124
)
126125

@@ -142,8 +141,7 @@ class TopLevelExtensionsGenerator(session: FirSession) : FirDeclarationGeneratio
142141
typeArguments = arrayOf(marker),
143142
isNullable = false
144143
),
145-
propertyName = propertyName,
146-
columnName = columnName,
144+
propertyName = PropertyName.of(name, columnName?.let { PropertyName.buildAnnotation(it) }),
147145
returnTypeRef = columnReturnType
148146
)
149147
listOf(rowExtension.symbol, columnsContainerExtension.symbol)

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

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,23 +21,25 @@ import org.jetbrains.kotlin.fir.toFirResolvedTypeRef
2121
import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
2222
import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
2323
import org.jetbrains.kotlin.name.CallableId
24-
import org.jetbrains.kotlin.name.Name
24+
import org.jetbrains.kotlinx.dataframe.plugin.PropertyName
2525
import org.jetbrains.kotlinx.dataframe.plugin.extensions.DataFramePlugin
2626

2727
internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
2828
callableId: CallableId,
2929
receiverType: ConeClassLikeTypeImpl,
30-
propertyName: Name,
30+
propertyName: PropertyName,
3131
returnTypeRef: FirResolvedTypeRef,
32-
columnName: String? = null,
3332
symbol: FirClassSymbol<*>? = null,
3433
effectiveVisibility: EffectiveVisibility = EffectiveVisibility.Public
3534
): FirProperty {
3635
val firPropertySymbol = FirPropertySymbol(callableId)
3736
return buildProperty {
37+
propertyName.columnNameAnnotation?.let {
38+
annotations += it
39+
}
3840
moduleData = session.moduleData
3941
resolvePhase = FirResolvePhase.BODY_RESOLVE
40-
origin = FirDeclarationOrigin.Plugin(DataFramePlugin(columnName))
42+
origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
4143
status = FirResolvedDeclarationStatusImpl(
4244
Visibilities.Public,
4345
Modality.FINAL,
@@ -67,7 +69,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
6769
getter = buildPropertyAccessor {
6870
moduleData = session.moduleData
6971
resolvePhase = FirResolvePhase.BODY_RESOLVE
70-
origin = FirDeclarationOrigin.Plugin(DataFramePlugin(columnName))
72+
origin = FirDeclarationOrigin.Plugin(DataFramePlugin)
7173
this.returnTypeRef = returnTypeRef
7274
dispatchReceiverType = receiverType
7375
this.symbol = firPropertyAccessorSymbol
@@ -79,7 +81,7 @@ internal fun FirDeclarationGenerationExtension.generateExtensionProperty(
7981
effectiveVisibility
8082
)
8183
}
82-
name = propertyName
84+
name = propertyName.identifier
8385
this.symbol = firPropertySymbol
8486
isVar = false
8587
isLocal = false

‎plugins/kotlin-dataframe/testData/box/columnName_invalidSymbol.kt

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
fun box(): String {
7+
val df = dataFrameOf("a.b")(1)
8+
df.`a b`
9+
return "OK"
10+
}

‎plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,12 @@ public void testColumnName() {
5252
runTest("testData/box/columnName.kt");
5353
}
5454

55+
@Test
56+
@TestMetadata("columnName_invalidSymbol.kt")
57+
public void testColumnName_invalidSymbol() {
58+
runTest("testData/box/columnName_invalidSymbol.kt");
59+
}
60+
5561
@Test
5662
@TestMetadata("columnWithStarProjection.kt")
5763
public void testColumnWithStarProjection() {

0 commit comments

Comments
 (0)