simplified CCTrie implementation

This commit is contained in:
Simon Cruanes 2015-02-18 11:32:09 +01:00
parent f426a97a31
commit 705fcff4ec

View file

@ -125,7 +125,7 @@ module Make(W : WORD) = struct
type 'a t = type 'a t =
| Empty | Empty
| Path of char_ list * 'a t | Cons of char_ * 'a t (* simple case *)
| Node of 'a option * 'a t M.t | Node of 'a option * 'a t M.t
(* invariants: (* invariants:
@ -136,7 +136,6 @@ module Make(W : WORD) = struct
let empty = Empty let empty = Empty
let _invariant = function let _invariant = function
| Path ([],_) -> false
| Node (None, map) when M.is_empty map -> false | Node (None, map) when M.is_empty map -> false
| _ -> true | _ -> true
@ -164,10 +163,6 @@ module Make(W : WORD) = struct
let _seq_map map k = let _seq_map map k =
M.iter (fun key v -> k (key,v)) map M.iter (fun key v -> k (key,v)) map
let _is_path = function
| Path _ -> true
| _ -> false
(* return common prefix, and disjoint suffixes *) (* return common prefix, and disjoint suffixes *)
let rec _merge_lists l1 l2 = match l1, l2 with let rec _merge_lists l1 l2 = match l1, l2 with
| [], _ | [], _
@ -180,41 +175,28 @@ module Make(W : WORD) = struct
else else
[], l1, l2 [], l1, l2
(* prefix [l] to the tree [t] *)
let _mk_path l t = match l, t with
| [], _ -> t
| _, Empty -> Empty
| _, Node _ -> Path (l, t)
| _, Path (l',t') ->
assert (not(_is_path t'));
Path (l@l', t')
let _mk_path_cons x t = match t with
| Empty -> Empty
| Node _ -> Path ([x], t)
| Path (l', t') ->
assert (not(_is_path t'));
Path (x::l', t')
(* sub-tree t prefixed with c *)
let _cons c t = Cons (c, t)
(* build a Node value *) (* build a Node value *)
let _mk_node value map = match value with let _mk_node value map = match value with
| Some _ -> Node (value, map) | Some _ -> Node (value, map)
| None -> | None ->
if M.is_empty map then Empty if M.is_empty map then Empty
else else
let high, t' = M.max_binding map in if M.cardinal map = 1
let low, _ = M.min_binding map in then
if W.compare low high = 0 let c, sub = M.min_binding map in
then _mk_path [high] t' (* only one element *) _cons c sub
else Node (value,map) else Node (value,map)
let _remove_sub c t = match t with (* remove key [c] from [t] *)
let _remove c t = match t with
| Empty -> t | Empty -> t
| Path ([], _) -> assert false | Cons (c', _) ->
| Path (c'::_, _) -> if W.compare c c' = 0
if W.compare c c' = 0 then Empty
then Empty else t
else t
| Node (value, map) -> | Node (value, map) ->
if M.mem c map if M.mem c map
then then
@ -223,29 +205,23 @@ module Make(W : WORD) = struct
else t else t
let update key f t = let update key f t =
(* [state]: current subtree and rebuild function; [x]: current char *) (* first arg: current subtree and rebuild function; [c]: current char *)
let goto (t, rebuild) c = let goto (t, rebuild) c =
match t with match t with
| Empty -> (t, fun t -> rebuild (_mk_path_cons c t)) | Empty -> empty, fun t -> rebuild (_cons c t)
| Path ([], _) -> assert false | Cons (c', t') ->
| Path (c'::l, t') -> if W.compare c c' = 0
if W.compare c c' = 0 then t', (fun t -> rebuild (_cons c t))
then else
(* follow the path *) let rebuild' new_child =
_mk_path l t', (fun t -> rebuild (_mk_path_cons c t)) rebuild (
else if is_empty new_child then t
(* exit the path, so we have an empty tree. Upon rebuild we else
potentially need to make a map *) let map = M.singleton c new_child in
let rebuild' new_child = let map = M.add c' t' map in
rebuild ( _mk_node None map
if is_empty new_child then t ) in
else empty, rebuild'
let map = M.singleton c new_child in
let map = M.add c' (_mk_path l t') map in
_mk_node None map
)
in
empty, rebuild'
| Node (value, map) -> | Node (value, map) ->
try try
let t' = M.find c map in let t' = M.find c map in
@ -271,13 +247,11 @@ module Make(W : WORD) = struct
in in
let finish (t,rebuild) = match t with let finish (t,rebuild) = match t with
| Empty -> rebuild (_mk_node (f None) M.empty) | Empty -> rebuild (_mk_node (f None) M.empty)
| Path ([], _) -> assert false | Cons (c, t') ->
| Path (c::l', t') -> rebuild
rebuild ( (match f None with
match f None with | None -> t
| None -> t (* TODO: raise exception & return original tree *) | Some _ as v -> _mk_node v (M.singleton c t')
| Some _ as v ->
_mk_node v (M.singleton c (_mk_path l' t'))
) )
| Node (value, map) -> | Node (value, map) ->
let value' = f value in let value' = f value in
@ -294,10 +268,9 @@ module Make(W : WORD) = struct
(* at subtree [t], and character [c] *) (* at subtree [t], and character [c] *)
let goto t c = match t with let goto t c = match t with
| Empty -> raise Not_found | Empty -> raise Not_found
| Path ([], _) -> assert false | Cons (c', t') ->
| Path (c'::l, t') ->
if W.compare c c' = 0 if W.compare c c' = 0
then _mk_path l t' then t'
else raise Not_found else raise Not_found
| Node (_, map) -> M.find c map | Node (_, map) -> M.find c map
and finish t = match t with and finish t = match t with
@ -311,7 +284,6 @@ module Make(W : WORD) = struct
try Some (find_exn k t) try Some (find_exn k t)
with Not_found -> None with Not_found -> None
let _difflist_append f l = fun l' -> f (l @ l')
let _difflist_add f x = fun l' -> f (x :: l') let _difflist_add f x = fun l' -> f (x :: l')
(* fold that also keeps the path from the root, so as to provide the list (* fold that also keeps the path from the root, so as to provide the list
@ -319,7 +291,7 @@ module Make(W : WORD) = struct
a function that prepends a list to some suffix *) a function that prepends a list to some suffix *)
let rec _fold f path t acc = match t with let rec _fold f path t acc = match t with
| Empty -> acc | Empty -> acc
| Path (l, t') -> _fold f (_difflist_append path l) t' acc | Cons (c, t') -> _fold f (_difflist_add path c) t' acc
| Node (v, map) -> | Node (v, map) ->
let acc = match v with let acc = match v with
| None -> acc | None -> acc
@ -350,7 +322,7 @@ module Make(W : WORD) = struct
let rec fold_values f acc t = match t with let rec fold_values f acc t = match t with
| Empty -> acc | Empty -> acc
| Path (_, t') -> fold_values f acc t' | Cons (_, t') -> fold_values f acc t'
| Node (v, map) -> | Node (v, map) ->
let acc = match v with let acc = match v with
| None -> acc | None -> acc
@ -365,29 +337,19 @@ module Make(W : WORD) = struct
let rec merge f t1 t2 = match t1, t2 with let rec merge f t1 t2 = match t1, t2 with
| Empty, _ -> t2 | Empty, _ -> t2
| _, Empty -> t1 | _, Empty -> t1
| Path (l1,t1'), Path (l2,t2') -> | Cons (c1,t1'), Cons (c2,t2') ->
let common, l1', l2' = _merge_lists l1 l2 in if W.compare c1 c2 = 0
begin match l1', l2' with then _cons c1 (merge f t1' t2')
| c1::l1'', c2::l2'' -> else
(* need to build a map here, to represent the choice let map = M.add c1 t1' M.empty in
between [c1] and [c2] *) let map = M.add c2 t2' map in
assert (W.compare c1 c2 <> 0); _mk_node None map
let map = M.add c1 (_mk_path l1'' t1') M.empty in
let map = M.add c2 (_mk_path l2'' t2') map in | Cons (c1, t1'), Node (value, map) ->
_mk_path common (Node (None, map))
| _ ->
_mk_path common
(merge f
(_mk_path l1' t1')
(_mk_path l2' t2')
)
end
| Path ([], _), _ -> assert false
| Path (c1::l1, t1'), Node (value, map) ->
begin try begin try
(* collision *) (* collision *)
let t2' = M.find c1 map in let t2' = M.find c1 map in
let new_t = merge f (_mk_path l1 t1') t2' in let new_t = merge f t1' t2' in
let map' = if is_empty new_t let map' = if is_empty new_t
then M.remove c1 map then M.remove c1 map
else M.add c1 new_t map else M.add c1 new_t map
@ -396,9 +358,9 @@ module Make(W : WORD) = struct
with Not_found -> with Not_found ->
(* no collision *) (* no collision *)
assert (not(is_empty t1')); assert (not(is_empty t1'));
Node (value, M.add c1 (_mk_path l1 t1') map) Node (value, M.add c1 t1' map)
end end
| Node _, Path _ -> merge f t2 t1 (* previous case *) | Node _, Cons _ -> merge f t2 t1 (* previous case *)
| Node(v1, map1), Node (v2, map2) -> | Node(v1, map1), Node (v2, map2) ->
let v = match v1, v2 with let v = match v1, v2 with
| None, _ -> v2 | None, _ -> v2
@ -419,7 +381,7 @@ module Make(W : WORD) = struct
let rec size t = match t with let rec size t = match t with
| Empty -> 0 | Empty -> 0
| Path (_, t') -> size t' | Cons (_, t') -> size t'
| Node (v, map) -> | Node (v, map) ->
let s = if v=None then 0 else 1 in let s = if v=None then 0 else 1 in
M.fold M.fold
@ -442,8 +404,7 @@ module Make(W : WORD) = struct
let _tree_node x l () = `Node (x,l) in let _tree_node x l () = `Node (x,l) in
match t with match t with
| Empty -> `Nil | Empty -> `Nil
| Path ([], _) -> assert false | Cons (c, t') -> `Node (`Char c, [to_tree t'])
| Path (c::l, t') -> `Node (`Char c, [to_tree (_mk_path l t')])
| Node (v, map) -> | Node (v, map) ->
let x = match v with let x = match v with
| None -> `Switch | None -> `Switch
@ -464,10 +425,9 @@ module Make(W : WORD) = struct
match cur with match cur with
| None -> (None, alternatives) | None -> (None, alternatives)
| Some (Empty,_) -> (None, alternatives) | Some (Empty,_) -> (None, alternatives)
| Some (Path ([], _),_) -> assert false | Some (Cons (c', t'), trail) ->
| Some (Path (c'::l, t'), trail) ->
if W.compare c c' = 0 if W.compare c c' = 0
then Some (_mk_path l t', _difflist_add trail c), alternatives then Some (t', _difflist_add trail c), alternatives
else None, alternatives else None, alternatives
| Some (Node (_, map), trail) -> | Some (Node (_, map), trail) ->
let alternatives = let alternatives =