Merge pull request #755 from Kotlin/col-kinds

Compiler plugin update
Kotlin · Jun 27, 2024 · 78795e5 · 78795e5
2 parents 33a9992 + da0f881
commit 78795e5
Show file tree

Hide file tree

Showing 171 changed files with 330 additions and 37,242 deletions.
diff --git a/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt b/core/generated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt
@@ -28,5 +28,18 @@ public annotation class Import
 @Target(AnnotationTarget.PROPERTY)
 public annotation class Order(val order: Int)
 
+/**
+ * For internal use
+ * Compiler plugin materializes schemas as classes.
+ * These classes have two kinds of properties:
+ * 1. Scope properties that only serve as a reference for internal property resolution
+ * 2. Schema properties that reflect dataframe structure
+ * Scope properties need
+ * to be excluded in IDE plugin and in [org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor.get]
+ * This annotation serves to distinguish between the two where needed
+ */
+@Target(AnnotationTarget.PROPERTY)
+public annotation class ScopeProperty
+
 @Target(AnnotationTarget.FUNCTION)
 internal annotation class Check
diff --git a/...rated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt b/...rated-sources/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
 import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertyOrderFromPrimaryConstructor
 import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
 import kotlin.reflect.KClass
@@ -54,7 +55,8 @@ internal object MarkersExtractor {
 
     private fun getFields(markerClass: KClass<*>, nullableProperties: Boolean): List<GeneratedField> {
         val order = getPropertyOrderFromPrimaryConstructor(markerClass) ?: emptyMap()
-        return markerClass.memberProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
+        val structuralProperties = markerClass.memberProperties.filter { it.findAnnotation<ScopeProperty>() == null }
+        return structuralProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
             val fieldName = ValidFieldName.of(it.name)
             val columnName = it.findAnnotation<ColumnName>()?.name ?: fieldName.unquoted
             val type = it.returnType

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/annotations/Plugin.kt
@@ -28,5 +28,18 @@ public annotation class Import
 @Target(AnnotationTarget.PROPERTY)
 public annotation class Order(val order: Int)
 
+/**
+ * For internal use
+ * Compiler plugin materializes schemas as classes.
+ * These classes have two kinds of properties:
+ * 1. Scope properties that only serve as a reference for internal property resolution
+ * 2. Schema properties that reflect dataframe structure
+ * Scope properties need
+ * to be excluded in IDE plugin and in [org.jetbrains.kotlinx.dataframe.codeGen.MarkersExtractor.get]
+ * This annotation serves to distinguish between the two where needed
+ */
+@Target(AnnotationTarget.PROPERTY)
+public annotation class ScopeProperty
+
 @Target(AnnotationTarget.FUNCTION)
 internal annotation class Check
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/codeGen/MarkersExtractor.kt
@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.dataframe.DataFrame
 import org.jetbrains.kotlinx.dataframe.DataRow
 import org.jetbrains.kotlinx.dataframe.annotations.ColumnName
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
+import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
 import org.jetbrains.kotlinx.dataframe.impl.schema.getPropertyOrderFromPrimaryConstructor
 import org.jetbrains.kotlinx.dataframe.schema.ColumnSchema
 import kotlin.reflect.KClass
@@ -54,7 +55,8 @@ internal object MarkersExtractor {
 
     private fun getFields(markerClass: KClass<*>, nullableProperties: Boolean): List<GeneratedField> {
         val order = getPropertyOrderFromPrimaryConstructor(markerClass) ?: emptyMap()
-        return markerClass.memberProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
+        val structuralProperties = markerClass.memberProperties.filter { it.hasAnnotation<ScopeProperty>() }
+        return structuralProperties.sortedBy { order[it.name] ?: Int.MAX_VALUE }.mapIndexed { _, it ->
             val fieldName = ValidFieldName.of(it.name)
             val columnName = it.findAnnotation<ColumnName>()?.name ?: fieldName.unquoted
             val type = it.returnType

diff --git a/.../src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ReturnTypeBasedReceiverInjector.kt b/.../src/org/jetbrains/kotlinx/dataframe/plugin/extensions/ReturnTypeBasedReceiverInjector.kt
@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.extensions
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.getAnnotationByClassId
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.extensions.FirExpressionResolutionExtension
 import org.jetbrains.kotlin.fir.scopes.collectAllProperties
@@ -21,7 +22,7 @@ class ReturnTypeBasedReceiverInjector(session: FirSession) : FirExpressionResolu
         val symbol = generatedTokenOrNull(functionCall) ?: return emptyList()
         return symbol.declaredMemberScope(session, FirResolvePhase.DECLARATIONS).collectAllProperties()
             .filterIsInstance<FirPropertySymbol>()
-            .filter { it.resolvedReturnType.classId?.shortClassName?.asString()?.startsWith("Scope") ?: false }
+            .filter { it.getAnnotationByClassId(Names.SCOPE_PROPERTY_ANNOTATION, session) != null }
             .map { it.resolvedReturnType }
     }
 

diff --git a/.../kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt b/.../kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/TokenGenerator.kt
@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 import org.jetbrains.kotlinx.dataframe.plugin.utils.generateExtensionProperty
 import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
 import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotationArgumentMapping
 import org.jetbrains.kotlin.fir.expressions.builder.buildLiteralExpression
@@ -51,8 +52,7 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
                 }
                 is CallShapeData.RefinedType -> callShapeData.scopes.associate {
                     val propertyName = Name.identifier(it.name.identifier.replaceFirstChar { it.lowercaseChar() })
-                    // making them var appeared to be the easiest way to filter
-                    propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isVal = false))
+                    propertyName to listOf(buildProperty(it.defaultType().toFirResolvedTypeRef(), propertyName, k, isScopeProperty = true))
                 }
                 is CallShapeData.Scope -> callShapeData.columns.associate { schemaProperty ->
                     val propertyName = Name.identifier(schemaProperty.name)
@@ -89,7 +89,6 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
 
     @OptIn(SymbolInternals::class)
     override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
-        // maybe Init needed not for everything
         val destination = mutableSetOf<Name>()
         when (classSymbol.fir.callShapeData) {
             is CallShapeData.RefinedType -> destination.add(SpecialNames.INIT)
@@ -110,28 +109,33 @@ class TokenGenerator(session: FirSession) : FirDeclarationGenerationExtension(se
         resolvedTypeRef: FirResolvedTypeRef,
         propertyName: Name,
         k: FirClassSymbol<*>,
-        isVal: Boolean = true,
+        isScopeProperty: Boolean = false,
         order: Int? = null,
     ): FirProperty {
-        return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type, isVal) {
+        return createMemberProperty(k, Key, propertyName, resolvedTypeRef.type) {
             modality = Modality.ABSTRACT
             visibility = Visibilities.Public
         }.apply {
+            val annotations = mutableListOf<FirAnnotation>()
             if (order != null) {
-                val orderAnnotation = buildAnnotation {
+                annotations += buildAnnotation {
                     annotationTypeRef = buildResolvedTypeRef {
-                        type = ConeClassLikeTypeImpl(
-                            ConeClassLikeLookupTagImpl(Names.ORDER_ANNOTATION),
-                            arrayOf(),
-                            isNullable = false
-                        )
+                        type = Names.ORDER_ANNOTATION.defaultType(emptyList())
                     }
                     argumentMapping = buildAnnotationArgumentMapping {
                         mapping[Names.ORDER_ARGUMENT] = buildLiteralExpression(null, ConstantValueKind.Int, order, setType = true)
                     }
                 }
-                replaceAnnotations(listOf(orderAnnotation))
             }
+            if (isScopeProperty) {
+                annotations += buildAnnotation {
+                    annotationTypeRef = buildResolvedTypeRef {
+                        type = Names.SCOPE_PROPERTY_ANNOTATION.defaultType(emptyList())
+                    }
+                    argumentMapping = buildAnnotationArgumentMapping()
+                }
+            }
+            replaceAnnotations(annotations)
         }
     }
 

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/SimpleCol.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/SimpleCol.kt
@@ -2,9 +2,9 @@
 
 package org.jetbrains.kotlinx.dataframe.plugin.impl
 
+import org.jetbrains.kotlin.fir.analysis.checkers.fullyExpandedClassId
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeNullability
-import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.isNullable
 import org.jetbrains.kotlinx.dataframe.impl.api.DataFrameLikeContainer
 import org.jetbrains.kotlinx.dataframe.impl.api.GenericColumn
@@ -101,17 +101,17 @@ data class SimpleColumnGroup(
 }
 
 fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleCol {
-    return if (type.classId == Names.DATA_ROW_CLASS_ID) {
-        val schema = pluginDataFrameSchema(type)
+    return if (type.fullyExpandedClassId(session) == Names.DATA_ROW_CLASS_ID) {
+        val schema = pluginDataFrameSchema(type.typeArguments[0])
         val group = SimpleColumnGroup(name, schema.columns())
         val column = if (type.isNullable) {
             makeNullable(group)
         } else {
             group
         }
         column
-    } else if (type.classId == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
-        val schema = pluginDataFrameSchema(type)
+    } else if (type.fullyExpandedClassId(session) == Names.DF_CLASS_ID && type.nullability == ConeNullability.NOT_NULL) {
+        val schema = pluginDataFrameSchema(type.typeArguments[0])
         SimpleFrameColumn(name, schema.columns())
     } else {
         SimpleDataColumn(name, type.wrap())
@@ -121,7 +121,7 @@ fun KotlinTypeFacade.simpleColumnOf(name: String, type: ConeKotlinType): SimpleC
 private fun KotlinTypeFacade.makeNullable(column: SimpleCol): SimpleCol {
     return when (column) {
         is SimpleColumnGroup -> {
-            SimpleColumnGroup(column.name, column.columns().map { makeNullable(column) })
+            SimpleColumnGroup(column.name, column.columns().map { makeNullable(it) })
         }
         is SimpleFrameColumn -> column
         is SimpleDataColumn -> SimpleDataColumn(column.name, column.type.changeNullability { true })

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/add.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/add.kt
@@ -7,8 +7,8 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
-import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
 import org.jetbrains.kotlinx.dataframe.plugin.impl.string
 import org.jetbrains.kotlinx.dataframe.plugin.impl.type
 
@@ -20,7 +20,7 @@ class Add : AbstractSchemaModificationInterpreter() {
     val Arguments.type: TypeApproximation by type(name("expression"))
 
     override fun Arguments.interpret(): PluginDataFrameSchema {
-        return PluginDataFrameSchema(receiver.columns() + SimpleDataColumn(name, type))
+        return PluginDataFrameSchema(receiver.columns() + simpleColumnOf(name, type.type))
     }
 }
 
@@ -30,7 +30,7 @@ class From : AbstractInterpreter<Unit>() {
     val Arguments.type: TypeApproximation by type(name("expression"))
 
     override fun Arguments.interpret() {
-        dsl.columns += SimpleDataColumn(receiver, type)
+        dsl.columns += simpleColumnOf(receiver, type.type)
     }
 }
 
@@ -40,7 +40,7 @@ class Into : AbstractInterpreter<Unit>() {
     val Arguments.name: String by string()
 
     override fun Arguments.interpret() {
-        dsl.columns += SimpleDataColumn(name, receiver)
+        dsl.columns += simpleColumnOf(name, receiver.type)
     }
 }
 

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/convert.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/convert.kt
@@ -1,24 +1,20 @@
 package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 
-import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
-import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType
-import org.jetbrains.kotlin.fir.types.classId
-import org.jetbrains.kotlin.fir.types.isNullable
+import org.jetbrains.kotlinx.dataframe.api.Infer
 import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
-import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
-import org.jetbrains.kotlinx.dataframe.api.Infer
 import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
-import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
+import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.enum
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
 import org.jetbrains.kotlinx.dataframe.plugin.impl.string
 import org.jetbrains.kotlinx.dataframe.plugin.impl.type
 import org.jetbrains.kotlinx.dataframe.plugin.impl.varargString
@@ -73,17 +69,8 @@ internal fun KotlinTypeFacade.convertImpl(
     type: TypeApproximation
 ): PluginDataFrameSchema {
     return pluginDataFrameSchema.map(columns.toSet()) { path, column ->
-        require(column is SimpleDataColumn) {
-            "$path must be ${SimpleDataColumn::class}, but was ${column::class}"
-        }
         val unwrappedType = type.type
-        // TODO: AnyRow
-        if (unwrappedType is ConeClassLikeType && unwrappedType.classId == Names.DF_CLASS_ID && !unwrappedType.isNullable) {
-            val f = unwrappedType.typeArguments.single()
-            SimpleFrameColumn(column.name, pluginDataFrameSchema(f).columns())
-        } else {
-            column.changeType(type)
-        }
+        simpleColumnOf(column.name, unwrappedType)
     }
 }
 

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/groupBy.kt
@@ -21,11 +21,11 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
 import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
 import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleCol
-import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleDataColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleColumnGroup
 import org.jetbrains.kotlinx.dataframe.plugin.impl.SimpleFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnWithPathApproximation
 import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.simpleColumnOf
 
 class GroupBy(val df: PluginDataFrameSchema, val keys: List<ColumnWithPathApproximation>, val moveToTop: Boolean)
 
@@ -82,27 +82,8 @@ fun KotlinTypeFacade.aggregate(
             )
         }
 
-        // important to create FrameColumns, nullable DataRows?
         val cols = createPluginDataFrameSchema(groupBy.keys, groupBy.moveToTop).columns() + dsl.columns.map {
-            when (it.type.classId) {
-                Names.DATA_ROW_CLASS_ID -> {
-                    when (it.type.nullability) {
-                        ConeNullability.NULLABLE -> SimpleDataColumn(
-                            it.name,
-                            TypeApproximation(it.type)
-                        )
-                        ConeNullability.UNKNOWN -> SimpleDataColumn(
-                            it.name,
-                            TypeApproximation(it.type)
-                        )
-                        ConeNullability.NOT_NULL -> {
-                            val typeProjection = it.type.typeArguments[0]
-                            SimpleColumnGroup(it.name, pluginDataFrameSchema(typeProjection).columns())
-                        }
-                    }
-                }
-                else -> SimpleDataColumn(it.name, TypeApproximation(it.type))
-            }
+            simpleColumnOf(it.name, it.type)
         }
         PluginDataFrameSchema(cols)
     } else {

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlinx.dataframe.annotations.Order
+import org.jetbrains.kotlinx.dataframe.annotations.ScopeProperty
 import kotlin.reflect.KClass
 
 object Names {
@@ -32,10 +33,12 @@ object Names {
         get() = Name.identifier("org.jetbrains.kotlinx.dataframe.annotations")
     val INTERPRETABLE_FQNAME: FqName
         get() = FqName("org.jetbrains.kotlinx.dataframe.annotations.Interpretable")
-    val ORDER_ANNOTATION = ClassId(FqName("org.jetbrains.kotlinx.dataframe.annotations"), Name.identifier(Order::class.simpleName!!))
+    private val annotationsPackage = FqName("org.jetbrains.kotlinx.dataframe.annotations")
+    val ORDER_ANNOTATION = ClassId(annotationsPackage, Name.identifier(Order::class.simpleName!!))
     val ORDER_ARGUMENT = Name.identifier(Order::order.name)
+    val SCOPE_PROPERTY_ANNOTATION = ClassId(annotationsPackage, Name.identifier(ScopeProperty::class.simpleName!!))
 
-    val DATA_SCHEMA_CLASS_ID = ClassId(FqName("org.jetbrains.kotlinx.dataframe.annotations"), Name.identifier("DataSchema"))
+    val DATA_SCHEMA_CLASS_ID = ClassId(annotationsPackage, Name.identifier("DataSchema"))
     val LIST = ClassId(FqName("kotlin.collections"), Name.identifier("List"))
     val DURATION_CLASS_ID = kotlin.time.Duration::class.classId()
     val LOCAL_DATE_CLASS_ID = kotlinx.datetime.LocalDate::class.classId()