Refactor JS optimizer to use RecursiveJsVisitor where possible and introduce more self-descriptive names in TemporaryAssignmentElimination. Add TODOs and comments for code pieces with non-obvious meaning.

js/js.inliner/src/org/jetbrains/kotlin/js/inline/clean/DeadCodeElimination.kt

@@ -86,6 +86,8 @@ internal class DeadCodeElimination(private val root: JsStatement) {
             EliminationVisitor().accept(x.initVars)
             EliminationVisitor().accept(x.incrementExpression)
 
+            // TODO: We may also check if condition is `true` constant or missing, which means this loop is infinite, i.e.
+            // code after this loop can be safely deleted when loop does not contain break statements.
             visitLoop(x.body) { true }
         }
 

js/js.inliner/src/org/jetbrains/kotlin/js/inline/clean/FunctionPostProcessor.kt

@@ -19,16 +19,19 @@ package org.jetbrains.kotlin.js.inline.clean
 import com.google.dart.compiler.backend.js.ast.JsBlock
 
 class FunctionPostProcessor(private val root: JsBlock) {
+    val optimizations = listOf(
+        { TemporaryAssignmentElimination(root).apply() },
+        { RedundantLabelRemoval(root).apply() },
+        { TemporaryVariableElimination(root).apply() },
+        { IfStatementReduction(root).apply() },
+        { DeadCodeElimination(root).apply() },
+        { RedundantVariableDeclarationElimination(root).apply() }
+    )
+    // TODO: reduce to A || B, A && B if possible
+
     fun apply() {
         do {
-            var hasChanges = false
+            val hasChanges = optimizations.fold(false) { existing, f -> existing or f() }
-            hasChanges = hasChanges or TemporaryAssignmentElimination(root).apply()
-            hasChanges = hasChanges or RedundantLabelRemoval(root).apply()
-            hasChanges = hasChanges or TemporaryVariableElimination(root).apply()
-            hasChanges = hasChanges or IfStatementReduction(root).apply()
-            // TODO: reduce to A || B, A && B if possible
-            hasChanges = hasChanges or DeadCodeElimination(root).apply()
-            hasChanges = hasChanges or RedundantVariableDeclarationElimination(root).apply()
         } while (hasChanges)
     }
 }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/clean/RedundantLabelRemoval.kt

@@ -100,7 +100,7 @@ internal class RedundantLabelRemoval(private val root: JsStatement) {
             }
         is JsIf -> {
             val thenRemoved = perform(statement.thenStatement, name) == null
-            val elseRemoved = statement.elseStatement?.let { perform(it, name) } == null && statement.elseStatement != null
+            val elseRemoved = statement.elseStatement?.let { perform(it, name) == null } ?: false
             when {
                 thenRemoved && elseRemoved -> null
                 elseRemoved -> {
@@ -119,6 +119,7 @@ internal class RedundantLabelRemoval(private val root: JsStatement) {
             }
         }
         is JsTry -> {
+            // TODO: optimize finally and catch blocks
             val finallyBlock = statement.finallyBlock
             val result = perform(statement.tryBlock, name)
             if (result != null) {

js/js.inliner/src/org/jetbrains/kotlin/js/inline/clean/TemporaryAssignmentElimination.kt

@@ -43,75 +43,78 @@ internal class TemporaryAssignmentElimination(private val root: JsBlock) {
     private fun analyze() {
         namesToProcess.addAll(collectDefinedNames(root))
 
-        object : JsVisitorWithContextImpl() {
+        object : RecursiveJsVisitor() {
-            override fun visit(x: JsReturn, ctx: JsContext<*>): Boolean {
+            override fun visitReturn(x: JsReturn) {
                 val returnExpr = x.expression
                 if (returnExpr != null) {
                     tryRecord(returnExpr, Usage.Return(x))
                 }
-                return super.visit(x, ctx)
+                super.visitReturn(x)
             }
 
-            override fun visit(x: JsExpressionStatement, ctx: JsContext<*>): Boolean {
+            override fun visitExpressionStatement(x: JsExpressionStatement) {
-                val assignment = JsAstUtils.decomposeAssignmentToVariable(x.expression)
+                val variableAssignment = JsAstUtils.decomposeAssignmentToVariable(x.expression)
-                if (assignment != null) {
+                if (variableAssignment != null) {
-                    val (name, value) = assignment
+                    val (name, value) = variableAssignment
-                    val usage = Usage.Assignment(x, name)
+                    val usage = Usage.VariableAssignment(x, name)
                     if (x.synthetic) {
                         syntheticNames += name
                     }
                     tryRecord(value, usage)
                     accept(value)
-                    return false
+                    // Don't visit LHS, since it's already treated as a temporary variable
+                    return
                 }
 
-                val mutation = JsAstUtils.decomposeAssignment(x.expression)
+                val propertyMutation = JsAstUtils.decomposeAssignment(x.expression)
-                if (mutation != null) {
+                if (propertyMutation != null) {
-                    val (target, value) = mutation
+                    val (target, value) = propertyMutation
                     if (!target.canHaveSideEffect()) {
-                        val usage = Usage.Mutation(x, target)
+                        val usage = Usage.PropertyMutation(x, target)
                         tryRecord(value, usage)
                         accept(value)
-                        return false
+                        return
                     }
                 }
 
-                return super.visit(x, ctx)
+                return super.visitExpressionStatement(x)
             }
 
-            override fun visit(x: JsVars, ctx: JsContext<*>): Boolean {
+            override fun visitVars(x: JsVars) {
+                // TODO: generalize for multliple declarations per one statement
                 if (x.vars.size == 1) {
                     val declaration = x.vars[0]
                     val initExpression = declaration.initExpression
                     if (initExpression != null) {
-                        tryRecord(initExpression, Usage.Declaration(x, declaration.name))
+                        tryRecord(initExpression, Usage.VariableDeclaration(x, declaration.name))
                     }
+                    // TODO: generalize for case when single JsVar is synthetic
                     if (x.synthetic) {
                         syntheticNames += declaration.name
                     }
                 }
 
                 x.vars.forEach { v -> v?.initExpression?.let { accept(it) } }
-                return false
             }
 
-            override fun visit(x: JsNameRef, ctx: JsContext<*>): Boolean {
+            override fun visitNameRef(nameRef: JsNameRef) {
-                val name = x.name
+                val name = nameRef.name
-                if (name != null && x.qualifier == null) {
+                if (name != null && nameRef.qualifier == null) {
                     use(name)
-                    return false
+                    return
                 }
-                return super.visit(x, ctx)
+                return super.visitNameRef(nameRef)
             }
 
-            override fun visit(x: JsFunction, ctx: JsContext<*>): Boolean {
+            override fun visitFunction(x: JsFunction) {
+                // Don't visit function body, but mark its free variables as used two times, so that further stages won't treat
+                // these variables as temporary.
                 x.collectFreeVariables().forEach { use(it); use(it) }
-                return false
             }
 
-            override fun visit(x: JsBreak, ctx: JsContext<*>) = false
+            override fun visitBreak(x: JsBreak) { }
 
-            override fun visit(x: JsContinue, ctx: JsContext<*>) = false
+            override fun visitContinue(x: JsContinue) { }
         }.accept(root)
 
         usages.keys.retainAll(syntheticNames)
@@ -119,11 +122,11 @@ internal class TemporaryAssignmentElimination(private val root: JsBlock) {
 
     private fun getUsage(name: JsName): Usage? {
         return mappedUsages.getOrPut(name) {
-            if (usageCount[name] ?: 0 != 1) return null
+            if (usageCount[name] != 1) return null
 
             val usage = usages[name]
             return when (usage) {
-                is Usage.Assignment -> {
+                is Usage.VariableAssignment -> {
                     val result = getUsage(usage.target)
                     if (result != null) {
                         result.statements.addAll(usage.statements)
@@ -133,7 +136,7 @@ internal class TemporaryAssignmentElimination(private val root: JsBlock) {
                         usage
                     }
                 }
-                is Usage.Declaration -> {
+                is Usage.VariableDeclaration -> {
                     val result = getUsage(usage.target)
                     if (result != null) {
                         result.statements.addAll(usage.statements)
@@ -151,22 +154,16 @@ internal class TemporaryAssignmentElimination(private val root: JsBlock) {
     private fun calculateDeclarations() {
         usages.keys.forEach { getUsage(it) }
 
-        object : JsVisitorWithContextImpl() {
+        object : RecursiveJsVisitor() {
-            override fun visit(x: JsExpressionStatement, ctx: JsContext<JsNode>): Boolean {
+            override fun visitExpressionStatement(x: JsExpressionStatement) {
-                if (x in statementsToRemove) {
-                    hasChanges = true
-                    ctx.removeMe()
-                    return false
-                }
-
                 val assignment = JsAstUtils.decomposeAssignmentToVariable(x.expression)
                 if (assignment != null) {
                     val usage = getUsage(assignment.first)
-                    if (usage is Usage.Declaration) {
+                    if (usage is Usage.VariableDeclaration) {
                         usage.count++
                     }
                 }
-                return super.visit(x, ctx)
+                super.visitExpressionStatement(x)
             }
         }.accept(root)
     }
@@ -187,13 +184,13 @@ internal class TemporaryAssignmentElimination(private val root: JsBlock) {
                     if (usage != null) {
                         val replacement = when (usage) {
                             is Usage.Return -> JsReturn(value).source(x.expression.source)
-                            is Usage.Assignment -> {
+                            is Usage.VariableAssignment -> {
                                 val expr = JsAstUtils.assignment(usage.target.makeRef(), value).source(x.expression.source)
                                 val statement = JsExpressionStatement(expr)
                                 statement.synthetic = usage.target in syntheticNames
                                 statement
                             }
-                            is Usage.Declaration -> {
+                            is Usage.VariableDeclaration -> {
                                 val statement: JsStatement = if (usage.count > 1) {
                                     val expr = JsAstUtils.assignment(usage.target.makeRef(), value).source(x.expression.source)
                                     val result = JsExpressionStatement(expr)
@@ -207,7 +204,7 @@ internal class TemporaryAssignmentElimination(private val root: JsBlock) {
                                 }
                                 statement
                             }
-                            is Usage.Mutation -> {
+                            is Usage.PropertyMutation -> {
                                 JsExpressionStatement(JsAstUtils.assignment(usage.target, value).source(x.expression.source))
                             }
                         }
@@ -245,8 +242,8 @@ internal class TemporaryAssignmentElimination(private val root: JsBlock) {
     private fun generateDeclarations() {
         var index = 0
         usages.values.asSequence()
-                .filter { it is Usage.Declaration && it.count > 1 }
+                .filter { it is Usage.VariableDeclaration && it.count > 1 }
-                .map { it as Usage.Declaration }
+                .map { it as Usage.VariableDeclaration }
                 .forEach {
                     val statement = JsAstUtils.newVar(it.target, null)
                     statement.synthetic = it.target in syntheticNames
@@ -272,12 +269,12 @@ internal class TemporaryAssignmentElimination(private val root: JsBlock) {
 
         class Return(statement: JsStatement) : Usage(statement)
 
-        class Assignment(statement: JsStatement, val target: JsName) : Usage(statement)
+        class VariableAssignment(statement: JsStatement, val target: JsName) : Usage(statement)
 
-        class Declaration(statement: JsStatement, val target: JsName) : Usage(statement) {
+        class VariableDeclaration(statement: JsStatement, val target: JsName) : Usage(statement) {
             var count = 0
         }
 
-        class Mutation(statement: JsStatement, val target: JsExpression) : Usage(statement)
+        class PropertyMutation(statement: JsStatement, val target: JsExpression) : Usage(statement)
     }
 }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/util/collectUtils.kt

@@ -30,23 +30,23 @@ fun collectFunctionReferencesInside(scope: JsNode): List<JsName> =
 private fun collectReferencedNames(scope: JsNode): Set<JsName> {
     val references = IdentitySet<JsName>()
 
-    object : JsVisitorWithContextImpl() {
+    object : RecursiveJsVisitor() {
-        override fun visit(x: JsBreak, ctx: JsContext<*>) = false
+        override fun visitBreak(x: JsBreak) { }
 
-        override fun visit(x: JsContinue, ctx: JsContext<*>) = false
+        override fun visitContinue(x: JsContinue) { }
 
-        override fun visit(x: JsVars.JsVar, ctx: JsContext<*>): Boolean {
+        override fun visit(x: JsVars.JsVar) {
             val initializer = x.initExpression
             if (initializer != null) {
                 accept(initializer)
             }
-            return false
         }
 
-        override fun endVisit(x: JsNameRef, ctx: JsContext<*>) {
+        override fun visitNameRef(nameRef: JsNameRef) {
-            val name = x.name
+            super.visitNameRef(nameRef)
+            val name = nameRef.name
             if (name != null) {
-                references.add(name)
+                references += name
             }
         }
     }.accept(scope)
@@ -57,29 +57,28 @@ private fun collectReferencedNames(scope: JsNode): Set<JsName> {
 fun collectUsedNames(scope: JsNode): Set<JsName> {
     val references = IdentitySet<JsName>()
 
-    object : JsVisitorWithContextImpl() {
+    object : RecursiveJsVisitor() {
-        override fun visit(x: JsBreak, ctx: JsContext<*>) = false
+        override fun visitBreak(x: JsBreak) { }
 
-        override fun visit(x: JsContinue, ctx: JsContext<*>) = false
+        override fun visitContinue(x: JsContinue) { }
 
-        override fun visit(x: JsVars.JsVar, ctx: JsContext<*>): Boolean {
+        override fun visit(x: JsVars.JsVar) {
             val initializer = x.initExpression
             if (initializer != null) {
                 accept(initializer)
             }
-            return false
         }
 
-        override fun endVisit(x: JsNameRef, ctx: JsContext<*>) {
+        override fun visitNameRef(nameRef: JsNameRef) {
-            val name = x.name
+            super.visitNameRef(nameRef)
-            if (name != null && x.qualifier == null) {
+            val name = nameRef.name
+            if (name != null && nameRef.qualifier == null) {
                 references.add(name)
             }
         }
 
-        override fun visit(x: JsFunction, ctx: JsContext<*>): Boolean {
+        override fun visitFunction(x: JsFunction) {
             references += x.collectFreeVariables()
-            return false
         }
     }.accept(scope)