Refactored Traverse implementation

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<T> TraverseBreadthFirst<T>(T root, Func<T, IEnumerable
         {
             if (childrenSelector == null) throw new ArgumentNullException(nameof(childrenSelector));
 
-            return _(); IEnumerable<T> _()
-            {
-                var queue = new Queue<T>();
-                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<T>();
+            return TraverseImpl(root, childrenSelector, queue.Enqueue, queue.Dequeue, queue);
         }
 
         /// <summary>
@@ -88,19 +78,30 @@ public static IEnumerable<T> TraverseDepthFirst<T>(T root, Func<T, IEnumerable<T
         {
             if (childrenSelector == null) throw new ArgumentNullException(nameof(childrenSelector));
 
+            // 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
+
+            var stack = new Stack<T>();
+            return TraverseImpl(root, x => childrenSelector(x).Reverse(), stack.Push, stack.Pop, stack);
+        }
+
+        static IEnumerable<T> TraverseImpl<T>(
+            T root,
+            Func<T, IEnumerable<T>> childrenSelector,
+            Action<T> Push,
+            Func<T> Pop,
+            ICollection collection)
+        {
             return _(); IEnumerable<T> _()
             {
-                var stack = new Stack<T>();
-                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);
                 }
             }
         }