diff --git a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/constructors.kt b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/constructors.kt index 992a67e1c..fbd77787b 100644 --- a/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/constructors.kt +++ b/core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/api/constructors.kt @@ -279,6 +279,8 @@ public inline fun dataFrameOf(vararg header: String, fill: (String) public fun dataFrameOf(header: Iterable): DataFrameBuilder = DataFrameBuilder(header.asList()) +@Refine +@Interpretable("DataFrameOf3") public fun dataFrameOf(vararg columns: Pair>): DataFrame<*> = columns.map { it.second.toColumn(it.first, Infer.Type) }.toDataFrame() diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt index 7bbb9da1d..c741de7fb 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/extensions/FunctionCallTransformer.kt @@ -190,7 +190,7 @@ class FunctionCallTransformer( val tokenFir = token.toClassSymbol(session)!!.fir tokenFir.callShapeData = CallShapeData.RefinedType(dataSchemaApis.map { it.scope.symbol }) - return buildLetCall(call, originalSymbol, dataSchemaApis, listOf(tokenFir)) + return buildScopeFunctionCall(call, originalSymbol, dataSchemaApis, listOf(tokenFir)) } } @@ -253,7 +253,7 @@ class FunctionCallTransformer( val keyToken = groupMarker.toClassSymbol(session)!!.fir keyToken.callShapeData = CallShapeData.RefinedType(groupApis.map { it.scope.symbol }) - return buildLetCall(call, originalSymbol, keyApis + groupApis, additionalDeclarations = listOf(groupToken, keyToken)) + return buildScopeFunctionCall(call, originalSymbol, keyApis + groupApis, additionalDeclarations = listOf(groupToken, keyToken)) } } @@ -305,18 +305,17 @@ class FunctionCallTransformer( private fun Name.asTokenName() = identifierOrNullIfSpecial?.titleCase() ?: DEFAULT_NAME @OptIn(SymbolInternals::class) - private fun buildLetCall( + private fun buildScopeFunctionCall( call: FirFunctionCall, originalSymbol: FirNamedFunctionSymbol, dataSchemaApis: List, additionalDeclarations: List ): FirFunctionCall { - val explicitReceiver = call.explicitReceiver ?: return call - val receiverType = explicitReceiver.resolvedType + val explicitReceiver = call.explicitReceiver + val receiverType = explicitReceiver?.resolvedType val returnType = call.resolvedType - val resolvedLet = findLet() - val parameter = resolvedLet.valueParameterSymbols[0] + val scopeFunction = if (explicitReceiver != null) findLet() else findRun() // original call is inserted later call.transformCalleeReference(object : FirTransformer() { @@ -350,20 +349,23 @@ class FunctionCallTransformer( returnTypeRef = buildResolvedTypeRef { type = returnType } - val itName = Name.identifier("it") - val parameterSymbol = FirValueParameterSymbol(itName) - valueParameters += buildValueParameter { - moduleData = session.moduleData - origin = FirDeclarationOrigin.Source - returnTypeRef = buildResolvedTypeRef { - type = receiverType + val parameterSymbol = receiverType?.let { + val itName = Name.identifier("it") + val parameterSymbol = FirValueParameterSymbol(itName) + valueParameters += buildValueParameter { + moduleData = session.moduleData + origin = FirDeclarationOrigin.Source + returnTypeRef = buildResolvedTypeRef { + type = receiverType + } + this.name = itName + this.symbol = parameterSymbol + containingFunctionSymbol = fSymbol + isCrossinline = false + isNoinline = false + isVararg = false } - this.name = itName - this.symbol = parameterSymbol - containingFunctionSymbol = fSymbol - isCrossinline = false - isNoinline = false - isVararg = false + parameterSymbol } body = buildBlock { this.coneTypeOrNull = returnType @@ -375,20 +377,23 @@ class FunctionCallTransformer( statements += additionalDeclarations statements += buildReturnExpression { - val itPropertyAccess = buildPropertyAccessExpression { - coneTypeOrNull = receiverType - calleeReference = buildResolvedNamedReference { - name = itName - resolvedSymbol = parameterSymbol + if (parameterSymbol != null) { + val itPropertyAccess = buildPropertyAccessExpression { + coneTypeOrNull = receiverType + calleeReference = buildResolvedNamedReference { + name = parameterSymbol.name + resolvedSymbol = parameterSymbol + } + } + if (callDispatchReceiver != null) { + call.replaceDispatchReceiver(itPropertyAccess) + } + call.replaceExplicitReceiver(itPropertyAccess) + if (callExtensionReceiver != null) { + call.replaceExtensionReceiver(itPropertyAccess) } } - if (callDispatchReceiver != null) { - call.replaceDispatchReceiver(itPropertyAccess) - } - call.replaceExplicitReceiver(itPropertyAccess) - if (callExtensionReceiver != null) { - call.replaceExtensionReceiver(itPropertyAccess) - } + result = call this.target = target } @@ -397,11 +402,19 @@ class FunctionCallTransformer( isLambda = true hasExplicitParameterList = false typeRef = buildResolvedTypeRef { - type = ConeClassLikeTypeImpl( - ConeClassLikeLookupTagImpl(ClassId(FqName("kotlin"), Name.identifier("Function1"))), - typeArguments = arrayOf(receiverType, returnType), - isNullable = false - ) + type = if (receiverType != null) { + ConeClassLikeTypeImpl( + ConeClassLikeLookupTagImpl(ClassId(FqName("kotlin"), Name.identifier("Function1"))), + typeArguments = arrayOf(receiverType, returnType), + isNullable = false + ) + } else { + ConeClassLikeTypeImpl( + ConeClassLikeLookupTagImpl(ClassId(FqName("kotlin"), Name.identifier("Function0"))), + typeArguments = arrayOf(returnType), + isNullable = false + ) + } } invocationKind = EventOccurrencesRange.EXACTLY_ONCE inlineStatus = InlineStatus.Inline @@ -413,11 +426,13 @@ class FunctionCallTransformer( val newCall1 = buildFunctionCall { source = call.source this.coneTypeOrNull = returnType - typeArguments += buildTypeProjectionWithVariance { - typeRef = buildResolvedTypeRef { - type = receiverType + if (receiverType != null) { + typeArguments += buildTypeProjectionWithVariance { + typeRef = buildResolvedTypeRef { + type = receiverType + } + variance = Variance.INVARIANT } - variance = Variance.INVARIANT } typeArguments += buildTypeProjectionWithVariance { @@ -429,11 +444,14 @@ class FunctionCallTransformer( dispatchReceiver = null this.explicitReceiver = callExplicitReceiver extensionReceiver = callExtensionReceiver ?: callDispatchReceiver - argumentList = buildResolvedArgumentList(original = null, linkedMapOf(argument to parameter.fir)) + argumentList = buildResolvedArgumentList( + original = null, + linkedMapOf(argument to scopeFunction.valueParameterSymbols[0].fir) + ) calleeReference = buildResolvedNamedReference { source = call.calleeReference.source - this.name = Name.identifier("let") - resolvedSymbol = resolvedLet + this.name = scopeFunction.name + resolvedSymbol = scopeFunction } } return newCall1 @@ -565,5 +583,9 @@ class FunctionCallTransformer( return session.symbolProvider.getTopLevelFunctionSymbols(FqName("kotlin"), Name.identifier("let")).single() } + private fun findRun(): FirFunctionSymbol<*> { + return session.symbolProvider.getTopLevelFunctionSymbols(FqName("kotlin"), Name.identifier("run")).single { it.typeParameterSymbols.size == 1 } + } + private fun String.titleCase() = replaceFirstChar { it.uppercaseChar() } } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dataFrameOf.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dataFrameOf.kt index bae8d7b6a..bc9270623 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dataFrameOf.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/dataFrameOf.kt @@ -1,9 +1,12 @@ @file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER") package org.jetbrains.kotlinx.dataframe.plugin.impl.api +import org.jetbrains.kotlin.fir.expressions.FirExpression +import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression import org.jetbrains.kotlin.fir.expressions.FirVarargArgumentsExpression import org.jetbrains.kotlin.fir.types.commonSuperTypeOrNull import org.jetbrains.kotlin.fir.types.resolvedType +import org.jetbrains.kotlin.fir.types.type import org.jetbrains.kotlin.fir.types.typeContext import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractSchemaModificationInterpreter @@ -36,3 +39,18 @@ class DataFrameBuilderInvoke0 : AbstractSchemaModificationInterpreter() { return PluginDataFrameSchema(columns) } } + +class DataFrameOf3 : AbstractSchemaModificationInterpreter() { + val Arguments.columns: List>> by arg() + + override fun Arguments.interpret(): PluginDataFrameSchema { + val res = columns.map { + val it = it.value + val name = (it.first as? FirLiteralExpression)?.value as? String + val type = (it.second as? FirExpression)?.resolvedType?.typeArguments?.getOrNull(0)?.type + if (name == null || type == null) return PluginDataFrameSchema(emptyList()) + simpleColumnOf(name, type) + } + return PluginDataFrameSchema(res) + } +} diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/stdlib.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/stdlib.kt new file mode 100644 index 000000000..e3c6bcc30 --- /dev/null +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/impl/api/stdlib.kt @@ -0,0 +1,13 @@ +package org.jetbrains.kotlinx.dataframe.plugin.impl.api + +import org.jetbrains.kotlinx.dataframe.plugin.impl.AbstractInterpreter +import org.jetbrains.kotlinx.dataframe.plugin.impl.Arguments +import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter + +class PairConstructor : AbstractInterpreter>() { + val Arguments.receiver: Any? by arg(lens = Interpreter.Id) + val Arguments.that: Any? by arg(lens = Interpreter.Id) + override fun Arguments.interpret(): Pair<*, *> { + return receiver to that + } +} diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt index 67d97c4f7..30de23def 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/interpret.kt @@ -136,7 +136,11 @@ fun KotlinTypeFacade.interpret( is FirCallableReferenceAccess -> { toKPropertyApproximation(it, session) } - + is FirFunctionCall -> { + it.loadInterpreter()?.let { processor -> + interpret(it, processor, emptyMap(), reporter) + } + } else -> null } 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 128b9fcb8..5a17575b4 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 @@ -60,13 +60,20 @@ import org.jetbrains.kotlin.fir.expressions.FirFunctionCall import org.jetbrains.kotlin.fir.expressions.FirGetClassCall import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier +import org.jetbrains.kotlin.fir.expressions.UnresolvedExpressionTypeAccess import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference +import org.jetbrains.kotlin.fir.references.resolved +import org.jetbrains.kotlin.fir.references.symbol +import org.jetbrains.kotlin.fir.references.toResolvedNamedFunctionSymbol import org.jetbrains.kotlin.fir.resolve.fqName import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol import org.jetbrains.kotlin.fir.types.classId import org.jetbrains.kotlin.fir.types.coneType +import org.jetbrains.kotlin.name.CallableId import org.jetbrains.kotlin.name.ClassId +import org.jetbrains.kotlin.name.FqName import org.jetbrains.kotlin.name.Name +import org.jetbrains.kotlin.name.StandardClassIds import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddDslStringInvoke import org.jetbrains.kotlinx.dataframe.plugin.impl.api.AddId import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Aggregate @@ -76,12 +83,14 @@ import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ColsOf1 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameBuilderInvoke0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf0 +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.DataFrameOf3 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FillNulls0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Flatten0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FlattenDefault import org.jetbrains.kotlinx.dataframe.plugin.impl.api.FrameCols0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.MapToFrame import org.jetbrains.kotlinx.dataframe.plugin.impl.api.Move0 +import org.jetbrains.kotlinx.dataframe.plugin.impl.api.PairConstructor import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ReadExcel import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrame import org.jetbrains.kotlinx.dataframe.plugin.impl.api.ToDataFrameColumn @@ -91,8 +100,16 @@ 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.Update0 import org.jetbrains.kotlinx.dataframe.plugin.impl.api.UpdateWith0 +import org.jetbrains.kotlinx.dataframe.plugin.utils.Names +@OptIn(UnresolvedExpressionTypeAccess::class) internal fun FirFunctionCall.loadInterpreter(session: FirSession): Interpreter<*>? { + if ( + calleeReference.toResolvedNamedFunctionSymbol()?.callableId == Names.TO && + coneTypeOrNull?.classId == Names.PAIR + ) { + return PairConstructor() + } val symbol = (calleeReference as? FirResolvedNamedReference)?.resolvedSymbol as? FirCallableSymbol ?: return null val argName = Name.identifier("interpreter") @@ -208,6 +225,7 @@ internal inline fun String.load(): T { "ToTop" -> ToTop() "Update0" -> Update0() "Aggregate" -> Aggregate() + "DataFrameOf3" -> DataFrameOf3() else -> error("$this") } as T } diff --git a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt index 7eb9d5971..9c766c21c 100644 --- a/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt +++ b/plugins/kotlin-dataframe/src/org/jetbrains/kotlinx/dataframe/plugin/utils/Names.kt @@ -5,6 +5,7 @@ package org.jetbrains.kotlinx.dataframe.plugin.utils +import org.jetbrains.kotlin.name.CallableId import org.jetbrains.kotlin.name.ClassId import org.jetbrains.kotlin.name.FqName import org.jetbrains.kotlin.name.Name @@ -50,6 +51,9 @@ object Names { val LOCAL_DATE_CLASS_ID = kotlinx.datetime.LocalDate::class.classId() val LOCAL_DATE_TIME_CLASS_ID = kotlinx.datetime.LocalDateTime::class.classId() val INSTANT_CLASS_ID = kotlinx.datetime.Instant::class.classId() + + val PAIR = ClassId(FqName("kotlin"), Name.identifier("Pair")) + val TO = CallableId(FqName("kotlin"), Name.identifier("to")) } private fun KClass<*>.classId(): ClassId { diff --git a/plugins/kotlin-dataframe/testData/box/dataFrameOf_to.kt b/plugins/kotlin-dataframe/testData/box/dataFrameOf_to.kt new file mode 100644 index 000000000..4d0954eb0 --- /dev/null +++ b/plugins/kotlin-dataframe/testData/box/dataFrameOf_to.kt @@ -0,0 +1,14 @@ +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( + "a" to listOf(1, 2), + "b" to listOf("str1", "str2"), + ) + val i: Int = df.a[0] + val str: String = df.b[0] + 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 2e5d5b5d2..605e49f49 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 @@ -82,6 +82,12 @@ public void testDataFrameOf() { runTest("testData/box/dataFrameOf.kt"); } + @Test + @TestMetadata("dataFrameOf_to.kt") + public void testDataFrameOf_to() { + runTest("testData/box/dataFrameOf_to.kt"); + } + @Test @TestMetadata("dataFrameOf_vararg.kt") public void testDataFrameOf_vararg() {