org.jetbrains.kotlinx.dataframe.plugin.analyzeRefinedCallShape.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of compiler-plugin-all Show documentation

Data processing in Kotlin

There is a newer version: 0.16.0-dev-5328

/*
 * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

package org.jetbrains.kotlinx.dataframe.plugin

import org.jetbrains.kotlin.fir.expressions.FirExpression
import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
import org.jetbrains.kotlin.fir.types.ConeClassLikeType
import org.jetbrains.kotlin.fir.types.ConeKotlinType
import org.jetbrains.kotlin.fir.types.ConeTypeProjection
import org.jetbrains.kotlin.fir.types.classId
import org.jetbrains.kotlin.fir.types.resolvedType
import org.jetbrains.kotlin.name.ClassId
import org.jetbrains.kotlin.name.Name
import org.jetbrains.kotlinx.dataframe.plugin.extensions.KotlinTypeFacade
import org.jetbrains.kotlinx.dataframe.plugin.impl.Interpreter

internal inline fun  KotlinTypeFacade.analyzeRefinedCallShape(
    call: FirFunctionCall,
    expectedReturnType: ClassId,
    reporter: InterpretationErrorReporter
): CallResult? {
    val callReturnType = call.resolvedType
    if (callReturnType.classId != expectedReturnType) return null
    // rootMarker is expected to be a token generated by the plugin.
    // it's implied by "refined call"
    // thus ConeClassLikeType
    val rootMarkers = callReturnType.typeArguments.filterIsInstance()
    if (rootMarkers.size != callReturnType.typeArguments.size) return null

    val newSchema: T = call.interpreterName(session)?.let { name ->
        when (name) {
            else -> name.load>().let { processor ->
                val dataFrameSchema = interpret(call, processor, reporter = reporter)
                    .let {
                        val value = it?.value
                        if (value !is T) {
                            if (!reporter.errorReported) {
                                reporter.reportInterpretationError(call, "${processor::class} must return ${T::class}, but was $value")
                            }
                            return null
                        }
                        value
                    }
                dataFrameSchema
            }
        }
    } ?: return null

    return CallResult(rootMarkers, newSchema)
}

data class CallResult(val markers: List, val result: T)

class RefinedArguments(val refinedArguments: List) : List by refinedArguments

data class RefinedArgument(val name: Name, val expression: FirExpression) {

    override fun toString(): String {
        return "RefinedArgument(name=$name, expression=${expression})"
    }
}

data class SchemaProperty(
    val marker: ConeTypeProjection,
    val name: String,
    val dataRowReturnType: ConeKotlinType,
    val columnContainerReturnType: ConeKotlinType,
    val override: Boolean = false
)