components/extlib/trie.ml

   1 (*
   2  * Trie: maps over lists.
   3  * Copyright (C) 2000 Jean-Christophe FILLIATRE
   4  *
   5  * This software is free software; you can redistribute it and/or
   6  * modify it under the terms of the GNU Library General Public
   7  * License version 2, as published by the Free Software Foundation.
   8  *
   9  * This software is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
  12  *
  13  * See the GNU Library General Public License version 2 for more details
  14  * (enclosed in the file LGPL).
  15  *)
  16
  17 (* $Id$ *)
  18
  19 (*s A trie is a tree-like structure to implement dictionaries over
  20     keys which have list-like structures. The idea is that each node
  21     branches on an element of the list and stores the value associated
  22     to the path from the root, if any. Therefore, a trie can be
  23     defined as soon as a map over the elements of the list is
  24     given. *)
  25
  26
  27 module Make (M : Map.S) = struct
  28
  29 (*s Then a trie is just a tree-like structure, where a possible
  30     information is stored at the node (['a option]) and where the sons
  31     are given by a map from type [key] to sub-tries, so of type
  32     ['a t M.t]. The empty trie is just the empty map. *)
  33
  34   type key = M.key list
  35
  36   type 'a t = Node of 'a option * 'a t M.t
  37
  38   let empty = Node (None, M.empty)
  39
  40 (*s To find a mapping in a trie is easy: when all the elements of the
  41     key have been read, we just inspect the optional info at the
  42     current node; otherwise, we descend in the appropriate sub-trie
  43     using [M.find]. *)
  44
  45   let rec find l t = match (l,t) with
  46     | [], Node (None,_)   -> raise Not_found
  47     | [], Node (Some v,_) -> v
  48     | x::r, Node (_,m)    -> find r (M.find x m)
  49
  50   let rec mem l t = match (l,t) with
  51     | [], Node (None,_)   -> false
  52     | [], Node (Some _,_) -> true
  53     | x::r, Node (_,m)    -> try mem r (M.find x m) with Not_found -> false
  54
  55 (*s Insertion is more subtle. When the final node is reached, we just
  56     put the information ([Some v]). Otherwise, we have to insert the
  57     binding in the appropriate sub-trie [t']. But it may not exists,
  58     and in that case [t'] is bound to an empty trie. Then we get a new
  59     sub-trie [t''] by a recursive insertion and we modify the
  60     branching, so that it now points to [t''], with [M.add]. *)
  61
  62   let add l v t =
  63     let rec ins = function
  64       | [], Node (_,m) -> Node (Some v,m)
  65       | x::r, Node (v,m) ->
  66           let t' = try M.find x m with Not_found -> empty in
  67           let t'' = ins (r,t') in
  68           Node (v, M.add x t'' m)
  69     in
  70     ins (l,t)
  71
  72 (*s When removing a binding, we take care of not leaving bindings to empty
  73     sub-tries in the nodes. Therefore, we test wether the result [t'] of
  74     the recursive call is the empty trie [empty]: if so, we just remove
  75     the branching with [M.remove]; otherwise, we modify it with [M.add]. *)
  76
  77   let rec remove l t = match (l,t) with
  78     | [], Node (_,m) -> Node (None,m)
  79     | x::r, Node (v,m) ->
  80         try
  81           let t' = remove r (M.find x m) in
  82           Node (v, if t' = empty then M.remove x m else M.add x t' m)
  83         with Not_found ->
  84           t
  85
  86 (*s The iterators [map], [mapi], [iter] and [fold] are implemented in
  87     a straigthforward way using the corresponding iterators [M.map],
  88     [M.mapi], [M.iter] and [M.fold]. For the last three of them,
  89     we have to remember the path from the root, as an extra argument
  90     [revp]. Since elements are pushed in reverse order in [revp],
  91     we have to reverse it with [List.rev] when the actual binding
  92     has to be passed to function [f]. *)
  93
  94   let rec map f = function
  95     | Node (None,m)   -> Node (None, M.map (map f) m)
  96     | Node (Some v,m) -> Node (Some (f v), M.map (map f) m)
  97
  98   let mapi f t =
  99     let rec maprec revp = function
 100     | Node (None,m) ->
 101         Node (None, M.mapi (fun x -> maprec (x::revp)) m)
 102     | Node (Some v,m) ->
 103         Node (Some (f (List.rev revp) v), M.mapi (fun x -> maprec (x::revp)) m)
 104     in
 105     maprec [] t
 106
 107   let iter f t =
 108     let rec traverse revp = function
 109       | Node (None,m) ->
 110           M.iter (fun x -> traverse (x::revp)) m
 111       | Node (Some v,m) ->
 112           f (List.rev revp) v; M.iter (fun x t -> traverse (x::revp) t) m
 113     in
 114     traverse [] t
 115
 116   let rec fold f t acc =
 117     let rec traverse revp t acc = match t with
 118       | Node (None,m) ->
 119           M.fold (fun x -> traverse (x::revp)) m acc
 120       | Node (Some v,m) ->
 121           f (List.rev revp) v (M.fold (fun x -> traverse (x::revp)) m acc)
 122     in
 123     traverse [] t acc
 124
 125   let compare cmp a b =
 126     let rec comp a b = match a,b with
 127       | Node (Some _, _), Node (None, _) -> 1
 128       | Node (None, _), Node (Some _, _) -> -1
 129       | Node (None, m1), Node (None, m2) ->
 130           M.compare comp m1 m2
 131       | Node (Some a, m1), Node (Some b, m2) ->
 132           let c = cmp a b in
 133           if c <> 0 then c else M.compare comp m1 m2
 134     in
 135     comp a b
 136
 137   let equal eq a b =
 138     let rec comp a b = match a,b with
 139       | Node (None, m1), Node (None, m2) ->
 140           M.equal comp m1 m2
 141       | Node (Some a, m1), Node (Some b, m2) ->
 142           eq a b && M.equal comp m1 m2
 143       | _ ->
 144           false
 145     in
 146     comp a b
 147
 148   (* The base case is rather stupid, but constructable *)
 149   let is_empty = function
 150     | Node (None, m1) -> M.is_empty m1
 151     | _ -> false
 152
 153 end