org.jetbrains.kotlinx.jupyter.codegen.RenderersProcessorImpl.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of kotlin-jupyter-kernel Show documentation
Kotlin Jupyter kernel published as artifact
There is a newer version: 0.12.0-356
Show newest version
package org.jetbrains.kotlinx.jupyter.codegen

import org.jetbrains.kotlinx.jupyter.api.Code
import org.jetbrains.kotlinx.jupyter.api.DisplayResult
import org.jetbrains.kotlinx.jupyter.api.FieldValue
import org.jetbrains.kotlinx.jupyter.api.PrecompiledRendererTypeHandler
import org.jetbrains.kotlinx.jupyter.api.ProcessingPriority
import org.jetbrains.kotlinx.jupyter.api.RendererFieldHandler
import org.jetbrains.kotlinx.jupyter.api.RendererHandlerWithPriority
import org.jetbrains.kotlinx.jupyter.api.libraries.ExecutionHost
import org.jetbrains.kotlinx.jupyter.exceptions.LibraryProblemPart
import org.jetbrains.kotlinx.jupyter.exceptions.rethrowAsLibraryException
import org.jetbrains.kotlinx.jupyter.repl.ContextUpdater
import org.jetbrains.kotlinx.jupyter.util.PriorityList

class RenderersProcessorImpl(
    private val contextUpdater: ContextUpdater,
) : ResultsRenderersProcessor {
    private var counter = 0
    private val renderers = PriorityList()

    override tailrec fun renderResult(
        host: ExecutionHost,
        field: FieldValue,
    ): Any? {
        val value = field.value
        if (value is DisplayResult) {
            return value
        }
        val (handler, id) =
            renderers.firstOrNull { it.handler.acceptsField(field) }
                ?: return value
        return if (id == null) {
            val newField =
                rethrowAsLibraryException(LibraryProblemPart.RENDERERS) {
                    handler.execution.execute(host, field)
                }
            renderResult(host, newField)
        } else {
            val methodName = getMethodName(id)
            contextUpdater.update()
            val functionInfo = contextUpdater.context.allFunctions[methodName]!!
            val resultValue =
                rethrowAsLibraryException(LibraryProblemPart.RENDERERS) {
                    functionInfo.function.call(functionInfo.line, value)
                }
            renderResult(host, FieldValue(resultValue, null))
        }
    }

    override fun renderValue(
        host: ExecutionHost,
        value: Any?,
    ): Any? {
        return renderResult(host, FieldValue(value, null))
    }

    override fun register(renderer: RendererFieldHandler): Code? {
        return register(renderer, ProcessingPriority.DEFAULT)
    }

    override fun register(
        renderer: RendererFieldHandler,
        priority: Int,
    ): Code? {
        return register(renderer, true, priority)
    }

    override fun registerWithoutOptimizing(renderer: RendererFieldHandler) {
        registerWithoutOptimizing(renderer, ProcessingPriority.DEFAULT)
    }

    override fun registerWithoutOptimizing(
        renderer: RendererFieldHandler,
        priority: Int,
    ) {
        register(renderer, false, priority)
    }

    private fun register(
        renderer: RendererFieldHandler,
        doOptimization: Boolean,
        priority: Int,
    ): Code? {
        if (!doOptimization || renderer !is PrecompiledRendererTypeHandler || !renderer.mayBePrecompiled) {
            renderers.add(HandlerWithInfo(renderer, null), priority)
            return null
        }

        val id = counter++
        renderers.add(HandlerWithInfo(renderer, id), priority)
        val methodName = getMethodName(id)
        return renderer.precompile(methodName, "___value")
    }

    override fun registeredRenderers(): List {
        return renderers.elementsWithPriority().map {
            RendererHandlerWithPriority(it.first.handler, it.second)
        }
    }

    override fun unregister(renderer: RendererFieldHandler) {
        val rendererInfosToRemove =
            renderers.elements().filter {
                it.handler === renderer
            }
        for (rendererInfo in rendererInfosToRemove) {
            renderers.remove(rendererInfo)
        }
    }

    override fun unregisterAll() {
        renderers.clear()
    }

    private fun getMethodName(id: Int) = "___renderResult$id"

    private data class HandlerWithInfo(
        val handler: RendererFieldHandler,
        val id: Int?,
    )
}