feat(iterators): add children arg for walk()/walkIterator()

packages/iterators/README.md

@@ -856,12 +856,22 @@ less than `n`, if input is too short...)
 ```
 
 ### `walk(fn: (x: any) => void, input: Iterable<any>, postOrder = false) => void`
+### `walk(fn: (x: any) => void, children: (x: any) => any, input: Iterable<any>, postOrder = false) => void`
 
 Recursively walks input and applies `fn` to each element (for in-place
 editing or side effects). Only iterable values and objects (but not
 strings) are traversed further. Traversal is pre-order by default, but
 can be changed to post-order via last arg.
 
+Note: Object traversal is done via `objectIterator()` and so will
+include pairs of `[k, v]` values.
+
+The optional `children` fn can be used to select child values of the
+currently visited value. If this function is given then only its return
+values will be traversed further (with the same constraint as mentioned
+above). If the fn returns `null` or `undefined`, no children will be
+visited.
+
 ```ts
 // dummy SVG document
 let doc = {
@@ -890,7 +900,7 @@ let circleTX = x => {
 };
 
 // transform doc
-ti.walk(circleTX, doc);
+ti.walk(circleTX, (x) => x.content, doc);
 
 doc.content[0].content[1]
 // { tag: "circle", attr: { x: 83.9269, y: 31.129, r: 5 } }
@@ -899,12 +909,22 @@ doc.content[1]
 ```
 
 ### `walkIterator(input: Iterable<any>, postOrder = false) => IterableIterator<any>`
+### `walkIterator(input: Iterable<any>, children: (x: any) => any, postOrder = false) => IterableIterator<any>`
 
 Yields an iterator performing either pre-order (default) or post-order
 [traversal](https://en.wikipedia.org/wiki/Tree_traversal#Pre-order) of
 input. Only iterable values and objects (but not strings) are traversed
 further.
 
+Note: Object traversal is done via `objectIterator()` and so will
+include pairs of `[k, v]` values.
+
+The optional `children` fn can be used to select child values of the
+currently visited value. If this function is given then only its return
+values will be traversed further (with the same constraint as mentioned
+above). If the fn returns `null` or `undefined`, no children will be
+visited.
+
 ```ts
 // pre-order traversal
 [...ti.map(JSON.stringify, ti.walkIterator([[[1, [2]], [3, [4]]], [5]], false))]

packages/iterators/src/walk.ts

@@ -1,33 +1,67 @@
 import { iterator, maybeIterator } from "./iterator";
 import { maybeObjectIterator } from "./object-iterator";
+import { Fn } from "@thi.ng/api/api";
 
 export function walkable(x) {
     return typeof x !== "string" ? maybeIterator(x) || maybeObjectIterator(x) : undefined;
 }
 
-export function walk(fn: (x: any) => void, input: Iterable<any>, postOrder = false) {
+export function walk(fn: Fn<any, void>, input: any, postOrder?: boolean);
+export function walk(fn: Fn<any, void>, children: Fn<any, any>, input: any, postOrder?: boolean);
+export function walk(fn: Fn<any, void>, ...args: any[]) {
+    let children: Fn<any, any>;
+    let input: any;
+    let postOrder: boolean;
+    if (args.length === 3) {
+        [children, input, postOrder] = args;
+    } else if (args.length === 2 && typeof args[0] === "function") {
+        [children, input] = args;
+    } else {
+        [input, postOrder] = args;
+    }
     let inner = (iter) => {
-        let v: IteratorResult<any>,
-            node;
+        let v: IteratorResult<any>;
         while (((v = iter.next()), !v.done)) {
             if (!postOrder) { fn(v.value); }
-            if ((node = walkable(v.value)) !== undefined) {
-                inner(node);
+            let cvals;
+            if (children) {
+                cvals = children(v.value);
+            } else {
+                cvals = v.value;
+            }
+            if ((cvals = walkable(cvals)) !== undefined) {
+                inner(cvals);
             }
             if (postOrder) { fn(v.value); }
         }
     };
     inner(iterator([input]));
 }
 
-export function walkIterator(input: Iterable<any>, postOrder = false) {
+export function walkIterator(input: any, postOrder?: boolean): IterableIterator<any>;
+export function walkIterator(input: any, children: Fn<any, any>, postOrder?: boolean): IterableIterator<any>;
+export function walkIterator(input: any, ...args: any[]) {
+    let children: Fn<any, any>;
+    let postOrder: boolean;
+    if (args.length === 2) {
+        [children, postOrder] = args;
+    } else if (typeof args[0] === "function") {
+        children = args[0];
+    } else {
+        postOrder = args[0];
+    }
     let walk = function* (iter) {
-        let v: IteratorResult<any>,
-            node;
+        let v: IteratorResult<any>;
         while (((v = iter.next()), !v.done)) {
             if (!postOrder) { yield v.value; }
-            if ((node = walkable(v.value)) !== undefined) {
-                yield* walk(node);
+            let cvals;
+            if (children) {
+                cvals = children(v.value);
+            } else {
+                cvals = v.value;
+            }
+            if ((cvals = walkable(cvals)) !== undefined) {
+                yield* walk(cvals);
             }
             if (postOrder) { yield v.value; }
         }

packages/iterators/test/index.ts

@@ -373,6 +373,18 @@ describe("iterators", function () {
                 [4],
                 4
             ]);
+        assert.deepEqual(
+            [...ti.walkIterator([[1, { a: [2] }], ["3", [4]]], (x) => Array.isArray(x) ? x : null, false)],
+            [
+                [[1, { a: [2] }], ["3", [4]]],
+                [1, { a: [2] }],
+                1,
+                { a: [2] },
+                ["3", [4]],
+                "3",
+                [4],
+                4
+            ]);
         assert.deepEqual(
             [...ti.walkIterator([[1, { a: [2] }], ["3", [4]]], true)],
             [1,