[K/JS] Rework IR deserialization and lowering phases to consume less memory

2023-04-19 13:10:19 +00:00
parent 1ee2d0814c
commit 33c5068b79
134 changed files with 964 additions and 616 deletions
@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.utils.*
 import org.jetbrains.kotlin.ir.visitors.*
 import org.jetbrains.kotlin.name.Name

@@ -141,8 +142,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
            }.apply {
                parent = irFunction.parent
                createParameterDeclarations()
-                typeParameters = irFunction.typeParameters.map { typeParam ->
-                    typeParam.copyToWithoutSuperTypes(this).apply { superTypes += typeParam.superTypes }
+                typeParameters = irFunction.typeParameters.memoryOptimizedMap { typeParam ->
+                    typeParam.copyToWithoutSuperTypes(this).apply { superTypes = superTypes memoryOptimizedPlus typeParam.superTypes }
                }
            }

@@ -186,11 +187,11 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                coroutineClass.declarations += this
                coroutineConstructors += this

-                valueParameters = functionParameters.mapIndexed { index, parameter ->
+                valueParameters = functionParameters.memoryOptimizedMapIndexed { index, parameter ->
                    parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
                }
                val continuationParameter = coroutineBaseClassConstructor.valueParameters[0]
-                valueParameters += continuationParameter.copyTo(
+                valueParameters = valueParameters memoryOptimizedPlus continuationParameter.copyTo(
                    this, DECLARATION_ORIGIN_COROUTINE_IMPL,
                    index = valueParameters.size, type = continuationType
                )
@@ -235,18 +236,18 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                parent = coroutineClass
                coroutineClass.declarations += this

-                typeParameters = stateMachineFunction.typeParameters.map { parameter ->
+                typeParameters = stateMachineFunction.typeParameters.memoryOptimizedMap { parameter ->
                    parameter.copyToWithoutSuperTypes(this, origin = DECLARATION_ORIGIN_COROUTINE_IMPL)
-                        .apply { superTypes += parameter.superTypes }
+                        .apply { superTypes = superTypes memoryOptimizedPlus parameter.superTypes }
                }

-                valueParameters = stateMachineFunction.valueParameters.mapIndexed { index, parameter ->
+                valueParameters = stateMachineFunction.valueParameters.memoryOptimizedMapIndexed { index, parameter ->
                    parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index)
                }

                this.createDispatchReceiverParameter()

-                overriddenSymbols += stateMachineFunction.symbol
+                overriddenSymbols = overriddenSymbols memoryOptimizedPlus stateMachineFunction.symbol
            }

            buildStateMachine(function, irFunction, argumentToPropertiesMap)
@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.ir.types.isArray
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.filterIsInstanceAnd

 val ANNOTATION_IMPLEMENTATION = object : IrDeclarationOriginImpl("ANNOTATION_IMPLEMENTATION", isSynthetic = true) {}

@@ -156,7 +157,8 @@ abstract class AnnotationImplementationTransformer(val context: BackendContext,
        // (although annotations imported from Java do have)
        val props = declarations.filterIsInstance<IrProperty>()
        if (props.isNotEmpty()) return props
-        return declarations.filterIsInstance<IrSimpleFunction>().filter { it.origin == IrDeclarationOrigin.DEFAULT_PROPERTY_ACCESSOR }
+        return declarations
+            .filterIsInstanceAnd<IrSimpleFunction> { it.origin == IrDeclarationOrigin.DEFAULT_PROPERTY_ACCESSOR }
            .mapNotNull { it.correspondingPropertySymbol?.owner }
    }

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrErrorExpressionImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.makeNullable
 import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.utils.*
 import org.jetbrains.kotlin.name.Name

 abstract class DefaultArgumentFunctionFactory(open val context: CommonBackendContext) {
@@ -40,7 +41,7 @@ abstract class DefaultArgumentFunctionFactory(open val context: CommonBackendCon
    }

    protected fun IrFunction.copyValueParametersFrom(original: IrFunction, wrapWithNullable: Boolean = true) {
-        valueParameters = original.valueParameters.map {
+        valueParameters = original.valueParameters.memoryOptimizedMap {
            val newType = it.type.remapTypeParameters(original.classIfConstructor, classIfConstructor)
            val makeNullable = wrapWithNullable && it.defaultValue != null &&
                    (context.ir.unfoldInlineClassType(it.type) ?: it.type) !in context.irBuiltIns.primitiveIrTypes
@@ -121,17 +122,18 @@ abstract class DefaultArgumentFunctionFactory(open val context: CommonBackendCon
                    context.mapping.defaultArgumentsOriginalFunction[defaultsFunction] = declaration

                    if (forceSetOverrideSymbols) {
-                        (defaultsFunction as IrSimpleFunction).overriddenSymbols += declaration.overriddenSymbols.mapNotNull {
-                            generateDefaultsFunction(
-                                it.owner,
-                                skipInlineMethods,
-                                skipExternalMethods,
-                                forceSetOverrideSymbols,
-                                visibility,
-                                useConstructorMarker,
-                                it.owner.copyAnnotations(),
-                            )?.symbol as IrSimpleFunctionSymbol?
-                        }
+                        (defaultsFunction as IrSimpleFunction).overriddenSymbols =
+                            defaultsFunction.overriddenSymbols memoryOptimizedPlus declaration.overriddenSymbols.mapNotNull {
+                                generateDefaultsFunction(
+                                    it.owner,
+                                    skipInlineMethods,
+                                    skipExternalMethods,
+                                    forceSetOverrideSymbols,
+                                    visibility,
+                                    useConstructorMarker,
+                                    it.owner.copyAnnotations(),
+                                )?.symbol as IrSimpleFunctionSymbol?
+                            }
                    }
                }
        }
@@ -200,7 +202,7 @@ abstract class DefaultArgumentFunctionFactory(open val context: CommonBackendCon
            parent = declaration.parent
            generateDefaultArgumentStubFrom(declaration, useConstructorMarker)
            // TODO some annotations are needed (e.g. @JvmStatic), others need different values (e.g. @JvmName), the rest are redundant.
-            annotations += copiedAnnotations
+            annotations = annotations memoryOptimizedPlus copiedAnnotations
        }
    }
 }
@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.memoryOptimizedPlus

 // TODO: fix expect/actual default parameters

@@ -488,7 +489,7 @@ class DefaultParameterPatchOverridenSymbolsLowering(
    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
        if (declaration is IrSimpleFunction) {
            (context.mapping.defaultArgumentsOriginalFunction[declaration] as? IrSimpleFunction)?.run {
-                declaration.overriddenSymbols += overriddenSymbols.mapNotNull {
+                declaration.overriddenSymbols = declaration.overriddenSymbols memoryOptimizedPlus overriddenSymbols.mapNotNull {
                    (context.mapping.defaultArgumentsDispatchFunction[it.owner] as? IrSimpleFunction)?.symbol
                }
            }
@@ -35,6 +35,8 @@ import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.memoryOptimizedMap
+import org.jetbrains.kotlin.utils.memoryOptimizedPlus

 interface VisibilityPolicy {
    fun forClass(declaration: IrClass, inInlineFunctionScope: Boolean): DescriptorVisibility =
@@ -708,7 +710,7 @@ class LocalDeclarationsLowering(
            )
            // Type parameters of oldDeclaration may depend on captured type parameters, so deal with that after copying.
            newDeclaration.typeParameters.drop(newTypeParameters.size).forEach { tp ->
-                tp.superTypes = tp.superTypes.map { localFunctionContext.remapType(it) }
+                tp.superTypes = tp.superTypes.memoryOptimizedMap { localFunctionContext.remapType(it) }
            }

            newDeclaration.parent = ownerParent
@@ -721,7 +723,7 @@ class LocalDeclarationsLowering(
            }
            newDeclaration.copyAttributes(oldDeclaration)

-            newDeclaration.valueParameters += createTransformedValueParameters(
+            newDeclaration.valueParameters = newDeclaration.valueParameters memoryOptimizedPlus createTransformedValueParameters(
                capturedValues, localFunctionContext, oldDeclaration, newDeclaration,
                isExplicitLocalFunction = oldDeclaration.origin == IrDeclarationOrigin.LOCAL_FUNCTION
            )
@@ -831,7 +833,7 @@ class LocalDeclarationsLowering(
                throw AssertionError("Local class constructor can't have extension receiver: ${ir2string(oldDeclaration)}")
            }

-            newDeclaration.valueParameters += createTransformedValueParameters(
+            newDeclaration.valueParameters = newDeclaration.valueParameters memoryOptimizedPlus createTransformedValueParameters(
                capturedValues, localClassContext, oldDeclaration, newDeclaration
            )
            newDeclaration.recordTransformedValueParameters(constructorContext)
@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.ir.types.classOrNull
 import org.jetbrains.kotlin.ir.types.getClass
 import org.jetbrains.kotlin.ir.types.isArray
 import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.utils.memoryOptimizedMapNotNull

 class MethodsFromAnyGeneratorForLowerings(val context: BackendContext, val irClass: IrClass, val origin: IrDeclarationOrigin) {
    private fun IrClass.addSyntheticFunction(name: String, returnType: IrType) =
@@ -43,7 +44,7 @@ class MethodsFromAnyGeneratorForLowerings(val context: BackendContext, val irCla

    companion object {
        fun IrClass.collectOverridenSymbols(predicate: (IrFunction) -> Boolean): List<IrSimpleFunctionSymbol> =
-            superTypes.mapNotNull { it.getClass()?.functions?.singleOrNull(predicate)?.symbol }
+            superTypes.memoryOptimizedMapNotNull { it.getClass()?.functions?.singleOrNull(predicate)?.symbol }
    }
 }

@@ -29,6 +29,9 @@ import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.utils.findIsInstanceAnd
+import org.jetbrains.kotlin.utils.memoryOptimizedMap
+import org.jetbrains.kotlin.utils.memoryOptimizedPlus

 abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
    // SAM wrappers are cached, either in the file class (if it exists), or in a top-level enclosing class.
@@ -79,8 +82,7 @@ abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) :
    override fun lower(irFile: IrFile) {
        cachedImplementations.clear()
        inlineCachedImplementations.clear()
-        enclosingContainer = irFile.declarations.filterIsInstance<IrClass>().find { it.isFileClass }
-            ?: irFile
+        enclosingContainer = irFile.declarations.findIsInstanceAnd<IrClass> { it.isFileClass } ?: irFile
        irFile.transformChildrenVoid()

        for (wrapper in cachedImplementations.values + inlineCachedImplementations.values) {
@@ -177,7 +179,7 @@ abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) :
            setSourceRange(createFor)
        }.apply {
            createImplicitParameterDeclarationWithWrappedDescriptor()
-            superTypes = listOf(superType) + getAdditionalSupertypes(superType)
+            superTypes = listOf(superType) memoryOptimizedPlus getAdditionalSupertypes(superType)
            parent = enclosingContainer!!
        }

@@ -221,7 +223,7 @@ abstract class SingleAbstractMethodLowering(val context: CommonBackendContext) :
            overriddenSymbols = listOf(originalSuperMethod.symbol)
            dispatchReceiverParameter = subclass.thisReceiver!!.copyTo(this)
            extensionReceiverParameter = originalSuperMethod.extensionReceiverParameter?.copyTo(this)
-            valueParameters = originalSuperMethod.valueParameters.map { it.copyTo(this) }
+            valueParameters = originalSuperMethod.valueParameters.memoryOptimizedMap { it.copyTo(this) }
            body = context.createIrBuilder(symbol).irBlockBody {
                +irReturn(
                    irCall(
@@ -28,6 +28,8 @@ import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.memoryOptimizedMap
+import org.jetbrains.kotlin.utils.memoryOptimizedPlus

 /**
 * Replaces suspend functions with regular non-suspend functions with additional
@@ -120,12 +122,12 @@ private fun IrSimpleFunction.createSuspendFunctionStub(context: CommonBackendCon
        val substitutionMap = makeTypeParameterSubstitutionMap(this, function)
        function.copyReceiverParametersFrom(this, substitutionMap)

-        function.overriddenSymbols += overriddenSymbols.map {
+        function.overriddenSymbols = function.overriddenSymbols memoryOptimizedPlus overriddenSymbols.map {
            factory.stageController.restrictTo(it.owner) {
                it.owner.getOrCreateFunctionWithContinuationStub(context).symbol
            }
        }
-        function.valueParameters = valueParameters.map { it.copyTo(function) }
+        function.valueParameters = valueParameters.memoryOptimizedMap { it.copyTo(function) }

        val mapping = mutableMapOf<IrValueSymbol, IrValueSymbol>()
        valueParameters.forEach { mapping[it.symbol] = function.valueParameters[it.index].symbol }
@@ -10,7 +10,6 @@ import org.jetbrains.kotlin.ir.declarations.IrDeclarationParent
 import org.jetbrains.kotlin.ir.declarations.IrTypeParametersContainer
 import org.jetbrains.kotlin.ir.declarations.copyAttributes
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
-import org.jetbrains.kotlin.ir.expressions.IrLoop
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
 import org.jetbrains.kotlin.ir.expressions.impl.IrTypeOperatorCallImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassifierSymbol
@@ -21,6 +20,7 @@ import org.jetbrains.kotlin.ir.types.impl.buildSimpleType
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
+import org.jetbrains.kotlin.utils.memoryOptimizedMap

 internal class DeepCopyIrTreeWithSymbolsForInliner(
    val typeArguments: Map<IrTypeParameterSymbol, IrType?>?,
@@ -54,7 +54,7 @@ internal class DeepCopyIrTreeWithSymbolsForInliner(
            arguments: List<IrTypeArgument>,
            erasedParameters: MutableSet<IrTypeParameterSymbol>?
        ) =
-            arguments.map { argument ->
+            arguments.memoryOptimizedMap { argument ->
                (argument as? IrTypeProjection)?.let { proj ->
                    remapTypeAndOptionallyErase(proj.type, erasedParameters)?.let { newType ->
                        makeTypeProjection(newType, proj.variance)
@@ -112,7 +112,7 @@ internal class DeepCopyIrTreeWithSymbolsForInliner(
                kotlinType = null
                this.classifier = symbolRemapper.getReferencedClassifier(classifier)
                arguments = remapTypeArguments(type.arguments, erasedParameters)
-                annotations = type.annotations.map { it.transform(copier, null) as IrConstructorCall }
+                annotations = type.annotations.memoryOptimizedMap { it.transform(copier, null) as IrConstructorCall }
            }
        }
    }