make tree-seq accept non-branch root, patch from cgrand

1 files changed, 13 insertions, 17 deletions

diff --git a/src/clj/clojure/core.clj b/src/clj/clojure/core.clj
index 8d4d863d..5c012155 100644
--- a/src/clj/clojure/core.clj
+++ b/src/clj/clojure/core.clj
@@ -2824,23 +2824,19 @@
         `(print-namespace-doc ~nspace)
         `(print-doc (var ~name))))))
 
-(defn tree-seq
-  "returns a lazy sequence of the nodes in a tree, via a depth-first walk.
-  branch? must be a fn of one arg that returns true if passed a node
-  that can have children (but may not).  children must be a fn of one
-  arg that returns a sequence of the children. Will only be called on
-  nodes for which branch? returns true. Root is the root node of the
-  tree, must be a branch."
-  [branch? children root]
-    (let [walk (fn walk [nodes]
-                   (when-first [node nodes]
-                     (lazy-cons
-                      node
-                      (if (branch? node)
-                        (lazy-cat (walk (children node))
-                                  (walk (rest nodes)))
-                        (walk (rest nodes))))))]
-      (lazy-cons root (walk (children root)))))
+ (defn tree-seq
+  "Returns a lazy sequence of the nodes in a tree, via a depth-first walk.
+   branch? must be a fn of one arg that returns true if passed a node
+   that can have children (but may not).  children must be a fn of one
+   arg that returns a sequence of the children. Will only be called on
+   nodes for which branch? returns true. Root is the root node of the
+  tree."  
+   [branch? children root]
+   (let [walk (fn walk [node]
+                (lazy-cons node
+                  (when (branch? node)
+                    (mapcat walk (children node)))))]
+     (walk root)))
 
 (defn file-seq
   "A tree seq on java.io.Files"