[Compiler plugin] support "replace unfold with DSL"

Kotlin · Jun 18, 2024 · 382d140 · 382d140
1 parent 18beb71
commit 382d140
Show file tree

Hide file tree

Showing 11 changed files with 1,009 additions and 8 deletions.
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/replace.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/replace.kt
@@ -6,6 +6,7 @@ import org.jetbrains.kotlinx.dataframe.ColumnsContainer
 import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
 import org.jetbrains.kotlinx.dataframe.columns.ColumnReference
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.get
@@ -14,6 +15,7 @@ import org.jetbrains.kotlinx.dataframe.impl.api.insertImpl
 import org.jetbrains.kotlinx.dataframe.impl.api.removeImpl
 import kotlin.reflect.KProperty
 
+@Interpretable("Replace0")
 public fun <T, C> DataFrame<T>.replace(columns: ColumnsSelector<T, C>): ReplaceClause<T, C> =
     ReplaceClause(this, columns)
 

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/unfold.kt
@@ -5,6 +5,8 @@ import org.jetbrains.kotlinx.dataframe.AnyColumnReference
 import org.jetbrains.kotlinx.dataframe.ColumnsSelector
 import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
+import org.jetbrains.kotlinx.dataframe.annotations.Refine
 import org.jetbrains.kotlinx.dataframe.columns.toColumnSet
 import org.jetbrains.kotlinx.dataframe.impl.api.unfoldImpl
 import kotlin.reflect.KProperty
@@ -18,6 +20,8 @@ public inline fun <reified T> DataColumn<T>.unfold(noinline body: CreateDataFram
 public inline fun <T, reified C> ReplaceClause<T, C>.unfold(vararg props: KProperty<*>, maxDepth: Int = 0): DataFrame<T> =
     with { it.unfold(props = props, maxDepth) }
 
+@Refine
+@Interpretable("ReplaceUnfold1")
 public inline fun <T, reified C> ReplaceClause<T, C>.unfold(noinline body: CreateDataFrameDsl<C>.() -> Unit): DataFrame<T> =
     with { it.unfoldImpl(skipPrimitive = false, body) }
 

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/rename.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/rename.kt
@@ -29,11 +29,11 @@ class RenameInto : AbstractSchemaModificationInterpreter() {
     override fun Arguments.interpret(): PluginDataFrameSchema {
         require(receiver.columns.size == newNames.size)
         var i = 0
-        return receiver.schema.map(receiver.columns.mapTo(mutableSetOf()) { it.path.path }, nextName = { newNames[i].also { i += 1 } })
+        return receiver.schema.rename(receiver.columns.mapTo(mutableSetOf()) { it.path.path }, nextName = { newNames[i].also { i += 1 } })
     }
 }
 
-internal fun PluginDataFrameSchema.map(selected: ColumnsSet, nextName: () -> String): PluginDataFrameSchema {
+internal fun PluginDataFrameSchema.rename(selected: ColumnsSet, nextName: () -> String): PluginDataFrameSchema {
     return PluginDataFrameSchema(
         f(columns(), nextName, selected, emptyList())
     )

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/toDataFrame.kt
@@ -65,7 +65,15 @@ class ToDataFrameDsl : AbstractSchemaModificationInterpreter() {
     val Arguments.body by dsl()
     override fun Arguments.interpret(): PluginDataFrameSchema {
         val dsl = CreateDataFrameDslImplApproximation()
-        body(dsl, mapOf("explicitReceiver" to Interpreter.Success(receiver)))
+        val receiver = receiver ?: return PluginDataFrameSchema(emptyList())
+        val arg = receiver.resolvedType.typeArguments.firstOrNull() ?: return PluginDataFrameSchema(emptyList())
+        when {
+            arg.isStarProjection -> PluginDataFrameSchema(emptyList())
+            else -> {
+                val classLike = arg.type as? ConeClassLikeType ?: return PluginDataFrameSchema(emptyList())
+                body(dsl, mapOf(Properties0.classExtraArgument to Interpreter.Success(classLike)))
+            }
+        }
         return PluginDataFrameSchema(dsl.columns)
     }
 }
@@ -90,15 +98,19 @@ class ToDataFrameDefault : AbstractSchemaModificationInterpreter() {
 private const val DEFAULT_MAX_DEPTH = 0
 
 class Properties0 : AbstractInterpreter<Unit>() {
+    companion object {
+        const val classExtraArgument = "explicitReceiver"
+    }
+
     val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
-    val Arguments.explicitReceiver: FirExpression? by arg()
+    val Arguments.coneKotlinType: ConeKotlinType by arg(name = name(classExtraArgument))
     val Arguments.maxDepth: Int by arg()
     val Arguments.body by dsl()
 
     override fun Arguments.interpret() {
         dsl.configuration.maxDepth = maxDepth
         body(dsl.configuration.traverseConfiguration, emptyMap())
-        val schema = toDataFrame(dsl.configuration.maxDepth, explicitReceiver, dsl.configuration.traverseConfiguration)
+        val schema = toDataFrame(dsl.configuration.maxDepth, coneKotlinType, dsl.configuration.traverseConfiguration)
         dsl.columns.addAll(schema.columns())
     }
 }
@@ -154,7 +166,7 @@ class Exclude1 : AbstractInterpreter<Unit>() {
 @OptIn(SymbolInternals::class)
 internal fun KotlinTypeFacade.toDataFrame(
     maxDepth: Int,
-    explicitReceiver: FirExpression?,
+    classLikeType: ConeKotlinType,
     traverseConfiguration: TraverseConfiguration
 ): PluginDataFrameSchema {
     fun ConeKotlinType.isValueType() =
@@ -269,14 +281,21 @@ internal fun KotlinTypeFacade.toDataFrame(
             }
     }
 
+    return PluginDataFrameSchema(convert(classLikeType, 0))
+}
+
+internal fun KotlinTypeFacade.toDataFrame(
+    maxDepth: Int,
+    explicitReceiver: FirExpression?,
+    traverseConfiguration: TraverseConfiguration
+): PluginDataFrameSchema {
     val receiver = explicitReceiver ?: return PluginDataFrameSchema(emptyList())
     val arg = receiver.resolvedType.typeArguments.firstOrNull() ?: return PluginDataFrameSchema(emptyList())
     return when {
         arg.isStarProjection -> PluginDataFrameSchema(emptyList())
         else -> {
             val classLike = arg.type as? ConeClassLikeType ?: return PluginDataFrameSchema(emptyList())
-            val columns = convert(classLike, 0)
-            PluginDataFrameSchema(columns)
+            return toDataFrame(maxDepth, classLike, traverseConfiguration)
         }
     }
 }

diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/unfold.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/unfold.kt
@@ -0,0 +1,38 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.api
+
+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.Interpreter
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+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.data.ReplaceClauseApproximation
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dataFrame
+import org.jetbrains.kotlinx.dataframe.plugin.impl.dsl
+
+class ReplaceUnfold1 : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: ReplaceClauseApproximation by arg()
+    val Arguments.body by dsl()
+    val Arguments.typeArg1: TypeApproximation by arg()
+
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        val configuration = CreateDataFrameDslImplApproximation()
+        body(configuration, mapOf(Properties0.classExtraArgument to Interpreter.Success(typeArg1.type)))
+
+        return receiver.df.map(receiver.columns.map { it.path.path }.toSet()) { a, column ->
+            if (column is SimpleFrameColumn || column is SimpleColumnGroup) return@map column
+            SimpleColumnGroup(column.name, configuration.columns)
+        }
+    }
+}
+
+class Replace0 : AbstractInterpreter<ReplaceClauseApproximation>() {
+    val Arguments.receiver: PluginDataFrameSchema by dataFrame()
+    val Arguments.columns: List<ColumnWithPathApproximation> by arg()
+
+    override fun Arguments.interpret(): ReplaceClauseApproximation {
+        return ReplaceClauseApproximation(receiver, columns)
+    }
+}
diff --git a/...aframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/data/ReplaceClauseApproximation.kt b/...aframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/data/ReplaceClauseApproximation.kt
@@ -0,0 +1,5 @@
+package org.jetbrains.kotlinx.dataframe.plugin.impl.data
+
+import org.jetbrains.kotlinx.dataframe.plugin.impl.PluginDataFrameSchema
+
+class ReplaceClauseApproximation(val df: PluginDataFrameSchema, val columns: List<ColumnWithPathApproximation>)
diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt
@@ -67,6 +67,8 @@ import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Replace0
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReplaceUnfold1
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDsl
@@ -167,6 +169,8 @@ internal inline fun <reified T> String.load(): T {
         "toDataFrameDsl" -> ToDataFrameDsl()
         "toDataFrame" -> ToDataFrame()
         "toDataFrameDefault" -> ToDataFrameDefault()
+        "Replace0" -> Replace0()
+        "ReplaceUnfold1" -> ReplaceUnfold1()
         else -> error("$this")
     } as T
 }