[Compiler plugin] Generate valid code in transform(call) when interpret(call) fails

Author: koperagen

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

@@ -30,7 +30,7 @@ internal inline fun <reified T> KotlinTypeFacade.analyzeRefinedCallShape(
     val rootMarkers = callReturnType.typeArguments.filterIsInstance<ConeClassLikeType>()
     if (rootMarkers.size != callReturnType.typeArguments.size) return null
 
-    val newSchema: T = call.interpreterName(session)?.let { name ->
+    val newSchema: T? = call.interpreterName(session)?.let { name ->
         when (name) {
             else -> name.load<Interpreter<*>>().let { processor ->
                 val dataFrameSchema = interpret(call, processor, reporter = reporter)
@@ -40,19 +40,20 @@ internal inline fun <reified T> KotlinTypeFacade.analyzeRefinedCallShape(
                             if (!reporter.errorReported) {
                                 reporter.reportInterpretationError(call, "${processor::class} must return ${T::class}, but was $value")
                             }
-                            return null
+                            null
+                        } else {
+                            value
                         }
-                        value
                     }
                 dataFrameSchema
             }
         }
-    } ?: return null
+    }
 
     return CallResult(rootMarkers, newSchema)
 }
 
-data class CallResult<T>(val markers: List<ConeClassLikeType>, val result: T)
+data class CallResult<T>(val markers: List<ConeClassLikeType>, val result: T?)
 
 class RefinedArguments(val refinedArguments: List<RefinedArgument>) : List<RefinedArgument> by refinedArguments
 

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

@@ -97,6 +97,11 @@ class FunctionCallTransformer(
 
     private interface CallTransformer {
         fun interceptOrNull(callInfo: CallInfo, symbol: FirNamedFunctionSymbol, hash: String): CallReturnType?
+
+        /**
+         * must still generate let with declared class from interceptOrNull when interpretation fails.
+         * it should only return null if later some frontend checker fails compilation in general
+         */
         fun transformOrNull(call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol): FirFunctionCall?
     }
 
@@ -180,7 +185,7 @@ class FunctionCallTransformer(
             val (tokens, dataFrameSchema) = callResult ?: return null
             val token = tokens[0]
             val firstSchema = token.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
-            val dataSchemaApis = materialize(dataFrameSchema, call, firstSchema)
+            val dataSchemaApis = materialize(dataFrameSchema ?: PluginDataFrameSchema.EMPTY, call, firstSchema)
 
             val tokenFir = token.toClassSymbol(session)!!.fir
             tokenFir.callShapeData = CallShapeData.RefinedType(dataSchemaApis.map { it.scope.symbol })
@@ -228,8 +233,13 @@ class FunctionCallTransformer(
             val keyMarker = rootMarkers[0]
             val groupMarker = rootMarkers[1]
 
-            val keySchema = createPluginDataFrameSchema(groupBy.keys, groupBy.moveToTop)
-            val groupSchema = PluginDataFrameSchema(groupBy.df.columns())
+            val (keySchema, groupSchema) = if (groupBy != null) {
+                val keySchema = createPluginDataFrameSchema(groupBy.keys, groupBy.moveToTop)
+                val groupSchema = PluginDataFrameSchema(groupBy.df.columns())
+                keySchema to groupSchema
+            } else {
+                PluginDataFrameSchema.EMPTY to PluginDataFrameSchema.EMPTY
+            }
 
             val firstSchema = keyMarker.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!
             val firstSchema1 = groupMarker.toClassSymbol(session)?.resolvedSuperTypes?.get(0)!!.toRegularClassSymbol(session)?.fir!!

plugins/kotlin-dataframe/testData/box/inventNamesForLocalClasses.kt

@@ -0,0 +1,9 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+fun box(): String {
+    val df = dataFrameOf("versions")(mapOf("a" to 1)).convert { versions }.with { dataFrameOf(it.keys)(it.values) }
+    return "OK"
+}

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

@@ -226,6 +226,12 @@ public void testInsert() {
     runTest("testData/box/insert.kt");
   }
 
+  @Test
+  @TestMetadata("inventNamesForLocalClasses.kt")
+  public void testInventNamesForLocalClasses() {
+    runTest("testData/box/inventNamesForLocalClasses.kt");
+  }
+
   @Test
   @TestMetadata("join.kt")
   public void testJoin() {