T8232: add extended fold over tree function to track level

The function fold_tree_with_path_cont allows the called function to
return a list of values, reset at each return to local root. The
motivating example is to carry a stack of boolean values to ignore paths
below a certain level.

2 files changed, 28 insertions, 0 deletions

diff --git a/src/vytree.ml b/src/vytree.ml
index 7c4e4fb..6ecc867 100644
--- a/src/vytree.ml
+++ b/src/vytree.ml
@@ -306,3 +306,28 @@ let fold_tree_with_path f (p', a) t =
         List.fold_left (fold_func f) (f (p, a) t) c in
         (Util.drop_first p), snd res
     in snd (fold_func f (p', a) t)
+
+(** Allow function called in fold to maintain a list of values for each
+    depth level of tree. A simple example is for the the function to cons a
+    boolean value to the list v at each call of the depth-first traversal;
+    at the return to local root, the value for that level is restored.
+    Note that if the function returns the empty list, this function reduces
+    to fold_tree_with_path.
+ *)
+
+let fold_tree_with_path_and_list f ((p', v), a) t =
+    let rec fold_func f ((p', v), a) t =
+    let p =
+        match name_of_node t with
+        | "" -> p'
+        | name -> name :: p'
+    in
+    let children = children_of_node t in
+    match children with
+    | [] -> let res =
+        f ((p, v), a) t in
+        (Util.drop_first p, Util.drop_first (snd (fst res))), snd res
+    | c -> let res =
+        List.fold_left (fold_func f) (f ((p, v), a) t) c in
+        (Util.drop_first p, Util.drop_first (snd (fst res))), snd res
+    in snd (fold_func f ((p', v), a) t)
diff --git a/src/vytree.mli b/src/vytree.mli
index 023d967..448aa84 100644
--- a/src/vytree.mli
+++ b/src/vytree.mli
@@ -96,3 +96,6 @@ val is_terminal_path : 'a t -> string list -> bool
 [@@alert exn "Vytree.Empty_path"]
 
 val fold_tree_with_path: (string list * 'acc -> 'b t -> string list * 'acc) -> string list * 'acc -> 'b t -> 'acc
+
+val fold_tree_with_path_and_list: ((string list * bool list) * 'acc -> 'b t ->
+    (string list * bool list) * 'acc) -> (string list * bool list) * 'acc -> 'b t -> 'acc