From 6b59bb99e25371040fc63066eafd96863ccd50b0 Mon Sep 17 00:00:00 2001
From: Lamberto Basti <basti.lamberto@gmail.com>
Date: Wed, 27 Sep 2023 17:38:55 +0200
Subject: [PATCH] Implement tree traversal methods in Tree class (#135)

* Implement tree traversal methods in Tree class

Enhanced the Tree class in the jewel foundation by adding two new methods for performing tree traversal: 'walkBreadthFirst' and 'walkDepthFirst'. These enhancements provide flexibility in tree exploration based on the application needs. Modifying the 'roots' variable from 'internal' to 'val' to facilitate tree manipulation.

* Refactor conditional statements in Tree.kt
---
 .../jetbrains/jewel/foundation/tree/Tree.kt   | 25 ++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/core/src/main/kotlin/org/jetbrains/jewel/foundation/tree/Tree.kt b/core/src/main/kotlin/org/jetbrains/jewel/foundation/tree/Tree.kt
index 27e29cbd9..89649c200 100644
--- a/core/src/main/kotlin/org/jetbrains/jewel/foundation/tree/Tree.kt
+++ b/core/src/main/kotlin/org/jetbrains/jewel/foundation/tree/Tree.kt
@@ -1,14 +1,37 @@
 package org.jetbrains.jewel.foundation.tree
 
+import org.jetbrains.jewel.foundation.tree.Tree.Element.Node
+
 @Suppress("UNCHECKED_CAST")
 fun <T> emptyTree() = Tree.EMPTY as Tree<T>
 
-class Tree<T> internal constructor(internal val roots: List<Element<T>>) {
+class Tree<T> internal constructor(val roots: List<Element<T>>) {
 
     companion object {
         internal val EMPTY = Tree(roots = emptyList<Element<Any?>>())
     }
 
+    fun isEmpty() = roots.isEmpty()
+
+    private fun walk(breathFirst: Boolean) = sequence {
+        val queue = roots.toMutableList()
+        while (queue.isNotEmpty()) {
+            val next = queue.removeFirst()
+            yield(next)
+            if (next is Node) {
+                next.open()
+                if (breathFirst) {
+                    queue.addAll(next.children.orEmpty())
+                } else {
+                    queue.addAll(0, next.children.orEmpty())
+                }
+            }
+        }
+    }
+
+    fun walkBreadthFirst() = walk(true)
+    fun walkDepthFirst() = walk(false)
+
     sealed interface Element<T> {
 
         val data: T