feat(Trie): map_tag (#132)

src/Trie.ml

@@ -314,6 +314,17 @@ let of_seq s = of_seq_with_merger ~prefix:Emp (fun _ _ y -> y) s
 
 type 'data untagged = ('data, unit) t
 
+let rec map_tag_node f n =
+  { tag_root = Option.map f n.tag_root
+  ; tag_default_child = Option.map f n.tag_default_child
+  ; tag_children = SegMap.map (map_tag_node f) n.tag_children
+  }
+
+let[@inline] map_tag f : _ t -> _ t =
+  function
+  | None -> None
+  | Some (d, t) -> Some (d, map_tag_node f t)
+
 let[@inline] retag t : _ t -> _ t =
   function
   | None -> None

src/Trie.mli

@@ -158,10 +158,13 @@ val of_seq_with_merger : ?prefix:bwd_path -> (bwd_path -> 'data * 'tag -> 'data
 (** Untagged tries (where all tags are [()]). *)
 type 'data untagged = ('data, unit) t
 
-(** [retag tag t] changes all tags within [t] to [tag] in O(1) time. The data remain intact. *)
+(** [map_tag f trie] applies [f] to all tags within [trie], which is equivalent to {!val:map}[ (fun _ (d, t) -> (d, f t)) trie] but can often be more efficient. The data remain intact. Note that if [f] is a constant function, use {!val:retag} for even more efficiency. *)
+val map_tag : ('tag1 -> 'tag2) -> ('data, 'tag1) t -> ('data, 'tag2) t
+
+(** [retag tag t] changes all tags within [t] to [tag] in O(1) time. It is equivalent to {!val:map_tag}[ (fun _ -> tag) t] but usually more efficient. The data remain intact. *)
 val retag : 'tag -> ('data, _) t -> ('data, 'tag) t
 
-(** [retag_subtree tag path t] changes all tags within the subtrie rooted at [path] to [tag] efficiently. The data remain intact. *)
+(** [retag_subtree path tag t] changes all tags within the subtrie rooted at [path] to [tag] efficiently. The data remain intact. *)
 val retag_subtree : path -> 'tag -> ('data, 'tag) t -> ('data, 'tag) t
 
 (** [untag t] is [retag () t]. *)

test/ListAsTrie.ml

@@ -100,6 +100,7 @@ let of_seq s = Seq.fold_left (union_singleton ~prefix:Emp (fun _ _ y -> y)) empt
 let of_seq_with_merger ?(prefix=Emp) m s = Seq.fold_left (union_singleton ~prefix m) empty s
 
 type 'a untagged = (path * ('a * unit)) list
+let map_tag f l = List.map (fun (p, (d, t)) -> p, (d, f t)) l
 let retag t l = List.map (fun (p, (d, _)) -> p, (d, t)) l
 let retag_subtree pre t l =
   List.map (fun ((p, (d, _)) as b) -> if Option.is_some (split_path pre p) then p, (d, t) else b) l

test/TestTrie.ml

@@ -231,6 +231,11 @@ let test_of_seq_with_merger =
        to_list (Trie.of_seq_with_merger ?prefix f (List.to_seq l))
        =
        ListAsTrie.of_seq_with_merger ?prefix f (List.to_seq l))
+let test_map_tag =
+  Q.Test.make ~count ~name:"map_tag"
+    Q.Gen.(pair (Q.fun1 Q.Observable.int int) gen_list)
+    ~print:Q.Print.(pair Q.Fn.print print_list)
+    (fun (Fun (_, f), l) -> to_list (Trie.map_tag f (of_list l)) = ListAsTrie.map_tag f l)
 let test_retag =
   Q.Test.make ~count ~name:"retag" Q.Gen.(pair int gen_list) ~print:Q.Print.(pair int print_list)
     (fun (t, l) -> to_list (Trie.retag t (of_list l)) = ListAsTrie.retag t l)
@@ -275,6 +280,7 @@ let () =
     ; test_to_seq_values
     ; test_of_seq
     ; test_of_seq_with_merger
+    ; test_map_tag
     ; test_retag
     ; test_retag_subtree
     ; test_untag