JS DCE: use less LinkedHashSets

compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/dce/K2JSDce.kt

@@ -78,11 +78,10 @@ class K2JSDce : CLITool<K2JSDceArguments>() {
 
         val dceResult = DeadCodeElimination.run(files, includedDeclarations, logConsumer)
         if (dceResult.status == DeadCodeEliminationStatus.FAILED) return ExitCode.COMPILATION_ERROR
-        val nodes = dceResult.reachableNodes.filterTo(mutableSetOf()) { it.reachable }
 
         val reachabilitySeverity = if (arguments.printReachabilityInfo) CompilerMessageSeverity.INFO else CompilerMessageSeverity.LOGGING
         messageCollector.report(reachabilitySeverity, "")
-        for (node in nodes.extractRoots()) {
+        for (node in dceResult.reachableNodes.extractReachableRoots(dceResult.context!!)) {
             printTree(
                 node, { messageCollector.report(reachabilitySeverity, it) },
                 printNestedMembers = false, showLocations = true

js/js.dce/src/org/jetbrains/kotlin/js/dce/Context.kt

@@ -126,6 +126,14 @@ class Context {
         }
     }
 
+    private var currentColor = 1.toByte()
+
+    fun clearVisited() {
+        currentColor++
+    }
+
+    fun visit(n: Node) = n.visit(currentColor)
+
     class Node private constructor(val localName: JsName?, qualifier: Qualifier?) {
         private var _dependenciesImpl: MutableSet<Node>? = null
         private var _expressionsImpl: MutableSet<JsExpression>? = null
@@ -178,6 +186,14 @@ class Context {
         var qualifier: Qualifier? = qualifier
             private set
 
+        private var color: Byte = 0
+
+        fun visit(c: Byte): Boolean {
+            val result = color != c
+            color = c
+            return result
+        }
+
         val memberNames: Set<String> get() = original._membersImpl?.keys ?: emptySet()
 
         constructor(localName: JsName? = null) : this(localName, null)

js/js.dce/src/org/jetbrains/kotlin/js/dce/DeadCodeElimination.kt

@@ -42,11 +42,14 @@ class DeadCodeElimination(private val logConsumer: (DCELogLevel, String) -> Unit
     val moduleMapping = mutableMapOf<JsBlock, String>()
     private val reachableNames = mutableSetOf<String>()
 
-    var reachableNodes = setOf<Node>()
+    var reachableNodes: Iterable<Node> = setOf()
         private set
 
+    var context: Context? = null
+
     fun apply(root: JsNode) {
         val context = Context()
+        this.context = context
 
         val topLevelVars = collectDefinedNames(root)
         context.addNodesForLocalVars(topLevelVars)
@@ -107,7 +110,7 @@ class DeadCodeElimination(private val logConsumer: (DCELogLevel, String) -> Unit
                 block
             }
 
-            if (hasErrors) return DeadCodeEliminationResult(emptySet(), DeadCodeEliminationStatus.FAILED)
+            if (hasErrors) return DeadCodeEliminationResult(dce.context, emptySet(), DeadCodeEliminationStatus.FAILED)
 
             program.globalBlock.statements += blocks
             program.globalBlock.fixForwardNameReferences()
@@ -139,7 +142,7 @@ class DeadCodeElimination(private val logConsumer: (DCELogLevel, String) -> Unit
                 }
             }
 
-            return DeadCodeEliminationResult(dce.reachableNodes, DeadCodeEliminationStatus.OK)
+            return DeadCodeEliminationResult(dce.context, dce.reachableNodes, DeadCodeEliminationStatus.OK)
         }
 
         private class Reporter(private val fileName: String, private val logConsumer: (DCELogLevel, String) -> Unit) : ErrorReporter {

js/js.dce/src/org/jetbrains/kotlin/js/dce/DeadCodeEliminationResult.kt

@@ -16,4 +16,4 @@
 
 package org.jetbrains.kotlin.js.dce
 
-class DeadCodeEliminationResult(val reachableNodes: Set<Context.Node>, val status: DeadCodeEliminationStatus)
+class DeadCodeEliminationResult(val context: Context?, val reachableNodes: Iterable<Context.Node>, val status: DeadCodeEliminationStatus)

js/js.dce/src/org/jetbrains/kotlin/js/dce/ReachabilityTracker.kt

@@ -29,11 +29,15 @@ class ReachabilityTracker(
         private val CALL_FUNCTIONS = setOf("call", "apply")
     }
 
+    init {
+        context.clearVisited()
+    }
+
     private var currentNodeWithLocation: JsNode? = null
     private var depth = 0
-    private val reachableNodesImpl = mutableSetOf<Node>()
+    private val reachableNodesImpl = mutableListOf<Node>()
 
-    val reachableNodes: Set<Node> get() = reachableNodesImpl
+    val reachableNodes: Iterable<Node> get() = reachableNodesImpl
 
     override fun visit(x: JsVars.JsVar) {
         if (shouldTraverse(x)) {
@@ -137,7 +141,9 @@ class ReachabilityTracker(
     fun reach(node: Node) {
         if (node.reachable) return
         node.reachable = true
-        reachableNodesImpl += node
+        if (context.visit(node)) {
+            reachableNodesImpl += node
+        }
 
         reachDeclaration(node)
 
@@ -192,7 +198,9 @@ class ReachabilityTracker(
         }
         else if (!node.declarationReachable) {
             node.declarationReachable = true
-            reachableNodesImpl += node
+            if (context.visit(node)) {
+                reachableNodesImpl += node
+            }
 
             node.original.qualifier?.parent?.let {
                 reportAndNest("reach-decl: parent $it", null) {

js/js.dce/src/org/jetbrains/kotlin/js/dce/util.kt

@@ -73,20 +73,21 @@ fun JsLocation.asString(): String {
     return "$simpleFileName:${startLine + 1}"
 }
 
-fun Set<Node>.extractRoots(): Set<Node> {
+fun Iterable<Node>.extractReachableRoots(context: Context): Iterable<Node> {
-    val result = mutableSetOf<Node>()
+    context.clearVisited()
-    val visited = mutableSetOf<Node>()
+
-    forEach { it.original.extractRootsImpl(result, visited) }
+    val result = mutableListOf<Node>()
+    forEach { if (it.reachable) it.original.extractRootsImpl(result, context) }
     return result
 }
 
-private fun Node.extractRootsImpl(target: MutableSet<Node>, visited: MutableSet<Node>) {
+private fun Node.extractRootsImpl(target: MutableList<Node>, context: Context) {
-    if (!visited.add(original)) return
+    if (!context.visit(original)) return
     val qualifier = original.qualifier
     if (qualifier == null) {
         target += original
     }
     else {
-        qualifier.parent.extractRootsImpl(target, visited)
+        qualifier.parent.extractRootsImpl(target, context)
     }
 }