[Compiler plugin] Add ability to extract schema from implicit this receiver

Author: koperagen

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

@@ -13,6 +13,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.data.ColumnPathApproximation
 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.ignore
 import org.jetbrains.kotlinx.dataframe.plugin.impl.type
 
 internal class Select0 : AbstractInterpreter<PluginDataFrameSchema>() {
@@ -25,6 +26,7 @@ internal class Select0 : AbstractInterpreter<PluginDataFrameSchema>() {
 }
 
 internal class Expr0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
     val Arguments.name: String by arg(defaultValue = Present("untitled"))
     val Arguments.infer: Infer by enum(defaultValue = Present(Infer.Nulls))
     val Arguments.expression: TypeApproximation by type()
@@ -53,6 +55,7 @@ internal class And0 : AbstractInterpreter<ColumnsResolver>() {
 }
 
 internal class All0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
     override fun Arguments.interpret(): ColumnsResolver {
         return object : ColumnsResolver {
             override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {
@@ -66,6 +69,7 @@ internal class All0 : AbstractInterpreter<ColumnsResolver>() {
 }
 
 internal class ColsOf0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
     val Arguments.typeArg0: TypeApproximation by arg()
 
     override fun Arguments.interpret(): ColumnsResolver {
@@ -90,6 +94,7 @@ private fun Arguments.colsOf(cols: List<ColumnWithPathApproximation>, type: Cone
         }
 
 internal class ColsAtAnyDepth0 : AbstractInterpreter<ColumnsResolver>() {
+    val Arguments.receiver by ignore()
     override fun Arguments.interpret(): ColumnsResolver {
         return object : ColumnsResolver {
             override fun resolve(df: PluginDataFrameSchema): List<ColumnWithPathApproximation> {

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

@@ -505,7 +505,7 @@ internal fun FirFunctionCall.collectArgumentExpressions(): RefinedArguments {
     val refinedArgument = mutableListOf<RefinedArgument>()
 
     val parameterName = Name.identifier("receiver")
-    explicitReceiver?.let {
+    (explicitReceiver ?: extensionReceiver)?.let {
         if (it is FirResolvedQualifier && it.resolvedToCompanionObject) {
             return@let
         }

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

@@ -0,0 +1,19 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+data class Nested(val d: Double)
+
+data class Record(val a: String, val b: Int, val nested: Nested)
+
+fun box(): String {
+    val df = dataFrameOf("a", "b", "c")(1, 2, 3)
+
+    df.groupBy { a }
+        .updateGroups { remove { a } }
+        .aggregate { c into "c" }
+    return "OK"
+}
+
+

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

@@ -0,0 +1,23 @@
+import org.jetbrains.kotlinx.dataframe.*
+import org.jetbrains.kotlinx.dataframe.annotations.*
+import org.jetbrains.kotlinx.dataframe.api.*
+import org.jetbrains.kotlinx.dataframe.io.*
+
+@DataSchema
+interface JoinLeaf {
+    val something: Int
+    val somethingElse: String
+}
+
+@DataSchema
+interface Join2 {
+    val c: DataRow<JoinLeaf>
+}
+
+fun selectionDsl(df: DataFrame<Join2>) {
+    df.ungroup { c }.select { colsOf<String>() }.somethingElse
+}
+
+fun box(): String {
+    return "OK"
+}

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

@@ -454,12 +454,24 @@ public void testSchema() {
     runTest("testData/box/Schema.kt");
   }
 
+  @Test
+  @TestMetadata("schemaFromImplicitReceiver.kt")
+  public void testSchemaFromImplicitReceiver() {
+    runTest("testData/box/schemaFromImplicitReceiver.kt");
+  }
+
   @Test
   @TestMetadata("select.kt")
   public void testSelect() {
     runTest("testData/box/select.kt");
   }
 
+  @Test
+  @TestMetadata("selectColsOf.kt")
+  public void testSelectColsOf() {
+    runTest("testData/box/selectColsOf.kt");
+  }
+
   @Test
   @TestMetadata("selectIt.kt")
   public void testSelectIt() {