From 838d419a84a1996c1120ea4f13e959054ac68138 Mon Sep 17 00:00:00 2001
From: Pavel Kunyavskiy <pavel.kunyavskiy@jetbrains.com>
Date: Sun, 13 Nov 2022 11:37:51 +0100
Subject: [PATCH] [K/N] Fix stack overflow in DirectedGraphCondensationBuilder

^KT-54732
---
 .../kotlin/backend/konan/GraphAlgorithms.kt   | 58 ++++++++++---------
 1 file changed, 30 insertions(+), 28 deletions(-)

diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/GraphAlgorithms.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/GraphAlgorithms.kt
index 8e08a573852..ca7cf6342d7 100644
--- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/GraphAlgorithms.kt
+++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/GraphAlgorithms.kt
@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.backend.konan
 
+import org.jetbrains.kotlin.utils.addToStdlib.popLast
+
 internal interface DirectedGraphNode<out K> {
     val key: K
     val directEdges: List<K>?
@@ -45,45 +47,45 @@ internal class DirectedGraphCondensationBuilder<K, out N: DirectedGraphNode<K>>(
             }
         }
 
-        // Topsort of built condensation.
-        multiNodes.forEach { multiNode ->
-            multiNode.nodes.forEach { nodeToMultiNodeMap.put(graph.get(it), multiNode) }
-        }
-        visited.clear()
-        multiNodes.forEach {
-            if (!visited.contains(it.nodes.first()))
-                findMultiNodesOrder(it)
-        }
-
-        return DirectedGraphCondensation(multiNodesOrder.reversed())
+        return DirectedGraphCondensation(multiNodes)
     }
 
     private fun findOrder(node: N) {
+        val stack = mutableListOf<Pair<N, Iterator<K>>>()
         visited += node.key
-        node.directEdges?.forEach {
-            if (!visited.contains(it))
-                findOrder(graph.get(it))
+        stack.add(node to (node.directEdges ?: emptyList()).iterator())
+        while (stack.isNotEmpty()) {
+            if (stack.last().second.hasNext()) {
+                val nextKey = stack.last().second.next()
+                if (!visited.contains(nextKey)) {
+                    visited += nextKey
+                    val nextNode = graph.get(nextKey)
+                    stack.add(nextNode to (nextNode.directEdges ?: emptyList()).iterator())
+                }
+            } else {
+                order += stack.last().first
+                stack.popLast()
+            }
         }
-        order += node
     }
 
     private fun paint(node: N, multiNode: MutableSet<K>) {
+        val stack = mutableListOf<Pair<N, Iterator<K>>>()
         visited += node.key
         multiNode += node.key
-        node.reversedEdges?.forEach {
-            if (!visited.contains(it))
-                paint(graph.get(it), multiNode)
-        }
-    }
-
-    private fun findMultiNodesOrder(node: DirectedGraphMultiNode<K>) {
-        visited.addAll(node.nodes)
-        node.nodes.forEach {
-            graph.get(it).directEdges?.forEach {
-                if (!visited.contains(it))
-                    findMultiNodesOrder(nodeToMultiNodeMap[graph.get(it)]!!)
+        stack.add(node to (node.reversedEdges ?: emptyList()).iterator())
+        while (stack.isNotEmpty()) {
+            if (stack.last().second.hasNext()) {
+                val nextKey = stack.last().second.next()
+                if (!visited.contains(nextKey)) {
+                    visited += nextKey
+                    multiNode += nextKey
+                    val nextNode = graph.get(nextKey)
+                    stack.add(nextNode to (nextNode.reversedEdges ?: emptyList()).iterator())
+                }
+            } else {
+                stack.popLast()
             }
         }
-        multiNodesOrder += node
     }
 }