oc.compiler-plugin.0.8.0.source-code.IrDestructuring.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of compiler-plugin Show documentation
Adds some utilities to code, developed originally for Advent of Code solving
There is a newer version: 0.9.0
package com.sschr15.aoc.compiler.internal

import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
import org.jetbrains.kotlin.backend.common.lower.irIfThen
import org.jetbrains.kotlin.backend.common.lower.irThrow
import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
import org.jetbrains.kotlin.cli.jvm.compiler.report
import org.jetbrains.kotlin.config.CompilerConfiguration
import org.jetbrains.kotlin.ir.IrStatement
import org.jetbrains.kotlin.ir.builders.*
import org.jetbrains.kotlin.ir.declarations.IrDeclarationBase
import org.jetbrains.kotlin.ir.declarations.IrFunction
import org.jetbrains.kotlin.ir.declarations.IrValueDeclaration
import org.jetbrains.kotlin.ir.declarations.IrVariable
import org.jetbrains.kotlin.ir.expressions.IrSyntheticBody
import org.jetbrains.kotlin.ir.types.starProjectedType
import org.jetbrains.kotlin.ir.util.*
import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.name.SpecialNames

class IrDestructuringFinder(
    val context: IrPluginContext,
    val config: CompilerConfiguration,
) : IrElementTransformerVoid() {
    override fun visitDeclaration(declaration: IrDeclarationBase): IrStatement {
        if (declaration.annotations.any { it.isAnnotationWithEqualFqName(FqName("com.sschr15.aoc.annotations.SkipDestructuringChecks")) }) {
            return declaration
        }

        return super.visitDeclaration(declaration)
    }

    val extra = ExtraPluginStuff(context)
    val destructuringMap = config.getMap(DestructuringFinder)

    override fun visitFunction(declaration: IrFunction): IrStatement {
        if (declaration.annotations.any { it.isAnnotationWithEqualFqName(FqName("com.sschr15.aoc.annotations.SkipDestructuringChecks")) }) {
            return declaration
        }

        val destructs = destructuringMap[declaration.getPackageFragment().packageFqName.child(declaration.name)]
            ?: return super.visitFunction(declaration)

        var currentDestruct = 0
        if (declaration.body is IrSyntheticBody) return super.visitFunction(declaration) // synthetic bodies have no statements
        val statements = declaration.body?.statements?.toMutableList() ?: return super.visitFunction(declaration)

        fun runOnStatement(stat: IrValueDeclaration): IrStatement {
            val expected = destructs[currentDestruct++]
            return context.irBuiltIns.createIrBuilder(stat.symbol, stat.startOffset, stat.endOffset).irBlock {
                +irIfThen(
                    condition = irCall(context.irBuiltIns.andandSymbol).apply {
                        putValueArgument(0, irIs(irGet(stat), context.irBuiltIns.collectionClass.starProjectedType))
                        putValueArgument(1, irNotEquals(
                            irCall(extra.sizeFunctionSymbol).apply {
                                dispatchReceiver = irGet(stat)
                            },
                            irInt(expected)
                        ))
                    },
                    thenPart=irThrow(irCallConstructor(extra.illegalArgumentExceptionCtor, emptyList()).apply {
                        putValueArgument(0, irConcat().apply {
                            arguments.addAll(listOf(
                                irString("expected to destruct exactly $expected elements, but "),
                                irCall(extra.sizeFunctionSymbol).apply {
                                    dispatchReceiver = irGet(stat)
                                },
                                irString(" were in the collection")
                            ))
                        })
                    }),
                )
            }
        }

        //TODO: handle destructuring in parameters
//        statements.addAll(
//            0,
//            declaration.valueParameters
//                .filter { it.name == SpecialNames.DESTRUCT }
//                .map { stat -> runOnStatement(stat) }
//        )

        statements.transformFlat { stat ->
            if (stat !is IrVariable) return@transformFlat null
            if (stat.name != SpecialNames.DESTRUCT) return@transformFlat null
            if (currentDestruct >= destructs.size) {
                config.report(CompilerMessageSeverity.STRONG_WARNING, "Destructuring ${stat.dumpKotlinLike()}, but no destructuring count was provided")
                return@transformFlat null
            }

            listOf(
                stat,
                runOnStatement(stat),
            )
        }
        return super.visitFunction(declaration)
    }
}