From db8fa8ae80dc28508cfd8a7c1b01fa2d6928ba7c Mon Sep 17 00:00:00 2001
From: Humza Shahid <humzasaur@gmail.com>
Date: Wed, 11 Sep 2024 11:25:33 +0100
Subject: [PATCH] add getPrefixList to bro-tree (just folds over the tree in
 O(n) time, which is the best we can do for that data structure). Will add a
 benchmark comparison of this.

---
 bench/bro-tree.sml     | 20 ++++++++++++++++++++
 src/zip-string-set.sml | 30 ++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/bench/bro-tree.sml b/bench/bro-tree.sml
index 7306e36..55bafec 100644
--- a/bench/bro-tree.sml
+++ b/bench/bro-tree.sml
@@ -64,4 +64,24 @@ struct
         else if str > k then exists (str, r)
         else true
     | _ => raise Match
+
+  fun foldr (f, acc, tree: bro) =
+    case tree of
+      N0 => acc
+    | N1 t => foldr (f, acc, t)
+    | N2 (l, k, r) =>
+        let
+          val acc = foldr (f, acc, r)
+          val acc = f (k, acc)
+        in
+          foldr (f, acc, l)
+        end
+    | _ => raise Match
+
+  fun getPrefixList (prefix, tree) =
+    foldr
+      ( (fn (k, acc) => if String.isSubstring prefix k then k :: acc else acc)
+      , []
+      , tree
+      )
 end
diff --git a/src/zip-string-set.sml b/src/zip-string-set.sml
index 7df0ade..19222d1 100644
--- a/src/zip-string-set.sml
+++ b/src/zip-string-set.sml
@@ -490,4 +490,34 @@ struct
   fun insertBinSearch (findChr, keyPos, children) =
     helpInsertBinSearch
       (findChr, keyPos, children, 0, Vector.length children - 1)
+
+  fun insertDifferenceFoundAtLeft
+    ( insKey
+    , insIdx
+    , splitTrieKeyStart
+    , trieChild
+    , childLeftKeys
+    , childLeftChildren
+    , childRightKeys
+    , childRightChildren
+    , parentKeys
+    , parentChildren
+    , parentConstructor
+    ) =
+    let
+      (* child node should always have CHILDREN case,
+       * because we are splitting prefix into two,
+       * when neither trieKey nor insKey match the prefix. *)
+      val childNode = CHILDREN {keys = childKeys, children = childChildren}
+      val keys =
+        Vector.mapi
+          (fn (idx, key) => if idx <> insIdx then key else splitTrieKeyStart)
+          parentKeys
+
+      val children =
+        Vector.mapi (fn (idx, elt) => if idx <> insIdx then elt else childNode)
+          parentChildren
+    in
+      parentConstructor {keys = keys, children = children}
+    end
 end