Skip to content

Commit 78795e5

Browse files
koperagenkoperagen
authored
Merge pull request #755 from Kotlin/col-kinds
Compiler plugin update
2 parents 33a9992 + da0f881 commit 78795e5

File tree

171 files changed

+330
-37242
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

171 files changed

+330
-37242
lines changed

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,18 @@ public annotation class Import
2828
@Target(AnnotationTarget.PROPERTY)
2929
public annotation class Order(val order: Int)
3030

31+
/**
32+
* For internal use
33+
* Compiler plugin materializes schemas as classes.
34+
* These classes have two kinds of properties:
35+
* 1. Scope properties that only serve as a reference for internal property resolution
36+
* 2. Schema properties that reflect dataframe structure
37+
* Scope properties need
38+
* to be excluded in IDE plugin and in [org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor.get]
39+
* This annotation serves to distinguish between the two where needed
40+
*/
41+
@Target(AnnotationTarget.PROPERTY)
42+
public annotation class ScopeProperty
43+
3144
@Target(AnnotationTarget.FUNCTION)
3245
internal annotation class Check

‎core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
44
import org.jetbrains.kotlinx.dataframe.DataRow
55
import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
66
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
7+
import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
78
import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertyOrderFromPrimaryConstructor
89
import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
910
import kotlin.reflect.KClass
@@ -54,7 +55,8 @@ internal object MarkersExtractor {
5455

5556
private fun getFields(markerClass: KClass<*>, nullableProperties: Boolean): List<GeneratedField> {
5657
val order = getPropertyOrderFromPrimaryConstructor(markerClass) ?: emptyMap()
57-
return markerClass.memberProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
58+
val structuralProperties = markerClass.memberProperties.filter { it.findAnnotation<ScopeProperty>() == null }
59+
return structuralProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
5860
val fieldName = ValidFieldName.of(it.name)
5961
val columnName = it.findAnnotation<ColumnName>()?.name ?: fieldName.unquoted
6062
val type = it.returnType

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,18 @@ public annotation class Import
2828
@Target(AnnotationTarget.PROPERTY)
2929
public annotation class Order(val order: Int)
3030

31+
/**
32+
* For internal use
33+
* Compiler plugin materializes schemas as classes.
34+
* These classes have two kinds of properties:
35+
* 1. Scope properties that only serve as a reference for internal property resolution
36+
* 2. Schema properties that reflect dataframe structure
37+
* Scope properties need
38+
* to be excluded in IDE plugin and in [org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor.get]
39+
* This annotation serves to distinguish between the two where needed
40+
*/
41+
@Target(AnnotationTarget.PROPERTY)
42+
public annotation class ScopeProperty
43+
3144
@Target(AnnotationTarget.FUNCTION)
3245
internal annotation class Check

‎core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
44
import org.jetbrains.kotlinx.dataframe.DataRow
55
import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
66
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
7+
import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
78
import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertyOrderFromPrimaryConstructor
89
import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
910
import kotlin.reflect.KClass
@@ -54,7 +55,8 @@ internal object MarkersExtractor {
5455

5556
private fun getFields(markerClass: KClass<*>, nullableProperties: Boolean): List<GeneratedField> {
5657
val order = getPropertyOrderFromPrimaryConstructor(markerClass) ?: emptyMap()
57-
return markerClass.memberProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
58+
val structuralProperties = markerClass.memberProperties.filter { it.hasAnnotation<ScopeProperty>() }
59+
return structuralProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
5860
val fieldName = ValidFieldName.of(it.name)
5961
val columnName = it.findAnnotation<ColumnName>()?.name ?: fieldName.unquoted
6062
val type = it.returnType

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions
33
import org.jetbrains.kotlin.fir.FirSession
44
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
55
import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
6+
import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
67
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
78
import org.jetbrains.kotlin.fir.extensions.FirExpressionResolutionExtension
89
import org.jetbrains.kotlin.fir.scopes.collectAllProperties
@@ -21,7 +22,7 @@ class ReturnTypeBasedReceiverInjector(session: FirSession) : FirExpressionResolu
2122
val symbol = generatedTokenOrNull(functionCall) ?: return emptyList()
2223
return symbol.declaredMemberScope(session, FirResolvePhase.DECLARATIONS).collectAllProperties()
2324
.filterIsInstance<FirPropertySymbol>()
24-
.filter { it.resolvedReturnType.classId?.shortClassName?.asString()?.startsWith("Scope") ?: false }
25+
.filter { it.getAnnotationByClassId(Names.SCOPE_PROPERTY_ANNOTATION, session) != null }
2526
.map { it.resolvedReturnType }
2627
}
2728

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

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.caches.getValue
1212
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
1313
import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
1414
import org.jetbrains.kotlin.fir.declarations.FirProperty
15+
import org.jetbrains.kotlin.fir.expressions.FirAnnotation
1516
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
1617
import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
1718
import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
@@ -51,8 +52,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
5152
}
5253
is CallShapeData.RefinedType -> callShapeData.scopes.associate {
5354
val propertyName = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
54-
// making them var appeared to be the easiest way to filter
55-
propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isVal = false))
55+
propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isScopeProperty = true))
5656
}
5757
is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty ->
5858
val propertyName = Name.identifier(schemaProperty.name)
@@ -89,7 +89,6 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
8989

9090
@OptIn(SymbolInternals::class)
9191
override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
92-
// maybe Init needed not for everything
9392
val destination = mutableSetOf<Name>()
9493
when (classSymbol.fir.callShapeData) {
9594
is CallShapeData.RefinedType -> destination.add(SpecialNames.INIT)
@@ -110,28 +109,33 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
110109
resolvedTypeRef: FirResolvedTypeRef,
111110
propertyName: Name,
112111
k: FirClassSymbol<*>,
113-
isVal: Boolean = true,
112+
isScopeProperty: Boolean = false,
114113
order: Int? = null,
115114
): FirProperty {
116-
return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type, isVal) {
115+
return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type) {
117116
modality = Modality.ABSTRACT
118117
visibility = Visibilities.Public
119118
}.apply {
119+
val annotations = mutableListOf<FirAnnotation>()
120120
if (order != null) {
121-
val orderAnnotation = buildAnnotation {
121+
annotations += buildAnnotation {
122122
annotationTypeRef = buildResolvedTypeRef {
123-
type = ConeClassLikeTypeImpl(
124-
ConeClassLikeLookupTagImpl(Names.ORDER_ANNOTATION),
125-
arrayOf(),
126-
isNullable = false
127-
)
123+
type = Names.ORDER_ANNOTATION.defaultType(emptyList())
128124
}
129125
argumentMapping = buildAnnotationArgumentMapping {
130126
mapping[Names.ORDER_ARGUMENT] = buildLiteralExpression(null, ConstantValueKind.Int, order, setType = true)
131127
}
132128
}
133-
replaceAnnotations(listOf(orderAnnotation))
134129
}
130+
if (isScopeProperty) {
131+
annotations += buildAnnotation {
132+
annotationTypeRef = buildResolvedTypeRef {
133+
type = Names.SCOPE_PROPERTY_ANNOTATION.defaultType(emptyList())
134+
}
135+
argumentMapping = buildAnnotationArgumentMapping()
136+
}
137+
}
138+
replaceAnnotations(annotations)
135139
}
136140
}
137141

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22

33
package org.jetbrains.kotlinx.dataframe.plugin.impl
44

5+
import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
56
import org.jetbrains.kotlin.fir.types.ConeKotlinType
67
import org.jetbrains.kotlin.fir.types.ConeNullability
7-
import org.jetbrains.kotlin.fir.types.classId
88
import org.jetbrains.kotlin.fir.types.isNullable
99
import org.jetbrains.kotlinx.dataframe.impl.api.DataFrameLikeContainer
1010
import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumn
@@ -101,17 +101,17 @@ data class SimpleColumnGroup(
101101
}
102102

103103
fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleCol {
104-
return if (type.classId == Names.DATA_ROW_CLASS_ID) {
105-
val schema = pluginDataFrameSchema(type)
104+
return if (type.fullyExpandedClassId(session) == Names.DATA_ROW_CLASS_ID) {
105+
val schema = pluginDataFrameSchema(type.typeArguments[0])
106106
val group = SimpleColumnGroup(name, schema.columns())
107107
val column = if (type.isNullable) {
108108
makeNullable(group)
109109
} else {
110110
group
111111
}
112112
column
113-
} else if (type.classId == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
114-
val schema = pluginDataFrameSchema(type)
113+
} else if (type.fullyExpandedClassId(session) == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
114+
val schema = pluginDataFrameSchema(type.typeArguments[0])
115115
SimpleFrameColumn(name, schema.columns())
116116
} else {
117117
SimpleDataColumn(name, type.wrap())
@@ -121,7 +121,7 @@ fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleC
121121
private fun KotlinTypeFacade.makeNullable(column: SimpleCol): SimpleCol {
122122
return when (column) {
123123
is SimpleColumnGroup -> {
124-
SimpleColumnGroup(column.name, column.columns().map { makeNullable(column) })
124+
SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
125125
}
126126
is SimpleFrameColumn -> column
127127
is SimpleDataColumn -> SimpleDataColumn(column.name, column.type.changeNullability { true })

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/add.kt

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
77
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
88
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
99
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
10-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1110
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
11+
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
1212
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
1313
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
1414

@@ -20,7 +20,7 @@ class Add : AbstractSchemaModificationInterpreter() {
2020
val Arguments.type: TypeApproximation by type(name("expression"))
2121

2222
override fun Arguments.interpret(): PluginDataFrameSchema {
23-
return PluginDataFrameSchema(receiver.columns() + SimpleDataColumn(name, type))
23+
return PluginDataFrameSchema(receiver.columns() + simpleColumnOf(name, type.type))
2424
}
2525
}
2626

@@ -30,7 +30,7 @@ class From : AbstractInterpreter<Unit>() {
3030
val Arguments.type: TypeApproximation by type(name("expression"))
3131

3232
override fun Arguments.interpret() {
33-
dsl.columns += SimpleDataColumn(receiver, type)
33+
dsl.columns += simpleColumnOf(receiver, type.type)
3434
}
3535
}
3636

@@ -40,7 +40,7 @@ class Into : AbstractInterpreter<Unit>() {
4040
val Arguments.name: String by string()
4141

4242
override fun Arguments.interpret() {
43-
dsl.columns += SimpleDataColumn(name, receiver)
43+
dsl.columns += simpleColumnOf(name, receiver.type)
4444
}
4545
}
4646

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/convert.kt

Lines changed: 5 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,20 @@
11
package org.jetbrains.kotlinx.dataframe.plugin.impl.api
22

3-
import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
4-
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
5-
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
6-
import org.jetbrains.kotlin.fir.types.classId
7-
import org.jetbrains.kotlin.fir.types.isNullable
3+
import org.jetbrains.kotlinx.dataframe.api.Infer
84
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
95
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
106
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
117
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
12-
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
13-
import org.jetbrains.kotlinx.dataframe.api.Infer
148
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
9+
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
1510
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
16-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1711
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
12+
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
1813
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
1914
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
2015
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
2116
import org.jetbrains.kotlinx.dataframe.plugin.impl.enum
17+
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
2218
import org.jetbrains.kotlinx.dataframe.plugin.impl.string
2319
import org.jetbrains.kotlinx.dataframe.plugin.impl.type
2420
import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString
@@ -73,17 +69,8 @@ internal fun KotlinTypeFacade.convertImpl(
7369
type: TypeApproximation
7470
): PluginDataFrameSchema {
7571
return pluginDataFrameSchema.map(columns.toSet()) { path, column ->
76-
require(column is SimpleDataColumn) {
77-
"$path must be ${SimpleDataColumn::class}, but was ${column::class}"
78-
}
7972
val unwrappedType = type.type
80-
// TODO: AnyRow
81-
if (unwrappedType is ConeClassLikeType && unwrappedType.classId == Names.DF_CLASS_ID && !unwrappedType.isNullable) {
82-
val f = unwrappedType.typeArguments.single()
83-
SimpleFrameColumn(column.name, pluginDataFrameSchema(f).columns())
84-
} else {
85-
column.changeType(type)
86-
}
73+
simpleColumnOf(column.name, unwrappedType)
8774
}
8875
}
8976

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt

Lines changed: 2 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,11 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
2121
import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
2222
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
2323
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
24-
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
2524
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
2625
import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
2726
import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
2827
import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
28+
import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
2929

3030
class GroupBy(val df: PluginDataFrameSchema, val keys: List<ColumnWithPathApproximation>, val moveToTop: Boolean)
3131

@@ -82,27 +82,8 @@ fun KotlinTypeFacade.aggregate(
8282
)
8383
}
8484

85-
// important to create FrameColumns, nullable DataRows?
8685
val cols = createPluginDataFrameSchema(groupBy.keys, groupBy.moveToTop).columns() + dsl.columns.map {
87-
when (it.type.classId) {
88-
Names.DATA_ROW_CLASS_ID -> {
89-
when (it.type.nullability) {
90-
ConeNullability.NULLABLE -> SimpleDataColumn(
91-
it.name,
92-
TypeApproximation(it.type)
93-
)
94-
ConeNullability.UNKNOWN -> SimpleDataColumn(
95-
it.name,
96-
TypeApproximation(it.type)
97-
)
98-
ConeNullability.NOT_NULL -> {
99-
val typeProjection = it.type.typeArguments[0]
100-
SimpleColumnGroup(it.name, pluginDataFrameSchema(typeProjection).columns())
101-
}
102-
}
103-
}
104-
else -> SimpleDataColumn(it.name, TypeApproximation(it.type))
105-
}
86+
simpleColumnOf(it.name, it.type)
10687
}
10788
PluginDataFrameSchema(cols)
10889
} else {

0 commit comments

Comments
 (0)