Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Compiler plugin] Lower frontend generated implicit receivers #869

Merged
merged 1 commit into from
Oct 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package org.jetbrains.kotlinx.dataframe.aggregation

import org.jetbrains.kotlinx.dataframe.DataFrame
import org.jetbrains.kotlinx.dataframe.annotations.HasSchema
import org.jetbrains.kotlinx.dataframe.annotations.Interpretable
import org.jetbrains.kotlinx.dataframe.api.ColumnSelectionDsl
import org.jetbrains.kotlinx.dataframe.api.pathOf
Expand All @@ -11,6 +12,7 @@ import org.jetbrains.kotlinx.dataframe.impl.columnName
import kotlin.reflect.KProperty
import kotlin.reflect.typeOf

@HasSchema(schemaArg = 0)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does HasSchema do again? There's no docs :/

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To help identify type parameter that represents dataframe schema
org.jetbrains.kotlinx.dataframe.plugin.InterpretKt#getSchema
org/jetbrains/kotlinx/dataframe/plugin/interpret.kt:255

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

right, so if we would take the type parameter at schemaArg, which in AggregateDsl<out T> is out T, and put it in DataFrame<> like DataFrame<T>, a schema generated for thát would represent this object. And this could be different for types like GroupedDataRow<out T, out G> where G might represent the data schema?

Could you explain this in a small piece of KDoc?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll think of something

public abstract class AggregateDsl<out T> :
DataFrame<T>,
ColumnSelectionDsl<T> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import org.jetbrains.kotlin.ir.expressions.IrConst
import org.jetbrains.kotlin.ir.expressions.IrErrorCallExpression
import org.jetbrains.kotlin.ir.expressions.IrExpression
import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
import org.jetbrains.kotlin.ir.expressions.impl.IrConstructorCallImpl
Expand All @@ -41,6 +42,7 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
import org.jetbrains.kotlin.ir.types.IrSimpleType
import org.jetbrains.kotlin.ir.types.IrType
import org.jetbrains.kotlin.ir.types.classFqName
import org.jetbrains.kotlin.ir.types.classOrFail
import org.jetbrains.kotlin.ir.types.classifierOrNull
Expand Down Expand Up @@ -235,16 +237,32 @@ private class DataFrameFileLowering(val context: IrPluginContext) : FileLowering
return true
}

@OptIn(UnsafeDuringIrConstructionAPI::class)
// org.jetbrains.kotlin.fir.backend.generators.CallAndReferenceGenerator#applyReceivers
override fun visitTypeOperator(expression: IrTypeOperatorCall): IrExpression {
if (isScope(expression.typeOperand)) {
return expression.replaceWithConstructorCall()
}
return super.visitTypeOperator(expression)
}

override fun visitErrorCallExpression(expression: IrErrorCallExpression): IrExpression {
val origin = (expression.type.classifierOrNull?.owner as? IrClass)?.origin ?: return expression
val fromPlugin = origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is DataFramePlugin
val scopeReference = expression.type.classFqName?.shortName()?.asString()?.startsWith("Scope") ?: false
if (!(fromPlugin || scopeReference)) {
if (!isScope(expression.type)) {
return expression
}
val constructor = expression.type.getClass()!!.constructors.toList().single()
val type = expression.type
return expression.replaceWithConstructorCall()
}

@OptIn(UnsafeDuringIrConstructionAPI::class)
private fun isScope(type: IrType): Boolean {
val origin = (type.classifierOrNull?.owner as? IrClass)?.origin ?: return false
val fromPlugin = origin is IrDeclarationOrigin.GeneratedByPlugin && origin.pluginKey is DataFramePlugin
val scopeReference = type.classFqName?.shortName()?.asString()?.startsWith("Scope") ?: false
return fromPlugin || scopeReference
}

@OptIn(UnsafeDuringIrConstructionAPI::class)
private fun IrExpression.replaceWithConstructorCall(): IrConstructorCallImpl {
val constructor = type.getClass()!!.constructors.toList().single()
return IrConstructorCallImpl(-1, -1, type, constructor.symbol, 0, 0, 0)
}
}
14 changes: 14 additions & 0 deletions plugins/kotlin-dataframe/testData/box/wrongReceiver.kt
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
import org.jetbrains.kotlinx.dataframe.api.*


@DataSchema
data class Record(val a: String, val b: Int)

fun box(): String {
val df = List(10) { Record(it.toString(), it) }.let { dataFrameOf(*it.toTypedArray()) }
val aggregate = df.pivot { b }.aggregate {
this.add("c") { 123 }.c
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

which part failed before and is now solved?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

execution of this expression this.add("c") { 123 }.c failed due to wrong receiver in generated bytecode. c extension property used this@aggregate instead of scope it comes from

}
return "OK"
}
Original file line number Diff line number Diff line change
Expand Up @@ -442,6 +442,12 @@ public void testUpdate() {
runTest("testData/box/update.kt");
}

@Test
@TestMetadata("wrongReceiver.kt")
public void testWrongReceiver() {
runTest("testData/box/wrongReceiver.kt");
}

@Nested
@TestMetadata("testData/box/colKinds")
@TestDataPath("$PROJECT_ROOT")
Expand Down
Loading