style(Trie): (possibly) improve time complexity of union

src/Trie.ml

@@ -27,7 +27,7 @@ type ('data, 'tag) node = 'data data_node * 'tag tag_node
 
   Non-invariants:
   1. The tag trie need not be minimum.
-  2. This implementation prefers removing tag_default_child.
+  2. This module prefers removing tag_default_child.
 *)
 
 type ('data, 'tag) t = ('data, 'tag) node option
@@ -53,6 +53,20 @@ let mk_tag_node' d t : _ tag_node = mk_tag_node d (t, (t, SegMap.empty))
 let mk_node d tag_params : _ node = d, mk_tag_node d tag_params
 let mk_node' d tag : _ node = d, mk_tag_node' d tag
 
+(* invariants: the input is already valid *)
+let drop_tag_default_child (d, t) =
+  if t.tag_default_child = None then (d, t) else
+    let tag_children =
+      SegMap.merge
+        (fun _ child tag_child ->
+           match child, tag_child with
+           | None, _ -> assert false
+           | Some d, None -> Some (mk_tag_node' d t.tag_default_child)
+           | Some _, Some t -> Some t)
+        d.children t.tag_children
+    in
+    d, { t with tag_default_child = None; tag_children }
+
 (* invariants: input tag tree must be a subset (if default tags were ignored) *)
 let mk_tree (root, children) tag_params : _ t =
   if Option.is_none root && SegMap.is_empty children
@@ -84,15 +98,15 @@ let get_children_node (d, t) =
   SegMap.merge
     (fun _ d' t' ->
        match d', t' with
-       | None, _ -> None
+       | None, _ -> assert false
        | Some d, None -> Some (d, mk_tag_node' d t.tag_default_child)
        | Some d, Some t -> Some (d, t))
     d.children t.tag_children
 let get_children_node2 (d, t1, t2) =
   SegMap.merge
     (fun _ d_t1 t2' ->
        match d_t1, t2' with
-       | None, _ -> None
+       | None, _ -> assert false
        | Some (d, t1), None -> Some (d, t1, mk_tag_node' d t2.tag_default_child)
        | Some (d, t1), Some t2 -> Some (d, t1, t2))
     (get_children_node (d, t1)) t2.tag_children
@@ -191,15 +205,32 @@ let union_option m r1 r2 =
   | Some r, None | None, Some r -> Some r
   | Some r1, Some r2 -> Some (m r1 r2)
 
+(* this function is optimized for the cases where the merging is rare *)
 let rec union_node ~prefix m n1 n2 =
+  let (nd1, nt1) as n1 = drop_tag_default_child n1
+  and (nd2, nt2) as n2 = drop_tag_default_child n2
+  in
   let root, tag_root = split_option @@
     union_option (m prefix) (find_root_node n1) (find_root_node n2)
   in
-  let children, tag_children =
-    split_children @@ SegMap.union
-      (fun seg n1 n2 -> Some (union_node ~prefix:(prefix #< seg) m n1 n2))
-      (get_children_node n1) (get_children_node n2)
+  let tag_exclusive_children =
+    SegMap.union
+      (fun _seg _t1 _t2 -> None)
+      nt1.tag_children nt2.tag_children
+  in
+  let tag_overlapping_children = ref SegMap.empty in
+  let children =
+    SegMap.union
+      (fun seg d1 d2 ->
+         let t1 = SegMap.find seg nt1.tag_children
+         and t2 = SegMap.find seg nt2.tag_children
+         in
+         let d, t = union_node ~prefix:(prefix #< seg) m (d1, t1) (d2, t2) in
+         tag_overlapping_children := SegMap.add seg t !tag_overlapping_children;
+         Some d)
+      nd1.children nd2.children
   in
+  let tag_children = SegMap.union (fun _ _ _ -> assert false) tag_exclusive_children !tag_overlapping_children in
   {root; children}, {tag_root; tag_default_child = None; tag_children}
 
 let union_ ~prefix m = union_option (union_node ~prefix m)