ExceptionKiller: coursework of kotlin compiler course

plugin-annotations/src/org/itmo/my/pretty/plugin/GenerateOrNull.kt

@@ -0,0 +1,3 @@
+package org.itmo.my.pretty.plugin
+
+annotation class GenerateOrNull

plugin-annotations/src/org/itmo/my/pretty/plugin/GenerateOrThrows.kt

@@ -0,0 +1,3 @@
+package org.itmo.my.pretty.plugin
+
+annotation class GenerateOrThrow

src/ru/itmo/kotlin/plugin/SimplePluginRegistrar.kt

@@ -1,10 +1,13 @@
 package ru.itmo.kotlin.plugin
 
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
-import ru.itmo.kotlin.plugin.fir.SimpleClassGenerator
+import ru.itmo.kotlin.plugin.fir.GenerateOrThrowExtension
+import ru.itmo.kotlin.plugin.fir.GenerateOrNullFunctionExtension
 
 class SimplePluginRegistrar : FirExtensionRegistrar() {
     override fun ExtensionRegistrarContext.configurePlugin() {
-        +::SimpleClassGenerator
+//        +::SimpleClassGenerator
+        +::GenerateOrNullFunctionExtension
+        +::GenerateOrThrowExtension
     }
 }

src/ru/itmo/kotlin/plugin/fir/FunctionGenerationExtension.kt

@@ -0,0 +1,121 @@
+package ru.itmo.kotlin.plugin.fir
+
+import org.jetbrains.kotlin.fir.FirFunctionTarget
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.FirValueParameter
+import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
+import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameterCopy
+import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.fir.expressions.buildResolvedArgumentList
+import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
+import org.jetbrains.kotlin.fir.expressions.builder.buildPropertyAccessExpression
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
+import org.jetbrains.kotlin.fir.moduleData
+import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.Name
+
+abstract class FunctionGenerationExtension(session: FirSession): FirDeclarationGenerationExtension(session) {
+    abstract val declPredicate: DeclarationPredicate
+
+    abstract fun generateNewFunctionName(oldName: String): String
+
+    abstract fun generateBody(returnFunTarget: FirFunctionTarget, baseFunSymbol: FirNamedFunctionSymbol): FirBlock
+
+    abstract fun makeReturnType(baseFunSymbol: FirNamedFunctionSymbol): FirResolvedTypeRef
+
+    @OptIn(SymbolInternals::class)
+    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+        if (callableId !in generatedCallableIds) {
+            return emptyList()
+        }
+        val baseFunSymbol = generatedCallableIds[callableId]!!
+        val returnFunTarget = FirFunctionTarget(callableId.callableName.asString(), false)
+
+        val funBody = generateBody(returnFunTarget, baseFunSymbol)
+
+        val genFun = buildSimpleFunction {
+            baseFunSymbol.fir.valueParameters.forEach { original ->
+                valueParameters.add(
+                    buildValueParameterCopy(original) {}
+                )
+            }
+
+            resolvePhase = FirResolvePhase.BODY_RESOLVE
+            moduleData = session.moduleData
+            origin = GenerateOrNullFunctionExtension.ThrowsFunctionKey.origin
+            status = baseFunSymbol.resolvedStatus
+            returnTypeRef = makeReturnType(baseFunSymbol)
+            name = callableId.callableName
+            symbol = FirNamedFunctionSymbol(callableId)
+
+            body = funBody
+        }
+
+        returnFunTarget.bind(genFun)
+
+
+        return listOf(genFun.symbol)
+    }
+
+    private val predicateBaseProvider = session.predicateBasedProvider
+
+    private val matchedFunction: List<FirNamedFunctionSymbol> by lazy {
+        predicateBaseProvider.getSymbolsByPredicate(declPredicate).filterIsInstance<FirNamedFunctionSymbol>()
+    }
+
+    private val generatedCallableIds: Map<CallableId, FirNamedFunctionSymbol> by lazy {
+        matchedFunction.associateBy { symbol ->
+            val genName = generateNewFunctionName(symbol.callableId.callableName.asString())
+            symbol.callableId.copy(callableName = Name.identifier(genName))
+        }
+    }
+
+    override fun getTopLevelCallableIds(): Set<CallableId> {
+        return generatedCallableIds.keys
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(declPredicate)
+    }
+
+    @OptIn(SymbolInternals::class)
+    protected fun buildFunctionCallWithSameArguments(baseFunSymbol: FirNamedFunctionSymbol): FirFunctionCall {
+        return buildFunctionCall {
+            typeRef = buildResolvedTypeRef {
+                type = baseFunSymbol.resolvedReturnType
+            }
+            calleeReference = buildResolvedNamedReference {
+                name = baseFunSymbol.name
+                resolvedSymbol = baseFunSymbol
+            }
+
+            val argumentMap = baseFunSymbol.fir.valueParameters.associateBy { param ->
+                buildPropertyAccessExpression {
+                    typeRef = param.returnTypeRef
+                    calleeReference = buildResolvedNamedReference {
+                        name = param.name
+                        resolvedSymbol = param.symbol
+                    }
+                }
+            }.let {
+                val result = linkedMapOf<FirExpression, FirValueParameter>()
+                result.putAll(it)
+                result
+            }
+
+            argumentList = buildResolvedArgumentList(argumentMap)
+        }
+    }
+}

src/ru/itmo/kotlin/plugin/fir/GenerateOrNullFunctionExtension.kt

@@ -0,0 +1,87 @@
+package ru.itmo.kotlin.plugin.fir
+
+import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.declarations.FirPluginKey
+import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
+import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+import org.jetbrains.kotlin.fir.extensions.predicate.has
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
+import org.jetbrains.kotlin.fir.types.*
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.withNullability
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.types.ConstantValueKind
+
+class GenerateOrNullFunctionExtension(session: FirSession): FunctionGenerationExtension(session) {
+    companion object {
+        private const val GENERATED_SUFFIX = "OrNull"
+    }
+
+    object ThrowsFunctionKey : FirPluginKey()
+
+    override val declPredicate: DeclarationPredicate = has("org.itmo.my.pretty.plugin.GenerateOrNull".fqn())
+
+    override fun generateNewFunctionName(oldName: String) =
+        oldName + GENERATED_SUFFIX
+
+    @OptIn(SymbolInternals::class)
+    override fun generateBody(returnFunTarget: FirFunctionTarget, baseFunSymbol: FirNamedFunctionSymbol): FirBlock {
+        // { return baseFun(...args...) }
+        val returnOriginalBlock = buildSimpleBlock(session) {
+            statements.add(
+                buildReturnExpression {
+                    target = returnFunTarget
+                    result = buildFunctionCallWithSameArguments(baseFunSymbol)
+                }
+            )
+        }
+
+        // try { returnOriginalBlock } catch { return null }
+        val tryCatch = buildTryExpression {
+            tryBlock = returnOriginalBlock
+            typeRef = buildResolvedTypeRef {
+                type = session.builtinTypes.nothingType.type
+            }
+            catches.add(
+                buildCatch {
+                    parameter = buildValueParameter {
+                        moduleData = session.moduleData
+                        origin = ThrowsFunctionKey.origin
+                        returnTypeRef = buildResolvedTypeRef {
+                            type = session.builtinTypes.throwableType.type
+                        }
+                        name = Name.identifier("e")
+                        symbol = FirValueParameterSymbol(Name.identifier("e"))
+                        isCrossinline = false
+                        isNoinline = false
+                        isVararg = false
+                    }
+
+                    val catchBlock = buildSimpleBlock(session) {
+                        statements.add(
+                            buildReturnExpression {
+                                target = returnFunTarget
+                                result = buildConstExpression(null, ConstantValueKind.Null, null, setType = true)
+                            }
+                        )
+                    }
+                    block = catchBlock
+                }
+            )
+        }
+
+        return buildSimpleBlock(session) {
+            statements.add(tryCatch)
+        }
+    }
+
+    override fun makeReturnType(baseFunSymbol: FirNamedFunctionSymbol): FirResolvedTypeRef {
+        return buildResolvedTypeRef {
+            type = baseFunSymbol.resolvedReturnType.withNullability(ConeNullability.NULLABLE, session.typeContext)
+        }
+    }
+}

src/ru/itmo/kotlin/plugin/fir/GenerateOrThrowExtension.kt

@@ -0,0 +1,86 @@
+package ru.itmo.kotlin.plugin.fir
+
+import org.jetbrains.kotlin.fir.FirFunctionTarget
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.buildResolvedArgumentList
+import org.jetbrains.kotlin.fir.expressions.builder.buildElvisExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildFunctionCall
+import org.jetbrains.kotlin.fir.expressions.builder.buildReturnExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildThrowExpression
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+import org.jetbrains.kotlin.fir.extensions.predicate.has
+import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
+import org.jetbrains.kotlin.fir.symbols.impl.ConeClassLikeLookupTagImpl
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.types.ConeNullability
+import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
+import org.jetbrains.kotlin.fir.types.typeContext
+import org.jetbrains.kotlin.fir.types.withNullability
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+class GenerateOrThrowExtension(session: FirSession): FunctionGenerationExtension(session) {
+    companion object {
+        private const val GEN_SUFFIX = "OrThrow"
+    }
+
+    override val declPredicate: DeclarationPredicate = has("org.itmo.my.pretty.plugin.GenerateOrThrow".fqn())
+
+    override fun generateNewFunctionName(oldName: String): String {
+        return oldName + GEN_SUFFIX
+    }
+
+    override fun generateBody(returnFunTarget: FirFunctionTarget, baseFunSymbol: FirNamedFunctionSymbol): FirBlock {
+        val elvis =
+            buildElvisExpression {
+                lhs = buildFunctionCallWithSameArguments(baseFunSymbol)
+                rhs = buildThrowExpression {
+                    exception = buildFunctionCall {
+                        val exceptionType = ConeClassLikeTypeImpl(
+                            ConeClassLikeLookupTagImpl(
+                                ClassId(
+                                    FqName("java.lang"),
+                                    FqName("IllegalStateException"),
+                                    false
+                                )
+                            ),
+                            arrayOf(),
+                            false
+                        )
+
+                        typeRef = buildResolvedTypeRef {
+                            type = exceptionType
+                        }
+                        argumentList = buildResolvedArgumentList(linkedMapOf())
+                        calleeReference = buildResolvedNamedReference {
+                            name = Name.identifier("IllegalStateException")
+                            resolvedSymbol = getEmptyParamsConstructor(exceptionType, session) ?: throw IllegalStateException()
+                        }
+                    }
+                }
+            }
+        elvis.replaceTypeRef(
+            buildResolvedTypeRef {
+                type = baseFunSymbol.resolvedReturnType.withNullability(ConeNullability.NOT_NULL, session.typeContext)
+            }
+        )
+        return buildSimpleBlock(session) {
+            statements.add(
+                buildReturnExpression {
+                    target = returnFunTarget
+                    result = elvis
+                }
+            )
+        }
+    }
+
+    override fun makeReturnType(baseFunSymbol: FirNamedFunctionSymbol): FirResolvedTypeRef {
+        return buildResolvedTypeRef {
+            type = baseFunSymbol.resolvedReturnType.withNullability(ConeNullability.NOT_NULL, session.typeContext)
+        }
+    }
+}

src/ru/itmo/kotlin/plugin/fir/Utils.kt

@@ -0,0 +1,47 @@
+package ru.itmo.kotlin.plugin.fir
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.builder.FirBlockBuilder
+import org.jetbrains.kotlin.fir.expressions.builder.buildBlock
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.scope
+import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
+import org.jetbrains.kotlin.fir.symbols.impl.FirConstructorSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
+import org.jetbrains.kotlin.name.FqName
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.InvocationKind
+import kotlin.contracts.contract
+
+@OptIn(ExperimentalContracts::class)
+inline fun buildSimpleBlock(session: FirSession,init: FirBlockBuilder.() -> Unit = {}): FirBlock {
+    contract {
+        callsInPlace(init, InvocationKind.EXACTLY_ONCE)
+    }
+
+    return buildBlock(init).also {
+        it.replaceTypeRef(
+            buildResolvedTypeRef {
+                type = session.builtinTypes.nothingType.type
+            }
+        )
+    }
+}
+
+fun getEmptyParamsConstructor(type: ConeKotlinType, session: FirSession): FirConstructorSymbol? {
+    val scope = type.scope(session, ScopeSession(), FakeOverrideTypeCalculator.DoNothing) ?: return null
+    var constructor: FirConstructorSymbol? = null
+
+    scope.processDeclaredConstructors {
+        if (it.valueParameterSymbols.isEmpty()) {
+            constructor = it
+        }
+    }
+    return constructor
+}
+
+fun String.fqn(): FqName {
+    return FqName(this)
+}

src/ru/itmo/kotlin/plugin/ir/SimpleIrGenerationExtension.kt

@@ -7,7 +7,9 @@ import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 
 class SimpleIrGenerationExtension: IrGenerationExtension {
     override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
-        val transformers = listOf(SimpleIrBodyGenerator(pluginContext))
+        val transformers = listOf<SimpleIrBodyGenerator>(
+//            SimpleIrBodyGenerator(pluginContext)
+        )
         for (transformer in transformers) {
             moduleFragment.acceptChildrenVoid(transformer)
         }