implement pixie.stdlib/tree-seq

pixie/stdlib.pxi

@@ -2126,3 +2126,16 @@ Expands to calls to `extend-type`."
    :added "0.1"}
   ([] (trace *e))
   ([e] (seq e)))
+
+(defn tree-seq
+  "Returns a lazy sequence of the nodes in a tree via a depth-first walk.
+   branch? - fn of node that should true when node has children
+   children - fn of node that should return a sequence of children (called if branch? true)
+   root - root node of the tree"
+  [branch? children root]
+  (let [walk (fn walk [node]
+               (lazy-seq
+                (cons node
+                  (when (branch? node)
+                    (mapcat walk (children node))))))]
+    (walk root)))

tests/pixie/tests/test-stdlib.pxi

@@ -378,10 +378,18 @@
   (t/assert= (transduce (drop-while even?) conj [0 2] [1 4 6]) [0 2 1 4 6])
   (t/assert= (transduce (drop-while even?) conj [0 2] [2 4 6 7 8]) [0 2 7 8]))
 
-
 (t/deftest test-trace
   (try
     (/ 0 0)
     (catch e
         (t/assert= (first (trace e)) {:type :runtime :data "Divide by zero"})
         (t/assert= (second (trace e)) {:type :native :name "_div"} ))))
+
+(t/deftest test-tree-seq
+  (t/assert= (vec (filter string?
+                          (tree-seq map?
+                                    :ch
+                                    {:ch [{:ch ["a" "b"]}
+                                          {:ch ["c" "d"]}
+                                          {:ch [{:ch ["e" {:ch ["f"]}]}]}]})))
+             ["a" "b" "c" "d" "e" "f"]))