[fir2ir] Support code fragment transformation

2023-06-20 18:22:08 +09:00
parent 02af189066
commit 9446118d7f
11 changed files with 261 additions and 4 deletions
@@ -12,7 +12,10 @@ import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.config.JvmSerializeIrMode
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.fir.backend.Fir2IrComponents
+import org.jetbrains.kotlin.fir.backend.Fir2IrConversionScope
 import org.jetbrains.kotlin.fir.backend.Fir2IrExtensions
+import org.jetbrains.kotlin.fir.backend.InjectedValue
+import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
@@ -64,6 +67,10 @@ class JvmFir2IrExtensions(
        }
    }

+    override fun findInjectedValue(calleeReference: FirReference, conversionScope: Fir2IrConversionScope): InjectedValue? {
+        return null
+    }
+
    override val irNeedsDeserialization: Boolean =
        configuration.get(JVMConfigurationKeys.SERIALIZE_IR, JvmSerializeIrMode.NONE) != JvmSerializeIrMode.NONE

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2023 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.kotlin.fir.backend
+
+import org.jetbrains.kotlin.fir.declarations.FirCodeFragment
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataKey
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationDataRegistry
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
+import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.Name
+
+class CodeFragmentConversionData(
+    val classId: ClassId,
+    val methodName: Name,
+    val injectedValues: List<InjectedValue>
+)
+
+class InjectedValue(val symbol: FirBasedSymbol<*>, val typeRef: FirTypeRef, val isMutated: Boolean) {
+    val irParameterSymbol: IrValueParameterSymbol = IrValueParameterSymbolImpl()
+}
+
+private object CodeFragmentTowerDataContext : FirDeclarationDataKey()
+
+private var FirCodeFragment.conversionDataOpt: CodeFragmentConversionData? by FirDeclarationDataRegistry.data(CodeFragmentTowerDataContext)
+
+var FirCodeFragment.conversionData: CodeFragmentConversionData
+    get() = conversionDataOpt ?: error("Conversion data is not set")
+    set(newValue) {
+        conversionDataOpt = newValue
+    }
@@ -58,6 +58,8 @@ class Fir2IrClassifierStorage(

    private val enumEntryCache: MutableMap<FirEnumEntry, IrEnumEntry> = commonMemberStorage.enumEntryCache

+    private val codeFragmentCache: MutableMap<FirCodeFragment, IrClass> = mutableMapOf()
+
    private val fieldsForContextReceivers: MutableMap<IrClass, List<IrField>> = mutableMapOf()

    private val localStorage: Fir2IrLocalClassStorage = Fir2IrLocalClassStorage(
@@ -189,6 +191,10 @@ class Fir2IrClassifierStorage(
        return localStorage[lookupTag.toSymbol(session)!!.fir as FirClass]
    }

+    fun getCachedIrCodeFragment(codeFragment: FirCodeFragment): IrClass? {
+        return codeFragmentCache[codeFragment]
+    }
+
    private fun FirRegularClass.enumClassModality(): Modality {
        return when {
            declarations.any { it is FirCallableDeclaration && it.modality == Modality.ABSTRACT } -> {
@@ -407,6 +413,44 @@ class Fir2IrClassifierStorage(
        return irAnonymousObject
    }

+    fun registerCodeFragmentClass(codeFragment: FirCodeFragment, containingFile: IrFile): IrClass {
+        val conversionData = codeFragment.conversionData
+        val signature = signatureComposer.composeSignature(codeFragment)
+
+        val irClass = codeFragment.convertWithOffsets { startOffset, endOffset ->
+            declareIrClass(signature) { symbol ->
+                irFactory.createClass(
+                    startOffset,
+                    endOffset,
+                    IrDeclarationOrigin.DEFINED,
+                    conversionData.classId.shortClassName,
+                    DescriptorVisibilities.PUBLIC,
+                    symbol,
+                    ClassKind.CLASS,
+                    Modality.FINAL,
+                    isExternal = false,
+                    isCompanion = false,
+                    isInner = false,
+                    isData = false,
+                    isValue = false,
+                    isExpect = false,
+                    isFun = false
+                ).apply {
+                    metadata = FirMetadataSource.CodeFragment(codeFragment)
+                    parent = containingFile
+                    typeParameters = emptyList()
+                    thisReceiver = declareThisReceiverParameter(
+                        thisType = IrSimpleTypeImpl(symbol, false, emptyList(), emptyList()),
+                        thisOrigin = IrDeclarationOrigin.INSTANCE_RECEIVER
+                    )
+                    superTypes = listOf(irBuiltIns.anyType)
+                }
+            }
+        }
+        codeFragmentCache[codeFragment] = irClass
+        return irClass
+    }
+
    private fun getIrAnonymousObjectForEnumEntry(anonymousObject: FirAnonymousObject, name: Name, irParent: IrClass?): IrClass {
        localStorage[anonymousObject]?.let { return it }
        val irAnonymousObject = registerIrAnonymousObject(anonymousObject, Visibilities.Private, name, irParent)
@@ -15,22 +15,29 @@ import org.jetbrains.kotlin.builtins.DefaultBuiltIns
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.config.AnalysisFlags
 import org.jetbrains.kotlin.config.LanguageFeature
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.backend.generators.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
 import org.jetbrains.kotlin.fir.declarations.utils.isSynthetic
 import org.jetbrains.kotlin.fir.descriptors.FirModuleDescriptor
+import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.extensions.declarationGenerators
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.generatedMembers
 import org.jetbrains.kotlin.fir.extensions.generatedNestedClassifiers
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.symbols.Fir2IrConstructorSymbol
+import org.jetbrains.kotlin.fir.symbols.Fir2IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.lazyDeclarationResolver
 import org.jetbrains.kotlin.ir.PsiIrFileEntry
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
 import org.jetbrains.kotlin.ir.declarations.impl.IrModuleFragmentImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrDelegatingConstructorCallImpl
 import org.jetbrains.kotlin.ir.interpreter.IrInterpreter
 import org.jetbrains.kotlin.ir.interpreter.IrInterpreterConfiguration
 import org.jetbrains.kotlin.ir.interpreter.IrInterpreterEnvironment
@@ -38,7 +45,10 @@ import org.jetbrains.kotlin.ir.interpreter.checker.EvaluationMode
 import org.jetbrains.kotlin.ir.interpreter.transformer.transformConst
 import org.jetbrains.kotlin.ir.util.KotlinMangler
 import org.jetbrains.kotlin.ir.util.NaiveSourceBasedFileEntryImpl
+import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance

 class Fir2IrConverter(
    private val moduleDescriptor: FirModuleDescriptor,
@@ -146,9 +156,11 @@ class Fir2IrConverter(
            moduleFragment
        )
        declarationStorage.registerFile(file, irFile)
-        file.declarations.forEach {
-            if (it is FirRegularClass) {
-                registerClassAndNestedClasses(it, irFile)
+        for (declaration in file.declarations) {
+            when (declaration) {
+                is FirRegularClass -> registerClassAndNestedClasses(declaration, irFile)
+                is FirCodeFragment -> classifierStorage.registerCodeFragmentClass(declaration, irFile)
+                else -> {}
            }
        }
        moduleFragment.files += irFile
@@ -227,6 +239,92 @@ class Fir2IrConverter(
        return irClass
    }

+    private fun processCodeFragmentMembers(
+        codeFragment: FirCodeFragment,
+        irClass: IrClass = classifierStorage.getCachedIrCodeFragment(codeFragment)!!
+    ): IrClass {
+        val conversionData = codeFragment.conversionData
+
+        declarationStorage.enterScope(irClass)
+
+        val signature = irClass.symbol.signature!!
+
+        val irPrimaryConstructor = symbolTable.declareConstructor(signature, { Fir2IrConstructorSymbol(signature) }) { irSymbol ->
+            irFactory.createConstructor(
+                UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                IrDeclarationOrigin.DEFINED,
+                Name.special("<init>"),
+                irClass.visibility,
+                isInline = false,
+                isExpect = false,
+                irClass.defaultType,
+                irSymbol,
+                isExternal = false,
+                isPrimary = true
+            ).apply {
+                parent = irClass
+                body = irFactory.createBlockBody(UNDEFINED_OFFSET, UNDEFINED_OFFSET).apply {
+                    statements += IrDelegatingConstructorCallImpl.fromSymbolOwner(
+                        UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                        irBuiltIns.unitType,
+                        irBuiltIns.anyClass.owner.declarations.firstIsInstance<IrConstructor>().symbol
+                    )
+                }
+            }
+        }
+
+        val irFragmentFunction = symbolTable.declareSimpleFunction(signature, { Fir2IrSimpleFunctionSymbol(signature) }) { irSymbol ->
+            val lastStatement = codeFragment.block.statements.lastOrNull()
+            val returnType = (lastStatement as? FirExpression)?.typeRef?.toIrType(typeConverter) ?: irBuiltIns.unitType
+
+            irFactory.createSimpleFunction(
+                UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                IrDeclarationOrigin.DEFINED,
+                conversionData.methodName,
+                DescriptorVisibilities.PUBLIC,
+                isInline = false,
+                isExpect = false,
+                returnType,
+                Modality.FINAL,
+                irSymbol,
+                isTailrec = false,
+                isSuspend = false,
+                isOperator = false,
+                isInfix = false,
+                isExternal = false,
+                containerSource = null,
+                isFakeOverride = false
+            ).apply fragmentFunction@{
+                parent = irClass
+                valueParameters = conversionData.injectedValues.mapIndexed { index, injectedValue ->
+                    val isMutated = injectedValue.isMutated
+
+                    irFactory.createValueParameter(
+                        UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                        if (isMutated) IrDeclarationOrigin.SHARED_VARIABLE_IN_EVALUATOR_FRAGMENT else IrDeclarationOrigin.DEFINED,
+                        Name.identifier("p$index"),
+                        injectedValue.typeRef.toIrType(typeConverter),
+                        isAssignable = isMutated,
+                        injectedValue.irParameterSymbol,
+                        index,
+                        varargElementType = null,
+                        isCrossinline = false,
+                        isNoinline = false,
+                        isHidden = false
+                    ).apply {
+                        parent = this@fragmentFunction
+                    }
+                }
+            }
+        }
+
+        irClass.declarations.add(irPrimaryConstructor)
+        irClass.declarations.add(irFragmentFunction)
+
+        declarationStorage.leaveScope(irClass)
+        return irClass
+    }
+
    private fun bindFakeOverridesInFile(file: FirFile) {
        val irFile = declarationStorage.getIrFile(file)
        for (irDeclaration in irFile.declarations) {
@@ -393,6 +491,9 @@ class Fir2IrConverter(
                // DO NOTHING
                null
            }
+            is FirCodeFragment -> {
+                processCodeFragmentMembers(declaration)
+            }
            else -> {
                error("Unexpected member: ${declaration::class}")
            }
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.fir.backend

+import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrMemberWithContainerSource
 import org.jetbrains.kotlin.ir.util.SymbolTable
@@ -15,6 +16,7 @@ interface Fir2IrExtensions {
    fun generateOrGetFacadeClass(declaration: IrMemberWithContainerSource, components: Fir2IrComponents): IrClass?
    fun deserializeToplevelClass(irClass: IrClass, components: Fir2IrComponents): Boolean
    fun registerDeclarations(symbolTable: SymbolTable)
+    fun findInjectedValue(calleeReference: FirReference, conversionScope: Fir2IrConversionScope): InjectedValue?

    object Default : Fir2IrExtensions {
        override val irNeedsDeserialization: Boolean
@@ -23,5 +25,6 @@ interface Fir2IrExtensions {
        override fun generateOrGetFacadeClass(declaration: IrMemberWithContainerSource, components: Fir2IrComponents): IrClass? = null
        override fun deserializeToplevelClass(irClass: IrClass, components: Fir2IrComponents): Boolean = false
        override fun registerDeclarations(symbolTable: SymbolTable) {}
+        override fun findInjectedValue(calleeReference: FirReference, conversionScope: Fir2IrConversionScope) = null
    }
 }
@@ -58,6 +58,7 @@ import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.kotlin.utils.addToStdlib.runUnless

 class Fir2IrVisitor(
@@ -316,6 +317,20 @@ class Fir2IrVisitor(
        }
    }

+    override fun visitCodeFragment(codeFragment: FirCodeFragment, data: Any?): IrElement {
+        val irClass = classifierStorage.getCachedIrCodeFragment(codeFragment)!!
+        val irFunction = irClass.declarations.firstIsInstance<IrSimpleFunction>()
+
+        declarationStorage.enterScope(irFunction)
+        conversionScope.withParent(irFunction) {
+            val irBlock = codeFragment.block.convertToIrBlock(forceUnitType = false)
+            irFunction.body = irFactory.createExpressionBody(irBlock)
+        }
+        declarationStorage.leaveScope(irFunction)
+
+        return irFunction
+    }
+
    override fun visitAnonymousObjectExpression(anonymousObjectExpression: FirAnonymousObjectExpression, data: Any?): IrElement {
        return visitAnonymousObject(anonymousObjectExpression.anonymousObject, data)
    }
@@ -639,6 +654,9 @@ class Fir2IrVisitor(
        data: Any?
    ): IrElement = whileAnalysing(session, thisReceiverExpression) {
        val calleeReference = thisReceiverExpression.calleeReference
+
+        callGenerator.injectGetValueCall(thisReceiverExpression, calleeReference)?.let { return it }
+
        val boundSymbol = calleeReference.boundSymbol
        when (boundSymbol) {
            is FirClassSymbol -> {
@@ -46,4 +46,6 @@ sealed class FirMetadataSource : MetadataSource {
    }

    class Script(override val fir: FirScript) : FirMetadataSource(), MetadataSource.Script
+
+    class CodeFragment(override val fir: FirCodeFragment) : FirMetadataSource(), MetadataSource.CodeFragment
 }
@@ -9,6 +9,7 @@ import com.intellij.openapi.progress.ProcessCanceledException
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.config.LanguageFeature
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.backend.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.isMethodOfAny
@@ -341,6 +342,19 @@ class CallAndReferenceGenerator(
            .applyCallArguments((qualifiedAccess as? FirCall)?.takeIf { !noArguments })
    }

+    internal fun injectGetValueCall(element: FirElement, calleeReference: FirReference): IrExpression? {
+        val injectedValue = extensions.findInjectedValue(calleeReference, conversionScope)
+        if (injectedValue != null) {
+            return element.convertWithOffsets { startOffset, endOffset ->
+                val type = injectedValue.typeRef.toIrType()
+                val origin = calleeReference.statementOrigin()
+                IrGetValueImpl(startOffset, endOffset, type, injectedValue.irParameterSymbol, origin)
+            }
+        }
+
+        return null
+    }
+
    fun convertToIrCall(
        qualifiedAccess: FirQualifiedAccessExpression,
        typeRef: FirTypeRef,
@@ -350,6 +364,8 @@ class CallAndReferenceGenerator(
        noArguments: Boolean = false
    ): IrExpression {
        try {
+            injectGetValueCall(qualifiedAccess, qualifiedAccess.calleeReference)?.let { return it }
+
            val type = typeRef.toIrType()
            val samConstructorCall = qualifiedAccess.tryConvertToSamConstructorCall(type)
            if (samConstructorCall != null) return samConstructorCall
@@ -535,12 +551,27 @@ class CallAndReferenceGenerator(
            ?: error("No receiver for dynamic call")
    }

+    private fun injectSetValueCall(element: FirElement, calleeReference: FirReference, assignedValue: IrExpression): IrExpression? {
+        val injectedValue = extensions.findInjectedValue(calleeReference, conversionScope)
+        if (injectedValue != null) {
+            return element.convertWithOffsets { startOffset, endOffset ->
+                val type = irBuiltIns.unitType
+                val origin = calleeReference.statementOrigin()
+                IrSetValueImpl(startOffset, endOffset, type, injectedValue.irParameterSymbol, assignedValue, origin)
+            }
+        }
+
+        return null
+    }
+
    fun convertToIrSetCall(variableAssignment: FirVariableAssignment, explicitReceiverExpression: IrExpression?): IrExpression {
        try {
            val type = irBuiltIns.unitType
            val calleeReference = variableAssignment.calleeReference ?: error("Reference not resolvable")
            val assignedValue = visitor.convertToIrExpression(variableAssignment.rValue)

+            injectSetValueCall(variableAssignment, calleeReference, assignedValue)?.let { return it }
+
            val firSymbol = calleeReference.toResolvedBaseSymbol()
            val isDynamicAccess = firSymbol?.origin == FirDeclarationOrigin.DynamicScope

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.common.serialization.mangle.MangleConstant
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.backend.FirMangler
+import org.jetbrains.kotlin.fir.backend.conversionData
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyGetter
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertySetter
@@ -47,6 +48,10 @@ class FirBasedSignatureComposer(val mangler: FirMangler) {
        return composeSignatureImpl(declaration, containingClass = null, forceExpect = false)
    }

+    fun composeSignature(declaration: FirCodeFragment): IdSignature? {
+        return composeSignatureImpl(declaration, containingClass = null, forceExpect = false)
+    }
+
    fun composeTypeParameterSignature(
        index: Int,
        containerSignature: IdSignature?
@@ -174,6 +179,12 @@ class FirBasedSignatureComposer(val mangler: FirMangler) {
                    description = builder.description,
                )
            }
+            is FirCodeFragment -> {
+                val conversionData = declaration.conversionData
+                val packageFqName = conversionData.classId.packageFqName.asString()
+                val classFqName = conversionData.classId.relativeClassName.asString()
+                IdSignature.CommonSignature(packageFqName, classFqName, builder.hashId, builder.mask, description = null)
+            }
            else -> error("Unsupported FIR declaration in signature composer: ${declaration.render()}")
        }
    }
@@ -206,6 +217,9 @@ class FirBasedSignatureComposer(val mangler: FirMangler) {
        override fun visitScript(script: FirScript, data: Any?) {
        }

+        override fun visitCodeFragment(codeFragment: FirCodeFragment, data: Any?) {
+        }
+
        override fun visitTypeAlias(typeAlias: FirTypeAlias, data: Any?) {
            setExpected(typeAlias.isExpect)
        }
@@ -469,7 +469,7 @@ fun FirBasedSymbol<*>.getOwnerLookupTag(): ConeClassLikeLookupTag? {
        is FirBackingFieldSymbol -> fir.propertySymbol.getOwnerLookupTag()
        is FirClassLikeSymbol<*> -> getContainingClassLookupTag()
        is FirCallableSymbol<*> -> containingClassLookupTag()
-        is FirScriptSymbol -> null
+        is FirScriptSymbol, is FirCodeFragmentSymbol -> null
        else -> errorWithAttachment("Unsupported owner search for ${fir::class.java}") {
            withFirEntry("ownerDeclaration", fir)
        }
@@ -18,6 +18,7 @@ interface MetadataSource {
        var serializedIr: ByteArray?
    }
    interface Script : MetadataSource
+    interface CodeFragment : MetadataSource
    interface Function : MetadataSource
    interface Property : MetadataSource {
        val isConst: Boolean