[JS IR] Cleanup StageController & Lowerings magic a bit

2021-12-23 14:36:31 +03:00
parent 7341a17a5d
commit 93e7eb2a8e
16 changed files with 62 additions and 166 deletions
@@ -138,10 +138,9 @@ private class DeclarationContainerLoweringVisitor(

 fun BodyLoweringPass.runOnFilePostfix(
    irFile: IrFile,
-    withLocalDeclarations: Boolean = false,
-    allowDeclarationModification: Boolean = false
+    withLocalDeclarations: Boolean = false
 ) {
-    val visitor = BodyLoweringVisitor(this, withLocalDeclarations, allowDeclarationModification)
+    val visitor = BodyLoweringVisitor(this, withLocalDeclarations)
    for (declaration in ArrayList(irFile.declarations)) {
        try {
            declaration.accept(visitor, null)
@@ -158,11 +157,8 @@ fun BodyLoweringPass.runOnFilePostfix(
    }
 }

-fun BodyAndScriptBodyLoweringPass.runOnFilePostfix(
-    irFile: IrFile,
-    allowDeclarationModification: Boolean = false
-) {
-    val visitor = ScriptBodyLoweringVisitor(this, allowDeclarationModification)
+fun BodyAndScriptBodyLoweringPass.runOnFilePostfix(irFile: IrFile) {
+    val visitor = ScriptBodyLoweringVisitor(this)
    for (declaration in ArrayList(irFile.declarations)) {
        declaration.accept(visitor, null)
    }
@@ -171,7 +167,6 @@ fun BodyAndScriptBodyLoweringPass.runOnFilePostfix(
 private open class BodyLoweringVisitor(
    private val loweringPass: BodyLoweringPass,
    private val withLocalDeclarations: Boolean,
-    private val allowDeclarationModification: Boolean,
 ) : IrElementVisitor<Unit, IrDeclaration?> {
    override fun visitElement(element: IrElement, data: IrDeclaration?) {
        element.acceptChildren(this, data)
@@ -191,13 +186,7 @@ private open class BodyLoweringVisitor(
        if (withLocalDeclarations) body.acceptChildren(this, null)
        val stageController = data!!.factory.stageController
        stageController.restrictTo(data) {
-            if (allowDeclarationModification) {
-                loweringPass.lower(body, data)
-            } else {
-                stageController.bodyLowering {
-                    loweringPass.lower(body, data)
-                }
-            }
+            loweringPass.lower(body, data)
        }
    }

@@ -208,25 +197,15 @@ private open class BodyLoweringVisitor(
 }

 private class ScriptBodyLoweringVisitor(
-    private val loweringPass: BodyAndScriptBodyLoweringPass,
-    private val allowDeclarationModification: Boolean
-) : BodyLoweringVisitor(loweringPass, false, allowDeclarationModification) {
+    private val loweringPass: BodyAndScriptBodyLoweringPass
+) : BodyLoweringVisitor(loweringPass, false) {

    override fun visitClass(declaration: IrClass, data: IrDeclaration?) {
        declaration.thisReceiver?.accept(this, declaration)
        declaration.typeParameters.forEach { it.accept(this, declaration) }
        ArrayList(declaration.declarations).forEach { it.accept(this, declaration) }
        if (declaration.origin == IrDeclarationOrigin.SCRIPT_CLASS) {
-            val stageController = declaration.factory.stageController
-            stageController.restrictTo(declaration) {
-                if (allowDeclarationModification) {
-                    loweringPass.lowerScriptBody(declaration, declaration)
-                } else {
-                    stageController.bodyLowering {
-                        loweringPass.lowerScriptBody(declaration, declaration)
-                    }
-                }
-            }
+            loweringPass.lowerScriptBody(declaration, declaration)
        }
    }
 }
@@ -267,13 +246,6 @@ interface DeclarationTransformer : FileLoweringPass {
            element.acceptChildrenVoid(this)
        }

-        override fun visitBody(body: IrBody) {
-            if (transformer.withLocalDeclarations) {
-                super.visitBody(body)
-            }
-            // else stop
-        }
-
        override fun visitFunction(declaration: IrFunction) {
            declaration.acceptChildrenVoid(this)

@@ -322,7 +294,9 @@ interface DeclarationTransformer : FileLoweringPass {
            declaration.typeParameters.forEach { it.accept(this, null) }
            ArrayList(declaration.declarations).forEach { it.accept(this, null) }

-            declaration.declarations.transformFlat(transformer::transformFlatRestricted)
+            declaration.declarations.transformFlat {
+                transformer.transformFlatRestricted(it)
+            }
        }

        override fun visitScript(declaration: IrScript) {
@@ -356,9 +356,7 @@ fun IrType.remapTypeParameters(

 /* Copied from K/N */
 fun IrDeclarationContainer.addChild(declaration: IrDeclaration) {
-    declaration.factory.stageController.unrestrictDeclarationListsAccess {
-        this.declarations += declaration
-    }
+    this.declarations += declaration
    declaration.setDeclarationsParent(this)
 }

@@ -33,9 +33,7 @@ class InnerClassesLowering(val context: BackendContext, private val innerClasses

    override fun transformFlat(declaration: IrDeclaration): List<IrDeclaration>? {
        if (declaration is IrClass && declaration.isInner) {
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
-                declaration.declarations += innerClassesSupport.getOuterThisField(declaration)
-            }
+            declaration.declarations += innerClassesSupport.getOuterThisField(declaration)
        } else if (declaration is IrConstructor) {
            val irClass = declaration.parentAsClass
            if (!irClass.isInner) return null
@@ -20,7 +20,7 @@ open class LocalClassPopupLowering(
    val recordExtractedLocalClasses: BackendContext.(IrClass) -> Unit = {},
 ) : BodyLoweringPass {
    override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, withLocalDeclarations = true, allowDeclarationModification = true)
+        runOnFilePostfix(irFile, withLocalDeclarations = true)
    }

    private data class ExtractedLocalClass(
@@ -85,7 +85,7 @@ class LocalDeclarationsLowering(
    BodyLoweringPass {

    override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, allowDeclarationModification = true)
+        runOnFilePostfix(irFile)
    }

    object DECLARATION_ORIGIN_FIELD_FOR_CAPTURED_VALUE :
@@ -34,7 +34,7 @@ import org.jetbrains.kotlin.ir.visitors.*

 class LocalClassesInInlineLambdasLowering(val context: CommonBackendContext) : BodyLoweringPass {
    override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, allowDeclarationModification = true)
+        runOnFilePostfix(irFile)
    }

    override fun lower(irBody: IrBody, container: IrDeclaration) {
@@ -84,7 +84,7 @@ class LocalClassesInInlineLambdasLowering(val context: CommonBackendContext) : B

 class LocalClassesInInlineFunctionsLowering(val context: CommonBackendContext) : BodyLoweringPass {
    override fun lower(irFile: IrFile) {
-        runOnFilePostfix(irFile, allowDeclarationModification = true)
+        runOnFilePostfix(irFile)
    }

    override fun lower(irBody: IrBody, container: IrDeclaration) {
@@ -33,18 +33,11 @@ fun eliminateDeadDeclarations(
    removeUnusedAssociatedObjects: Boolean = true,
 ) {

-    val allRoots = context.irFactory.stageController.withInitialIr { buildRoots(modules, context) }
+    val allRoots = buildRoots(modules, context)

    val usefulDeclarations = usefulDeclarations(allRoots, context, removeUnusedAssociatedObjects)

-    context.irFactory.stageController.unrestrictDeclarationListsAccess {
-        processUselessDeclarations(
-            modules,
-            usefulDeclarations,
-            context,
-            removeUnusedAssociatedObjects,
-        )
-    }
+    processUselessDeclarations(modules, usefulDeclarations, context, removeUnusedAssociatedObjects)
 }

 private fun IrField.isConstant(): Boolean {
@@ -276,32 +269,30 @@ fun usefulDeclarations(
    }

    // use withInitialIr to avoid ConcurrentModificationException in dce-driven lowering when adding roots' nested declarations (members)
-    context.irFactory.stageController.withInitialIr {
-        // Add roots
-        roots.forEach {
-            it.enqueue(null, null, altFromFqn = "<ROOT>")
-        }
+    // Add roots
+    roots.forEach {
+        it.enqueue(null, null, altFromFqn = "<ROOT>")
+    }

-        // Add roots' nested declarations
-        roots.forEach {
-            it.acceptVoid(
-                object : IrElementVisitorVoid {
-                    override fun visitElement(element: IrElement) {
-                        element.acceptChildrenVoid(this)
-                    }
-
-                    override fun visitBody(body: IrBody) {
-                        // Skip
-                    }
-
-                    override fun visitDeclaration(declaration: IrDeclarationBase) {
-                        if (declaration !== it) declaration.enqueue(it, "roots' nested declaration")
-
-                        super.visitDeclaration(declaration)
-                    }
+    // Add roots' nested declarations
+    roots.forEach {
+        it.acceptVoid(
+            object : IrElementVisitorVoid {
+                override fun visitElement(element: IrElement) {
+                    element.acceptChildrenVoid(this)
                }
-            )
-        }
+
+                override fun visitBody(body: IrBody) {
+                    // Skip
+                }
+
+                override fun visitDeclaration(declaration: IrDeclarationBase) {
+                    if (declaration !== it) declaration.enqueue(it, "roots' nested declaration")
+
+                    super.visitDeclaration(declaration)
+                }
+            }
+        )
    }

    val toStringMethod =
@@ -5,8 +5,6 @@

 package org.jetbrains.kotlin.ir.backend.js

-import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
-import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.StageController
 import org.jetbrains.kotlin.ir.util.IdSignature
@@ -20,8 +18,6 @@ class WholeWorldStageController : StageController() {
    override var currentDeclaration: IrDeclaration? = null
    private var index: Int = 0

-    var declarationListsRestricted: Boolean = false
-
    override fun <T> restrictTo(declaration: IrDeclaration, fn: () -> T): T {
        val previousCurrentDeclaration = currentDeclaration
        val previousIndex = index
@@ -37,56 +33,20 @@ class WholeWorldStageController : StageController() {
        }
    }

-    override fun createSignature(parentSignature: IdSignature): IdSignature {
-        return IdSignature.LoweredDeclarationSignature(parentSignature, currentStage, index++)
-    }
-
-    override fun <T> withStage(stage: Int, fn: () -> T): T {
+    override fun <T> withInitialIr(block: () -> T): T {
        val oldStage = currentStage
-        currentStage = stage
+        currentStage = 0
        val oldCurrentDeclaration = currentDeclaration
        currentDeclaration = null
        return try {
-            fn()
+            block()
        } finally {
            currentStage = oldStage
            currentDeclaration = oldCurrentDeclaration
        }
    }

-    override fun <T> withInitialIr(block: () -> T): T {
-        return withStage(0) {
-            declarationRestriction(true, block)
-        }
-    }
-
-    override fun <T> bodyLowering(fn: () -> T): T {
-        val wasRestricted = declarationListsRestricted
-        declarationListsRestricted = true
-
-        return try {
-            fn()
-        } finally {
-            declarationListsRestricted = wasRestricted
-        }
-    }
-
-    override fun <T> unrestrictDeclarationListsAccess(fn: () -> T): T {
-        return declarationRestriction(false, fn)
-    }
-
-    private fun <T> declarationRestriction(value: Boolean, fn: () -> T): T{
-        val wasRestricted = declarationListsRestricted
-        declarationListsRestricted = value
-
-        return try {
-            fn()
-        } finally {
-            declarationListsRestricted = wasRestricted
-        }
-    }
-
-    override fun canAccessDeclarationsOf(irClass: IrClass): Boolean {
-        return !declarationListsRestricted || irClass.visibility == DescriptorVisibilities.LOCAL/* && irClass !in context.extractedLocalClasses*/
+    override fun createSignature(parentSignature: IdSignature): IdSignature {
+        return IdSignature.LoweredDeclarationSignature(parentSignature, currentStage, index++)
    }
 }
@@ -72,9 +72,7 @@ abstract class AbstractBlockDecomposerLowering(
                    val lastStatement = newBody.statements.last()
                    val actualParent = if (newBody.statements.size > 1 || lastStatement !is IrReturn || lastStatement.value != expression) {
                        expression = JsIrBuilder.buildCall(initFunction.symbol, expression.type)
-                        context.irFactory.stageController.unrestrictDeclarationListsAccess {
-                            (container.parent as IrDeclarationContainer).declarations += initFunction
-                        }
+                        (container.parent as IrDeclarationContainer).declarations += initFunction
                        initFunction
                    } else {
                        container
@@ -364,10 +364,8 @@ class EnumClassCreateInitializerLowering(val context: JsCommonBackendContext) :

            // TODO Why not move to upper level?
            // TODO Also doesn't fit the transformFlat-ish API
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
-                declaration.declarations += entryInstancesInitializedVar
-                declaration.declarations += initEntryInstancesFun
-            }
+            declaration.declarations += entryInstancesInitializedVar
+            declaration.declarations += initEntryInstancesFun

            return null
        }
@@ -428,9 +426,7 @@ class EnumEntryCreateGetInstancesFunsLowering(val context: JsCommonBackendContex

                // TODO prettify
                entryGetInstanceFun.parent = irClass.parent
-                context.irFactory.stageController.unrestrictDeclarationListsAccess {
-                    (irClass.parent as IrDeclarationContainer).declarations += entryGetInstanceFun
-                }
+                (irClass.parent as IrDeclarationContainer).declarations += entryGetInstanceFun

                return listOf(declaration) // TODO not null?
            }
@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.utils.getJsNameOrKotlinName
-import org.jetbrains.kotlin.ir.backend.js.utils.hasStableJsName
 import org.jetbrains.kotlin.ir.backend.js.utils.realOverrideTarget
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
@@ -10,7 +10,10 @@ import org.jetbrains.kotlin.backend.common.getOrPut
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.builders.declarations.buildField
-import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrEnumEntry
+import org.jetbrains.kotlin.ir.declarations.IrField
 import org.jetbrains.kotlin.ir.expressions.IrBody
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrGetEnumValue
@@ -51,10 +54,6 @@ class ExternalEnumUsagesLowering(val context: JsIrBackendContext) : BodyLowering
            isStatic = true
        }.also {
            it.parent = irClass
-
-            // TODO need a way to emerge local declarations from BodyLoweringPass
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
-                irClass.declarations += it
-            }
+            irClass.declarations += it
        }
 }
@@ -38,9 +38,7 @@ class JsAnnotationImplementationTransformer(val jsContext: JsIrBackendContext) :

    override fun visitClassNew(declaration: IrClass): IrStatement {
        if (declaration.isAnnotationClass) {
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
-                implementGeneratedFunctions(declaration, declaration)
-            }
+            implementGeneratedFunctions(declaration, declaration)
        }
        return super.visitClassNew(declaration)
    }
@@ -45,9 +45,7 @@ class JsSingleAbstractMethodLowering(context: CommonBackendContext) : SingleAbst

        for (wrapper in cachedImplementations.values + inlineCachedImplementations.values) {
            val parentClass = wrapper.parent as IrDeclarationContainer
-            context.irFactory.stageController.unrestrictDeclarationListsAccess {
-                parentClass.declarations += wrapper
-            }
+            parentClass.declarations += wrapper
        }
    }

@@ -43,13 +43,10 @@ class PrimaryConstructorLowering(val context: JsCommonBackendContext) : Declarat
    private val unitType = context.irBuiltIns.unitType

    private fun createPrimaryConstructor(irClass: IrClass): IrConstructor {
-        // TODO better API for declaration creation. This case doesn't fit the usual transformFlat-like API.
-        val declaration = context.irFactory.stageController.unrestrictDeclarationListsAccess {
-            irClass.addConstructor {
-                origin = SYNTHETIC_PRIMARY_CONSTRUCTOR
-                isPrimary = true
-                visibility = DescriptorVisibilities.PRIVATE
-            }
+        val declaration = irClass.addConstructor {
+            origin = SYNTHETIC_PRIMARY_CONSTRUCTOR
+            isPrimary = true
+            visibility = DescriptorVisibilities.PRIVATE
        }

        declaration.body = irClass.run {
@@ -8,19 +8,9 @@ package org.jetbrains.kotlin.ir.declarations
 import org.jetbrains.kotlin.ir.util.IdSignature

 open class StageController(open val currentStage: Int = 0) {
-    open fun <T> withStage(stage: Int, fn: () -> T): T = fn()
-
-    open val bodiesEnabled: Boolean get() = true
-
-    open fun <T> withInitialIr(block: () -> T): T = block()
-
    open fun <T> restrictTo(declaration: IrDeclaration, fn: () -> T): T = fn()

-    open fun <T> bodyLowering(fn: () -> T): T = fn()
-
-    open fun <T> unrestrictDeclarationListsAccess(fn: () -> T): T = fn()
-
-    open fun canAccessDeclarationsOf(irClass: IrClass): Boolean = true
+    open fun <T> withInitialIr(block: () -> T): T = block()

    // Used in JS IC. Declarations created during lowerings need meaningful signatures.
    open fun createSignature(parentSignature: IdSignature): IdSignature? = null