[Compiler plugin] Remove code duplication caused by aggregate implementation

Author: koperagen

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

@@ -35,15 +35,6 @@ fun KotlinTypeFacade.analyzeRefinedCallShape(call: FirFunctionCall, reporter: In
 
     val newSchema: PluginDataFrameSchema = call.interpreterName(session)?.let { name ->
         when (name) {
-            "Aggregate" -> {
-                val groupByCall = call.explicitReceiver as? FirFunctionCall
-                val interpreter = groupByCall?.loadInterpreter(session)
-                if (interpreter != null) {
-                    aggregate(groupByCall, interpreter, reporter, call)
-                } else {
-                    PluginDataFrameSchema(emptyList())
-                }
-            }
             else -> name.load<Interpreter<*>>().let { processor ->
                 val dataFrameSchema = interpret(call, processor, reporter = reporter)
                     .let {

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

@@ -279,8 +279,8 @@ class FunctionCallTransformer(
             val firstSchema = keyMarker.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
             val firstSchema1 = groupMarker.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
 
-            val dataSchemaApis = materialize(keySchema, call, firstSchema)
-            val dataSchemaApis1 = materialize(groupSchema, call, firstSchema1, dataSchemaApis.size)
+            val dataSchemaApis = materialize(keySchema, call, firstSchema, "Key")
+            val dataSchemaApis1 = materialize(groupSchema, call, firstSchema1, "Group", dataSchemaApis.size)
 
             val tokenFir = keyMarker.toClassSymbol(session)!!.fir
             tokenFir.callShapeData = CallShapeData.RefinedType(dataSchemaApis.map { it.scope.symbol })
@@ -431,6 +431,7 @@ class FunctionCallTransformer(
         dataFrameSchema: PluginDataFrameSchema,
         call: FirFunctionCall,
         firstSchema: FirRegularClass,
+        prefix: String = "",
         i: Int = 0
     ): List<DataSchemaApi> {
         var i = i
@@ -469,7 +470,7 @@ class FunctionCallTransformer(
                     val text = call.source?.text ?: call.calleeReference.name
                     val name =
                         "${column.name.titleCase().replEscapeLineBreaks()}_${hashToTwoCharString(abs(text.hashCode()))}"
-                    return materialize(suggestedName = name)
+                    return materialize(suggestedName = "$prefix$name")
                 }
 
                 when (it) {

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

@@ -3,15 +3,11 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api
 import org.jetbrains.kotlinx.dataframe.plugin.InterpretationErrorReporter
 import org.jetbrains.kotlinx.dataframe.plugin.interpret
 import org.jetbrains.kotlinx.dataframe.plugin.loadInterpreter
-import org.jetbrains.kotlinx.dataframe.plugin.pluginDataFrameSchema
-import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 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.resolvedType
 import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
 import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
@@ -55,21 +51,30 @@ class GroupByInto : AbstractInterpreter<Unit>() {
     }
 }
 
+class Aggregate : AbstractSchemaModificationInterpreter() {
+    val Arguments.receiver: GroupBy by arg()
+    val Arguments.body: FirAnonymousFunctionExpression by arg(lens = Interpreter.Id)
+    override fun Arguments.interpret(): PluginDataFrameSchema {
+        return aggregate(
+            receiver,
+            InterpretationErrorReporter.DEFAULT,
+            body
+        )
+    }
+}
+
 fun KotlinTypeFacade.aggregate(
-    groupByCall: FirFunctionCall,
-    interpreter: Interpreter<*>,
+    groupBy: GroupBy,
     reporter: InterpretationErrorReporter,
-    call: FirFunctionCall
-): PluginDataFrameSchema? {
-    val groupBy = interpret(groupByCall, interpreter, reporter = reporter)?.value as? GroupBy ?: return null
-    val aggregate = call.argumentList.arguments.singleOrNull() as? FirAnonymousFunctionExpression
-    val body = aggregate?.anonymousFunction?.body ?: return null
-    val lastExpression = (body.statements.lastOrNull() as? FirReturnExpression)?.result
+    firAnonymousFunctionExpression: FirAnonymousFunctionExpression
+): PluginDataFrameSchema {
+    val body = firAnonymousFunctionExpression.anonymousFunction.body
+    val lastExpression = (body?.statements?.lastOrNull() as? FirReturnExpression)?.result
     val type = lastExpression?.resolvedType
     return if (type != session.builtinTypes.unitType) {
         val dsl = GroupByDsl()
         val calls = buildList {
-            body.statements.filterIsInstance<FirFunctionCall>().let { addAll(it) }
+            body?.statements?.filterIsInstance<FirFunctionCall>()?.let { addAll(it) }
             if (lastExpression is FirFunctionCall) add(lastExpression)
         }
         calls.forEach { call ->
@@ -87,7 +92,7 @@ fun KotlinTypeFacade.aggregate(
         }
         PluginDataFrameSchema(cols)
     } else {
-        null
+        PluginDataFrameSchema(emptyList())
     }
 }
 

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

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.expressions.FirThisReceiverExpression
 import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression
+import org.jetbrains.kotlin.fir.expressions.arguments
 import org.jetbrains.kotlin.fir.expressions.impl.FirResolvedArgumentList
 import org.jetbrains.kotlin.fir.references.FirResolvedCallableReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
@@ -144,10 +145,22 @@ fun <T> KotlinTypeFacade.interpret(
                     }
 
                     is FirFunctionCall -> {
-                        val interpreter = expression.loadInterpreter()
-                        interpreter?.let {
-                            val result = interpret(expression, interpreter, emptyMap(), reporter)
-                            result
+                        var interpreter = expression.loadInterpreter()
+                        if (interpreter == null) {
+                            val r = expression.arguments[0]
+                            val last = (r as? FirAnonymousFunctionExpression)?.anonymousFunction?.body?.statements?.lastOrNull()
+                            val call = (last as? FirReturnExpression)?.result as? FirFunctionCall
+                            val interpreter = call?.loadInterpreter()
+                            if (interpreter != null) {
+                                interpret(call, interpreter, emptyMap(), reporter)
+                            } else {
+                                null
+                            }
+                        } else {
+                            interpreter?.let {
+                                val result = interpret(expression, interpreter, emptyMap(), reporter)
+                                result
+                            }
                         }
                     }
 

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

@@ -69,6 +69,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Aggregate
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.All0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsAtAnyDepth0
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0
@@ -206,6 +207,7 @@ internal inline fun <reified T> String.load(): T {
         "Move0" -> Move0()
         "ToTop" -> ToTop()
         "Update0" -> Update0()
+        "Aggregate" -> Aggregate()
         else -> error("$this")
     } as T
 }