From 1349a8d8acc1a39e258eaf1b9a06ac2401138cb8 Mon Sep 17 00:00:00 2001
From: Nikita Klimenko <nikita.klimenko@jetbrains.com>
Date: Thu, 23 Jan 2025 00:22:33 +0100
Subject: [PATCH] [Compiler plugin] DSL for creating nested column groups in
 toDataFrame

---
 .../jetbrains/kotlinx/dataframe/api/insert.kt |  1 +
 .../kotlinx/dataframe/api/toDataFrame.kt      |  1 +
 .../dataframe/plugin/impl/api/toDataFrame.kt  | 12 ++++++++++++
 .../dataframe/plugin/loadInterpreter.kt       |  2 ++
 .../testData/box/toDataFrame_nested.kt        | 19 +++++++++++++++++++
 ...DataFrameBlackBoxCodegenTestGenerated.java |  6 ++++++
 6 files changed, 41 insertions(+)
 create mode 100644 plugins/kotlin-dataframe/testData/box/toDataFrame_nested.kt

diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt
index 61f3b59e68..03efde6593 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/insert.kt
@@ -83,6 +83,7 @@ public fun <T> InsertClause<T>.after(column: ColumnSelector<T, *>): DataFrame<T>
 
 public fun <T> InsertClause<T>.after(column: String): DataFrame<T> = df.add(this.column).move(this.column).after(column)
 
+@AccessApiOverload
 public fun <T> InsertClause<T>.after(column: ColumnAccessor<*>): DataFrame<T> = after(column.path())
 
 @AccessApiOverload
diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt
index b26eaf36a9..6c81fc73fc 100644
--- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt
+++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/toDataFrame.kt
@@ -202,6 +202,7 @@ public abstract class CreateDataFrameDsl<T> : TraversePropertiesDsl {
 
     public inline fun <reified R> inferType(noinline expression: (T) -> R): InferType<T, R> = InferType(expression)
 
+    @Interpretable("ToDataFrameDslStringInvoke")
     public abstract operator fun String.invoke(builder: CreateDataFrameDsl<T>.() -> Unit)
 }
 
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
index 60fb38b76a..4b632b80f0 100644
--- 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
@@ -131,6 +131,18 @@ class Properties0 : AbstractInterpreter<Unit>() {
     }
 }
 
+class ToDataFrameDslStringInvoke : AbstractInterpreter<Unit>() {
+    val Arguments.dsl: CreateDataFrameDslImplApproximation by arg()
+    val Arguments.receiver: String by arg()
+    val Arguments.builder by dsl()
+
+    override fun Arguments.interpret() {
+        val addDsl = CreateDataFrameDslImplApproximation()
+        builder(addDsl, emptyMap())
+        dsl.columns.add(SimpleColumnGroup(receiver, addDsl.columns))
+    }
+}
+
 class CreateDataFrameConfiguration {
     var maxDepth = DEFAULT_MAX_DEPTH
     var traverseConfiguration: TraverseConfiguration = TraverseConfiguration()
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
index 8d13a2d37f..9c3f0e90fa 100644
--- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt
+++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/loadInterpreter.kt
@@ -103,6 +103,7 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDsl
+import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDslStringInvoke
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TrimMargin
@@ -242,6 +243,7 @@ internal inline fun <reified T> String.load(): T {
         "toDataFrameDsl" -> ToDataFrameDsl()
         "toDataFrame" -> ToDataFrame()
         "toDataFrameDefault" -> ToDataFrameDefault()
+        "ToDataFrameDslStringInvoke" -> ToDataFrameDslStringInvoke()
         "DataFrameOf0" -> DataFrameOf0()
         "DataFrameBuilderInvoke0" -> DataFrameBuilderInvoke0()
         "ToDataFrameColumn" -> ToDataFrameColumn()
diff --git a/plugins/kotlin-dataframe/testData/box/toDataFrame_nested.kt b/plugins/kotlin-dataframe/testData/box/toDataFrame_nested.kt
new file mode 100644
index 0000000000..34409a1699
--- /dev/null
+++ b/plugins/kotlin-dataframe/testData/box/toDataFrame_nested.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.*
+
+fun box(): String {
+    val df = listOf("a").toDataFrame {
+        "group" {
+            "anotherGroup" {
+                "a" from { 1 }
+            }
+            "b" from { "c" }
+        }
+    }
+
+    df.group.b
+    df.group.anotherGroup.a
+    return "OK"
+}
diff --git a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java
index 12cb1289f7..d736fc4081 100644
--- a/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java
+++ b/plugins/kotlin-dataframe/tests-gen/org/jetbrains/kotlin/fir/dataframe/DataFrameBlackBoxCodegenTestGenerated.java
@@ -508,6 +508,12 @@ public void testToDataFrame_from() {
     runTest("testData/box/toDataFrame_from.kt");
   }
 
+  @Test
+  @TestMetadata("toDataFrame_nested.kt")
+  public void testToDataFrame_nested() {
+    runTest("testData/box/toDataFrame_nested.kt");
+  }
+
   @Test
   @TestMetadata("toDataFrame_nullableList.kt")
   public void testToDataFrame_nullableList() {