Skip to content

Commit 61865d8

Browse files
koperagenkoperagen
authored
Merge pull request #914 from Kotlin/stdlib-interpret
[Compiler plugin] Add a mechanism to handle function calls to stdlib that can appear as df api arguments
2 parents 0e3fdcd + 650222e commit 61865d8

File tree

7 files changed

+100
-13
lines changed

7 files changed

+100
-13
lines changed

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

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,38 @@ package org.jetbrains.kotlinx.dataframe.plugin.impl.api
33
import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter
44
import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments
55
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter
6+
import org.jetbrains.kotlinx.dataframe.plugin.impl.Present
67

7-
class PairConstructor : AbstractInterpreter<Pair<*, *>>() {
8+
class PairToConstructor : AbstractInterpreter<Pair<*, *>>() {
89
val Arguments.receiver: Any? by arg(lens = Interpreter.Id)
910
val Arguments.that: Any? by arg(lens = Interpreter.Id)
1011
override fun Arguments.interpret(): Pair<*, *> {
1112
return receiver to that
1213
}
1314
}
15+
16+
class PairConstructor : AbstractInterpreter<Pair<*, *>>() {
17+
val Arguments.first: Any? by arg(lens = Interpreter.Id)
18+
val Arguments.second: Any? by arg(lens = Interpreter.Id)
19+
override fun Arguments.interpret(): Pair<*, *> {
20+
return first to second
21+
}
22+
}
23+
24+
class TrimMargin : AbstractInterpreter<String>() {
25+
val Arguments.receiver: String by arg(lens = Interpreter.Value)
26+
val Arguments.marginPrefix: String by arg(lens = Interpreter.Value, defaultValue = Present("|"))
27+
28+
override fun Arguments.interpret(): String {
29+
return receiver.trimMargin(marginPrefix)
30+
}
31+
}
32+
33+
class TrimIndent : AbstractInterpreter<String>() {
34+
val Arguments.receiver: String by arg(lens = Interpreter.Value)
35+
36+
override fun Arguments.interpret(): String {
37+
return receiver.trimIndent()
38+
}
39+
}
40+

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

Lines changed: 31 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -62,16 +62,13 @@ import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
6262
import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
6363
import org.jetbrains.kotlin.fir.expressions.UnresolvedExpressionTypeAccess
6464
import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
65-
import org.jetbrains.kotlin.fir.references.resolved
66-
import org.jetbrains.kotlin.fir.references.symbol
67-
import org.jetbrains.kotlin.fir.references.toResolvedNamedFunctionSymbol
65+
import org.jetbrains.kotlin.fir.references.toResolvedFunctionSymbol
6866
import org.jetbrains.kotlin.fir.resolve.fqName
6967
import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
7068
import org.jetbrains.kotlin.fir.types.classId
7169
import org.jetbrains.kotlin.fir.types.coneType
7270
import org.jetbrains.kotlin.name.CallableId
7371
import org.jetbrains.kotlin.name.ClassId
74-
import org.jetbrains.kotlin.name.FqName
7572
import org.jetbrains.kotlin.name.Name
7673
import org.jetbrains.kotlin.name.StandardClassIds
7774
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke
@@ -91,25 +88,22 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0
9188
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame
9289
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Move0
9390
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairConstructor
91+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairToConstructor
9492
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReadExcel
9593
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame
9694
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn
9795
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDefault
9896
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameDsl
9997
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameFrom
10098
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToTop
99+
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.TrimMargin
101100
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Update0
102101
import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0
103102
import org.jetbrains.kotlinx.dataframe.plugin.utils.Names
104103

105-
@OptIn(UnresolvedExpressionTypeAccess::class)
106104
internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? {
107-
if (
108-
calleeReference.toResolvedNamedFunctionSymbol()?.callableId == Names.TO &&
109-
coneTypeOrNull?.classId == Names.PAIR
110-
) {
111-
return PairConstructor()
112-
}
105+
val interpreter = Stdlib.interpreter(this)
106+
if (interpreter != null) return interpreter
113107
val symbol =
114108
(calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirCallableSymbol ?: return null
115109
val argName = Name.identifier("interpreter")
@@ -121,6 +115,32 @@ internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*
121115
}
122116
}
123117

118+
private object Stdlib {
119+
private val map: MutableMap<Key, Interpreter<*>> = mutableMapOf()
120+
init {
121+
register(Names.TO, Names.PAIR, PairToConstructor())
122+
register(Names.PAIR_CONSTRUCTOR, Names.PAIR, PairConstructor())
123+
register(Names.TRIM_MARGIN, StandardClassIds.String, TrimMargin())
124+
register(Names.TRIM_INDENT, StandardClassIds.String, TrimMargin())
125+
}
126+
127+
@OptIn(UnresolvedExpressionTypeAccess::class)
128+
fun interpreter(call: FirFunctionCall): Interpreter<*>? {
129+
val id = call.calleeReference.toResolvedFunctionSymbol()?.callableId ?: return null
130+
val returnType = call.coneTypeOrNull?.classId ?: return null
131+
return map[Key(id, returnType)]
132+
}
133+
134+
fun register(id: CallableId, returnType: ClassId, interpreter: Interpreter<*>) {
135+
map[Key(id, returnType)] = interpreter
136+
}
137+
}
138+
139+
private data class Key(
140+
val id: CallableId,
141+
val returnType: ClassId,
142+
)
143+
124144
internal fun FirFunctionCall.interpreterName(session: FirSession): String? {
125145
val symbol =
126146
(calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirCallableSymbol ?: return null

‎plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
package org.jetbrains.kotlinx.dataframe.plugin.utils
77

8+
import org.jetbrains.kotlin.builtins.StandardNames
89
import org.jetbrains.kotlin.name.CallableId
910
import org.jetbrains.kotlin.name.ClassId
1011
import org.jetbrains.kotlin.name.FqName
@@ -53,7 +54,10 @@ object Names {
5354
val INSTANT_CLASS_ID = kotlinx.datetime.Instant::class.classId()
5455

5556
val PAIR = ClassId(FqName("kotlin"), Name.identifier("Pair"))
57+
val PAIR_CONSTRUCTOR = CallableId(FqName("kotlin"), FqName("Pair"), Name.identifier("Pair"))
5658
val TO = CallableId(FqName("kotlin"), Name.identifier("to"))
59+
val TRIM_MARGIN = CallableId(StandardNames.TEXT_PACKAGE_FQ_NAME, Name.identifier("trimMargin"))
60+
val TRIM_INDENT = CallableId(StandardNames.TEXT_PACKAGE_FQ_NAME, Name.identifier("trimIndent"))
5761
}
5862

5963
private fun KClass<*>.classId(): ClassId {

‎plugins/kotlin-dataframe/testData/box/dataFrameOf_to.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import org.jetbrains.kotlinx.dataframe.io.*
55

66
fun box(): String {
77
val df = dataFrameOf(
8-
"a" to listOf(1, 2),
8+
Pair("a", listOf(1, 2))
99
"b" to listOf("str1", "str2"),
1010
)
1111
val i: Int = df.a[0]

‎plugins/kotlin-dataframe/testData/box/trimIndent.kt

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
fun box(): String {
7+
val df = DataFrame.readJsonStr("""
8+
{"a": 123}
9+
""".trimIndent())
10+
df.a
11+
return "OK"
12+
}

‎plugins/kotlin-dataframe/testData/box/trimMargin.kt

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
import org.jetbrains.kotlinx.dataframe.*
2+
import org.jetbrains.kotlinx.dataframe.annotations.*
3+
import org.jetbrains.kotlinx.dataframe.api.*
4+
import org.jetbrains.kotlinx.dataframe.io.*
5+
6+
fun box(): String {
7+
val df = DataFrame.readJsonStr("""
8+
|{"a": 123}
9+
|""".trimMargin())
10+
df.a
11+
return "OK"
12+
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -442,6 +442,18 @@ public void testTransformReplaceFunctionCall() {
442442
runTest("testData/box/transformReplaceFunctionCall.kt");
443443
}
444444

445+
@Test
446+
@TestMetadata("trimIndent.kt")
447+
public void testTrimIndent() {
448+
runTest("testData/box/trimIndent.kt");
449+
}
450+
451+
@Test
452+
@TestMetadata("trimMargin.kt")
453+
public void testTrimMargin() {
454+
runTest("testData/box/trimMargin.kt");
455+
}
456+
445457
@Test
446458
@TestMetadata("ungroup.kt")
447459
public void testUngroup() {

0 commit comments

Comments
 (0)