diff --git a/MoreLinq/Traverse.cs b/MoreLinq/Traverse.cs index 4fb46b5b8..2745ab653 100644 --- a/MoreLinq/Traverse.cs +++ b/MoreLinq/Traverse.cs @@ -18,6 +18,7 @@ namespace MoreLinq { using System; + using System.Collections; using System.Collections.Generic; using System.Linq; @@ -48,19 +49,8 @@ public static IEnumerable TraverseBreadthFirst(T root, Func _() - { - var queue = new Queue(); - queue.Enqueue(root); - - while (queue.Count != 0) - { - var current = queue.Dequeue(); - yield return current; - foreach (var child in childrenSelector(current)) - queue.Enqueue(child); - } - } + var queue = new Queue(); + return TraverseImpl(root, childrenSelector, queue.Enqueue, queue.Dequeue, queue); } /// @@ -88,19 +78,30 @@ public static IEnumerable TraverseDepthFirst(T root, Func(); + return TraverseImpl(root, x => childrenSelector(x).Reverse(), stack.Push, stack.Pop, stack); + } + + static IEnumerable TraverseImpl( + T root, + Func> childrenSelector, + Action Push, + Func Pop, + ICollection collection) + { return _(); IEnumerable _() { - var stack = new Stack(); - stack.Push(root); + Push(root); - while (stack.Count != 0) + while (collection.Count != 0) { - var current = stack.Pop(); + var current = Pop(); yield return current; - // because a stack pops the elements out in LIFO order, we need to push them in reverse - // if we want to traverse the returned list in the same order as was returned to us - foreach (var child in childrenSelector(current).Reverse()) - stack.Push(child); + foreach (var child in childrenSelector(current)) + Push(child); } } }