Merge pull request #907 from Kotlin/invent-local-names-fix

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/src/org/jetbrains/kotlinx/dataframe/plugin/impl/SimpleCol.kt

@@ -18,6 +18,9 @@ import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
 data class PluginDataFrameSchema(
     private val columns: List<SimpleCol>
 ) : DataFrameLikeContainer<SimpleCol> {
+    companion object {
+        val EMPTY = PluginDataFrameSchema(emptyList())
+    }
     override fun columns(): List<SimpleCol> {
         return columns
     }

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

@@ -92,7 +92,7 @@ fun KotlinTypeFacade.aggregate(
         }
         PluginDataFrameSchema(cols)
     } else {
-        PluginDataFrameSchema(emptyList())
+        PluginDataFrameSchema.EMPTY
     }
 }
 

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

@@ -132,7 +132,7 @@ internal fun PluginDataFrameSchema.insertImpl(
         columns.firstOrNull()?.referenceNode?.getRoot(),
         0,
         factory = { PluginDataFrameSchema(it) },
-        empty = PluginDataFrameSchema(emptyList()),
+        empty = PluginDataFrameSchema.EMPTY,
         rename = { rename(it) },
         createColumnGroup = { name, columns ->
             SimpleColumnGroup(name, columns)

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

@@ -288,12 +288,12 @@ internal fun KotlinTypeFacade.toDataFrame(
             }
     }
 
-    val receiver = explicitReceiver ?: return PluginDataFrameSchema(emptyList())
-    val arg = receiver.resolvedType.typeArguments.firstOrNull() ?: return PluginDataFrameSchema(emptyList())
+    val receiver = explicitReceiver ?: return PluginDataFrameSchema.EMPTY
+    val arg = receiver.resolvedType.typeArguments.firstOrNull() ?: return PluginDataFrameSchema.EMPTY
     return when {
-        arg.isStarProjection -> PluginDataFrameSchema(emptyList())
+        arg.isStarProjection -> PluginDataFrameSchema.EMPTY
         else -> {
-            val classLike = arg.type as? ConeClassLikeType ?: return PluginDataFrameSchema(emptyList())
+            val classLike = arg.type as? ConeClassLikeType ?: return PluginDataFrameSchema.EMPTY
             val columns = convert(classLike, 0)
             PluginDataFrameSchema(columns)
         }

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

@@ -314,16 +314,16 @@ interface InterpretationErrorReporter {
 
 fun KotlinTypeFacade.pluginDataFrameSchema(schemaTypeArg: ConeTypeProjection): PluginDataFrameSchema {
     val schema = if (schemaTypeArg.isStarProjection) {
-        PluginDataFrameSchema(emptyList())
+        PluginDataFrameSchema.EMPTY
     } else {
-        val coneClassLikeType = schemaTypeArg.type as? ConeClassLikeType ?: return PluginDataFrameSchema(emptyList())
+        val coneClassLikeType = schemaTypeArg.type as? ConeClassLikeType ?: return PluginDataFrameSchema.EMPTY
         pluginDataFrameSchema(coneClassLikeType)
     }
     return schema
 }
 
 fun KotlinTypeFacade.pluginDataFrameSchema(coneClassLikeType: ConeClassLikeType): PluginDataFrameSchema {
-    val symbol = coneClassLikeType.toSymbol(session) as? FirRegularClassSymbol ?: return PluginDataFrameSchema(emptyList())
+    val symbol = coneClassLikeType.toSymbol(session) as? FirRegularClassSymbol ?: return PluginDataFrameSchema.EMPTY
     val declarationSymbols = if (symbol.isLocal && symbol.resolvedSuperTypes.firstOrNull() != session.builtinTypes.anyType.type) {
         val rootSchemaSymbol = symbol.resolvedSuperTypes.first().toSymbol(session) as? FirRegularClassSymbol
         rootSchemaSymbol?.declaredMemberScope(session, FirResolvePhase.DECLARATIONS)

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() {