extracted/untrusted/prioritySet.ml

   1 (* Pasted from Pottier's PP compiler *)
   2
   3 (* This module offers sets of elements where each element carries an
   4    integer priority. All operations execute in logarithmic time with
   5    respect to the number of elements in the set. *)
   6
   7 module Make (X : Set.OrderedType)
   8 = struct
   9
  10   (* First, define normal sets and maps. *)
  11
  12   module Set = Set.Make(X)
  13
  14   module Map = MyMap.Make(X)
  15
  16   (* Next, define maps of integers to nonempty sets of elements. *)
  17
  18   module IntMap = struct
  19
  20     module M = MyMap.Make (struct
  21       type t = int
  22       let compare = compare
  23     end)
  24
  25     include M
  26
  27     module H = SetMap.MakeHetero(Set)(M)
  28
  29     let update = H.update
  30
  31   end
  32
  33   (* Now, define priority sets. *)
  34
  35   type t = {
  36
  37       (* A mapping of elements to priorities. *)
  38
  39       priority: int Map.t;
  40
  41       (* A mapping of priorities to sets of elements. By convention, a
  42          priority has no entry in this table if that entry would be an
  43          empty set of elements. This allows finding the
  44          lowest-priority element in logarithmic time. *)
  45
  46       level: Set.t IntMap.t
  47
  48     }
  49
  50   (* [empty] is the empty set. *)
  51
  52   let empty =
  53     {
  54       priority = Map.empty;
  55       level = IntMap.empty
  56     }
  57
  58   (* [priority x s] looks up the priority of element [x]. *)
  59
  60   let priority x s =
  61     try
  62       Map.find x s.priority
  63     with Not_found ->
  64       assert false
  65
  66   (* [add x p s] inserts element [x] with priority [p]. *)
  67
  68   let add x p s =
  69     assert (not (Map.mem x s.priority));
  70     {
  71       priority = Map.add x p s.priority;
  72       level = IntMap.update p (Set.add x) s.level
  73     }
  74
  75   (* [remove x s] removes element [x]. *)
  76
  77   let remove x s =
  78     let p, priority =
  79       try
  80         Map.find_remove x s.priority
  81       with Not_found ->
  82         assert false
  83     in
  84     let level =
  85       IntMap.update p (function xs ->
  86         assert (Set.mem x xs);
  87         Set.remove x xs
  88       ) s.level
  89     in
  90     {
  91       priority = priority;
  92       level = level
  93     }
  94
  95   (* [change x p s] changes the priority of element [x] to [p]. *)
  96
  97   let change x p1 s =
  98     let p0 = priority x s in
  99     if p0 = p1 then
 100       s
 101     else
 102       {
 103         priority = Map.add x p1 s.priority; (* overriding previous entry *)
 104         level = IntMap.update p1 (Set.add x) (IntMap.update p0 (Set.remove x) s.level)
 105       }
 106
 107   (* [increment x d s] increases the priority of element [x] by [d]. *)
 108
 109   let increment x d s =
 110     change x (priority x s + d) s
 111
 112   (* [incrementifx x p s] increases the priority of element [x] by [d]
 113      if [x] is a member of the priority set. *)
 114
 115   let incrementifx x d s =
 116     if Map.mem x s.priority then
 117       increment x d s
 118     else
 119       s
 120
 121   (* [lowest s] returns [Some (x, p)], where element [x] has minimum
 122      priority [p] among all elements of [s]. It returns [None] if [s]
 123      is empty. *)
 124
 125   let lowest s =
 126     try
 127       let p, xs = IntMap.minimum s.level in (* can fail if set is empty *)
 128       try
 129         Some (Set.choose xs, p) (* cannot fail *)
 130       with Not_found ->
 131         assert false
 132     with Not_found ->
 133       None
 134
 135   (* [fold f s accu] fold over the set [s]. Elements are presented
 136      to [f] in increasing order of priority. *)
 137
 138   let fold f s accu =
 139     IntMap.fold (fun p xs accu ->
 140       Set.fold (fun x accu ->
 141         f x p accu
 142       ) xs accu
 143     ) s.level accu
 144
 145 end
 146