1 files changed, 19 insertions, 11 deletions

diff --git a/lib/trie.ml b/lib/trie.ml
index f6dc657cc..a4a348c19 100644
--- a/lib/trie.ml
+++ b/lib/trie.ml
@@ -12,7 +12,7 @@ sig
   type data
   type t
   val empty : t
-  val get : t -> data list
+  val get : t -> data
   val next : t -> label -> t
   val labels : t -> label list
   val add : label list -> data -> t -> t
@@ -20,25 +20,33 @@ sig
   val iter : (label list -> data -> unit) -> t -> unit
 end
 
-module Make (Y : Map.OrderedType) (X : Set.OrderedType) =
+module type Grp =
+sig
+  type t
+  val nil : t
+  val is_nil : t -> bool
+  val add : t -> t -> t
+  val sub : t -> t -> t
+end
+
+module Make (Y : Map.OrderedType) (X : Grp) =
 struct
 
-module T_dom = Set.Make(X)
 module T_codom = Map.Make(Y)
 
 type data = X.t
 type label = Y.t
-type t = Node of T_dom.t * t T_codom.t
+type t = Node of X.t * t T_codom.t
 
 let codom_for_all f m =
   let fold key v accu = f v && accu in
   T_codom.fold fold m true
 
-let empty = Node (T_dom.empty, T_codom.empty)
+let empty = Node (X.nil, T_codom.empty)
 
 let next (Node (_,m)) lbl = T_codom.find lbl m
 
-let get (Node (hereset,_)) = T_dom.elements hereset
+let get (Node (hereset,_)) = hereset
 
 let labels (Node (_,m)) =
   (** FIXME: this is order-dependent. Try to find a more robust presentation? *)
@@ -46,13 +54,13 @@ let labels (Node (_,m)) =
 
 let in_dom (Node (_,m)) lbl = T_codom.mem lbl m
 
-let is_empty_node (Node(a,b)) = (T_dom.is_empty a) && (T_codom.is_empty b)
+let is_empty_node (Node(a,b)) = (X.is_nil a) && (T_codom.is_empty b)
 
 let assure_arc m lbl =
   if T_codom.mem lbl m then
     m
   else
-    T_codom.add lbl (Node (T_dom.empty,T_codom.empty)) m
+    T_codom.add lbl (Node (X.nil,T_codom.empty)) m
 
 let cleanse_arcs (Node (hereset,m)) =
   let m = if codom_for_all is_empty_node m then T_codom.empty else m in
@@ -67,15 +75,15 @@ let rec at_path f (Node (hereset,m)) = function
                          T_codom.add h (at_path f (T_codom.find h m) t) m))
 
 let add path v tm =
-  at_path (fun hereset -> T_dom.add v hereset) tm path
+  at_path (fun hereset -> X.add v hereset) tm path
 
 let remove path v tm =
-  at_path (fun hereset -> T_dom.remove v hereset) tm path
+  at_path (fun hereset -> X.sub hereset v) tm path
 
 let iter f tlm =
   let rec apprec pfx (Node(hereset,m)) =
     let path = List.rev pfx in
-    T_dom.iter (fun v -> f path v) hereset;
+    f path hereset;
     T_codom.iter (fun l tm -> apprec (l::pfx) tm) m
   in
   apprec [] tlm