hphp/hack/src/utils/trie.ml

   1 (*
   2  * Copyright (c) 2015, Facebook, Inc.
   3  * All rights reserved.
   4  *
   5  * This source code is licensed under the MIT license found in the
   6  * LICENSE file in the "hack" directory of this source tree.
   7  *
   8  *)
   9
  10 open Hh_core
  11
  12 (* Utility functions *)
  13
  14 let make_pair (a: 'a) (b: 'b) : ('a*'b) =
  15   (a,b)
  16
  17 let common_prefix (s1: string) (s2: string) : int =
  18   let i=ref 0 in
  19   let l1 = String.length s1 in
  20   let l2 = String.length s2 in
  21   while (!i < l1) && (!i < l2) && (s1.[!i] = s2.[!i]) do
  22       i:=!i+1
  23   done;
  24   !i
  25
  26 let drop s c =
  27   let l = String.length s in
  28   String.sub s c (l-c)
  29
  30 let take s c =
  31   String.sub s 0 c
  32
  33 let (|>) (o : 'a)  (f : 'a -> 'b) : 'b = f o
  34
  35 let id (x : 'a) : 'a = x
  36
  37 type 'a return = { return : 'b . 'a -> 'b }
  38
  39 let with_return (type t) (f : _ -> t) =
  40   let module Capture =
  41       struct exception Return of t end
  42   in
  43   let return = { return = (fun x -> raise (Capture.Return x)); } in
  44   try f return with Capture.Return x -> x
  45
  46 (* Trie implementation *)
  47
  48 type 'a t = Leaf of 'a | Node of ('a t) SMap.t ref
  49 let create (): 'a t = Node (ref (SMap.empty))
  50
  51 exception Inconsistent_trie of string
  52
  53 let get_node (trie : 'a t) : ('a t) SMap.t ref = match trie with
  54   | Node n -> n
  55   | _      -> raise (Inconsistent_trie "Cannot match to leaf")
  56
  57 let get_leaf (trie : 'a t) : 'a = match trie with
  58   | Leaf v -> v
  59   | _      -> raise (Inconsistent_trie "Cannot match to node")
  60
  61 (* Match a string s with a key; return a tuple:
  62     i : int    -- position where the match ends
  63     k : string -- the full key matched
  64     n : 'a t   -- the node associated with key k
  65 *)
  66 let trie_assoc_partial (trie : 'a t) (w : string) : (int * string * 'a t) option =
  67   with_return (fun e ->
  68     !(get_node trie)
  69     |> SMap.iter (fun key elt ->
  70                    let c = common_prefix key w in
  71                    if (not (c = 0)) || ((key = "") && w = "") then
  72                      e.return (Some (c, key, elt)));
  73     None)
  74
  75
  76 let rec mem (trie : 'a t) (w : string) : bool =
  77   with_return (fun e ->
  78     let (i, key, child) = match trie_assoc_partial trie w with
  79                            | Some x -> x
  80                            | None -> e.return false in
  81
  82     if key = "" then
  83       e.return true;
  84
  85     if String.length key = i then
  86       e.return (mem child (drop w i));
  87
  88     false)
  89
  90 let add_one (node : 'a t) (c : string) (inner :'a t): unit =
  91   let elts = get_node node in
  92   elts := SMap.add c inner !elts
  93
  94 (* split key in position c, put left part as new key to a new node n
  95  * and put right part as child of n, then return n *)
  96 let split_key (parent : 'a t) (key : string) (child : 'a t) (c : int) : 'a t =
  97   let left_key = take key c in
  98   let right_key = drop key c in
  99   let parent_list = get_node parent in
 100   parent_list := SMap.remove key !parent_list;
 101
 102   let n = create () in
 103   add_one parent left_key n;
 104   add_one n right_key child;
 105   n
 106
 107 let add_leaf (node : 'a t) (key : string) (v : 'a) : unit =
 108   let leaf = match key with "" -> Leaf v
 109                          | _  -> let res = create () in
 110                                  add_one res "" (Leaf v);
 111                                  res in
 112   add_one node key leaf
 113
 114
 115 let rec add ?(if_exist : 'b -> 'a -> unit = fun _ _ -> ())
 116             ~(transform : 'a -> 'b)
 117             (trie : 'b t)
 118             (w : string)
 119             (v : 'a) : unit =
 120
 121   with_return (fun e ->
 122     let (c, key, child) = match trie_assoc_partial trie w with
 123                            | Some x -> x
 124                            | None -> e.return (add_leaf trie w (transform v)) in
 125
 126     if (String.length key = c) && (w = "") then     (* leaf exists; use if_exists callback *)
 127       e.return (if_exist (get_leaf child) v);
 128
 129     if c = String.length key then    (* full key match; do final recursive call *)
 130       e.return (add child (drop w c) v ~if_exist ~transform);
 131
 132     (* Partial match: need to split key with common parts *)
 133     let n = split_key trie key child c in
 134     add_leaf n (drop w c) (transform v))
 135
 136
 137 let to_list (limit : int option)
 138             (trie : 'b t)
 139             (kmap : string -> 'a)
 140             (vmap : 'a -> 'b -> 'c) : 'c list =
 141
 142   with_return (fun e ->
 143     let reslist  = ref [] in
 144     let rescount = ref 0 in
 145
 146     let more () = match limit with Some i -> i > !rescount | None -> true in
 147
 148     let rec to_list_aux t s = match t with
 149       | Leaf v  -> if more () then
 150                      (reslist := (vmap (kmap s) v)::!reslist;
 151                       incr rescount)
 152                    else
 153                     e.return (List.rev !reslist)
 154       | Node cs ->
 155           SMap.fold (fun tail rhs _acc ->
 156                        to_list_aux rhs (s ^ tail))
 157                     !cs ()
 158
 159     in
 160     to_list_aux trie "";
 161     List.rev !reslist)
 162
 163
 164 let find_impl ?(limit : int option = None)
 165               (exact : bool)
 166               (trie : 'a t)
 167               (pre : string)
 168               (vmap : string -> 'a -> 'c) : 'c list =
 169
 170   with_return (fun e ->
 171     let append = (^) pre in
 172     let rec find_impl_aux trie p =
 173       let (c, key, child) = match trie_assoc_partial trie p with
 174                               | Some x -> x
 175                               | None -> e.return [] in
 176
 177       match (String.length key = c), (not exact) && (String.length p = c) with
 178        | true, _     when String.length p = 0 -> to_list limit child append vmap
 179        | true, true                           -> to_list limit child append vmap
 180        | true, _                              -> find_impl_aux child (drop p c)
 181
 182        | false, true -> to_list limit child (fun k -> pre ^ (drop key c) ^ k) vmap
 183
 184        | _ -> []
 185     in
 186     find_impl_aux trie pre)
 187
 188
 189 let find (trie : 'a t) (s : string) : 'a =
 190   match find_impl true trie s make_pair with
 191   | (_s, v)::_tl -> v
 192   | _ -> raise Not_found
 193
 194
 195 let find_prefix (trie : 'a t)
 196                 (s : string)
 197                 (vmap : string -> 'a -> 'b): 'b list =
 198   find_impl false trie s vmap
 199
 200 let find_prefix_limit (i : int)
 201                       (trie : 'a t)
 202                       (s : string)
 203                       (vmap : string -> 'a -> 'b) : 'b list =
 204   find_impl false trie s vmap ~limit:(Some i)
 205
 206 let remove_one (trie : 'a t) (key : string) : unit =
 207   let elts = get_node trie in
 208   elts := SMap.remove key !elts
 209
 210 let rec remove_impl (exact: bool) (trie : 'a t) (s : string): unit =
 211   with_return (fun e ->
 212     let (c, key, child) = match trie_assoc_partial trie s with
 213                            | Some x -> x
 214                            | None   -> e.return () in
 215
 216     match (String.length key = c), exact, (String.length s = c) with
 217      | true, true, true when c = 0 -> remove_one trie (take s c)
 218      | true, false, true           -> remove_one trie (take s c)
 219
 220      | true, _, _ -> remove_impl exact child (drop s c)
 221      | _          -> ())
 222
 223
 224 let remove (trie : 'a t) (s : string) : unit =
 225   remove_impl true trie s
 226
 227 let remove_prefix (trie : 'a t) (s : string) : unit =
 228   remove_impl false trie s
 229
 230 (* let rec merge ?(if_exist : 'a -> 'a -> unit = fun _ _ -> ()) *)
 231 (*               (trieDes : 'a t) *)
 232 (*               (trieSrc : 'a t) : unit = *)
 233 (*   let rec merge_one des s n = *)
 234 (*     try *)
 235 (*       let (c, key, child) = trie_assoc_partial des s in *)
 236 (*       if String.length key = c then begin *)
 237 (*         (\* matched whole key. *)
 238 (*          * Either continue to merge child, if whole s matched, *)
 239 (*          * or continue to match in des's path, if some s left*\) *)
 240 (*         if (String.length s = c) && (not (s = "")) then *)
 241 (*           merge child n ~if_exist *)
 242 (*         else if (not (s = "")) then *)
 243 (*           merge_one child (string_drop s c) n *)
 244 (*         else *)
 245 (*           (\* empty key match, means same value exist *)
 246 (*            * resolve use if_exist *\) *)
 247 (*           let desv = get_leaf child in *)
 248 (*           let srcv = get_leaf n in *)
 249 (*           if_exist desv srcv *)
 250 (*         end *)
 251 (*       else *)
 252 (*         (\* partially match key, need to split key, and merge into *\) *)
 253 (*         let common = string_take s c in *)
 254 (*         remove_one des key; *)
 255 (*         if (c = String.length s) then begin *)
 256 (*           (\* when s is fully match, just take the node keyed by s *)
 257 (*            * no need to create new one *\) *)
 258 (*           add_one des common n; *)
 259 (*           merge_one n (string_drop key c) child *)
 260 (*         end *)
 261 (*         else begin *)
 262 (*           let t = create () in *)
 263 (*           add_one t (string_drop key c) child; *)
 264 (*           add_one t (string_drop s c) n; *)
 265 (*           add_one des common t *)
 266 (*         end *)
 267 (*     with Not_found -> *)
 268 (*        add_one des s n *)
 269 (*     in *)
 270 (*   let l = get_node trieSrc in *)
 271 (*   List.iter (fun (s, n) -> merge_one trieDes s n) !l *)
 272
 273 let rec to_string_impl (buf : Buffer.t) (trie : 'a t) : unit =
 274   match trie with
 275    | Node elts -> SMap.fold
 276                     (fun k v _ ->
 277                       Printf.bprintf buf "%S:{" k;
 278                       to_string_impl buf v;
 279                       Printf.bprintf buf "}")
 280                       !elts ()
 281    | Leaf _v -> ()
 282
 283
 284 let to_string (trie : 'a t) : string =
 285   let buf = Buffer.create 250 in
 286   to_string_impl buf trie;
 287   Buffer.contents buf