src/Tools/Metis/src/RandomSet.sml

merged

(* ========================================================================= *)
(* FINITE SETS IMPLEMENTED WITH RANDOMLY BALANCED TREES                      *)
(* Copyright (c) 2004-2006 Joe Hurd, distributed under the BSD License *)
(* ========================================================================= *)

structure RandomSet :> Set =
struct

exception Bug = Useful.Bug;

exception Error = Useful.Error;

val pointerEqual = Portable.pointerEqual;

val K = Useful.K;

val snd = Useful.snd;

val randomInt = Portable.randomInt;

val randomWord = Portable.randomWord;

(* ------------------------------------------------------------------------- *)
(* Random search trees.                                                      *)
(* ------------------------------------------------------------------------- *)

type priority = Word.word;

datatype 'a tree =
    E
  | T of
    {size : int,
     priority : priority,
     left : 'a tree,
     key : 'a,
     right : 'a tree};
    
type 'a node =
     {size : int,
      priority : priority,
      left : 'a tree,
      key : 'a,
      right : 'a tree};

datatype 'a set = Set of ('a * 'a -> order) * 'a tree;

(* ------------------------------------------------------------------------- *)
(* Random priorities.                                                        *)
(* ------------------------------------------------------------------------- *)

local
  val randomPriority = randomWord;

  val priorityOrder = Word.compare;
in
  fun treeSingleton key =
      T {size = 1, priority = randomPriority (),
         left = E, key = key, right = E};

  fun nodePriorityOrder cmp (x1 : 'a node, x2 : 'a node) =
      let
        val {priority = p1, key = k1, ...} = x1
        and {priority = p2, key = k2, ...} = x2
      in
        case priorityOrder (p1,p2) of
          LESS => LESS
        | EQUAL => cmp (k1,k2)
        | GREATER => GREATER
      end;
end;

(* ------------------------------------------------------------------------- *)
(* Debugging functions.                                                      *)
(* ------------------------------------------------------------------------- *)

local
  fun checkSizes E = 0
    | checkSizes (T {size,left,right,...}) =
      let
        val l = checkSizes left
        and r = checkSizes right
        val () = if l + 1 + r = size then () else raise Error "wrong size"
      in
        size
      end

  fun checkSorted _ x E = x
    | checkSorted cmp x (T {left,key,right,...}) =
      let
        val x = checkSorted cmp x left
        val () =
            case x of
              NONE => ()
            | SOME k =>
              case cmp (k,key) of
                LESS => ()
              | EQUAL => raise Error "duplicate keys"
              | GREATER => raise Error "unsorted"
      in
        checkSorted cmp (SOME key) right
      end;

  fun checkPriorities _ E = NONE
    | checkPriorities cmp (T (x as {left,right,...})) =
      let
        val () =
            case checkPriorities cmp left of
              NONE => ()
            | SOME l =>
              case nodePriorityOrder cmp (l,x) of
                LESS => ()
              | EQUAL => raise Error "left child has equal key"
              | GREATER => raise Error "left child has greater priority"
        val () =
            case checkPriorities cmp right of
              NONE => ()
            | SOME r =>
              case nodePriorityOrder cmp (r,x) of
                LESS => ()
              | EQUAL => raise Error "right child has equal key"
              | GREATER => raise Error "right child has greater priority"
      in
        SOME x
      end;
in
  fun checkWellformed s (set as Set (cmp,tree)) =
      (let
         val _ = checkSizes tree
         val _ = checkSorted cmp NONE tree
         val _ = checkPriorities cmp tree
       in
         set
       end
       handle Error err => raise Bug err)
      handle Bug bug => raise Bug (s ^ "\nRandomSet.checkWellformed: " ^ bug);
end;

(* ------------------------------------------------------------------------- *)
(* Basic operations.                                                         *)
(* ------------------------------------------------------------------------- *)

fun comparison (Set (cmp,_)) = cmp;

fun empty cmp = Set (cmp,E);

fun treeSize E = 0
  | treeSize (T {size = s, ...}) = s;

fun size (Set (_,tree)) = treeSize tree;

fun mkT p l k r =
    T {size = treeSize l + 1 + treeSize r, priority = p,
       left = l, key = k, right = r};

fun singleton cmp key = Set (cmp, treeSingleton key);

local
  fun treePeek cmp E pkey = NONE
    | treePeek cmp (T {left,key,right,...}) pkey =
      case cmp (pkey,key) of
        LESS => treePeek cmp left pkey
      | EQUAL => SOME key
      | GREATER => treePeek cmp right pkey
in
  fun peek (Set (cmp,tree)) key = treePeek cmp tree key;
end;

(* treeAppend assumes that every element of the first tree is less than *)
(* every element of the second tree. *)

fun treeAppend _ t1 E = t1
  | treeAppend _ E t2 = t2
  | treeAppend cmp (t1 as T x1) (t2 as T x2) =
    case nodePriorityOrder cmp (x1,x2) of
      LESS =>
      let
        val {priority = p2, left = l2, key = k2, right = r2, ...} = x2
      in
        mkT p2 (treeAppend cmp t1 l2) k2 r2
      end
    | EQUAL => raise Bug "RandomSet.treeAppend: equal keys"
    | GREATER =>
      let
        val {priority = p1, left = l1, key = k1, right = r1, ...} = x1
      in
        mkT p1 l1 k1 (treeAppend cmp r1 t2)
      end;

(* nodePartition splits the node into three parts: the keys comparing less *)
(* than the supplied key, an optional equal key, and the keys comparing *)
(* greater. *)

local
  fun mkLeft [] t = t
    | mkLeft (({priority,left,key,...} : 'a node) :: xs) t =
      mkLeft xs (mkT priority left key t);

  fun mkRight [] t = t
    | mkRight (({priority,key,right,...} : 'a node) :: xs) t =
      mkRight xs (mkT priority t key right);

  fun treePart _ _ lefts rights E = (mkLeft lefts E, NONE, mkRight rights E)
    | treePart cmp pkey lefts rights (T x) = nodePart cmp pkey lefts rights x
  and nodePart cmp pkey lefts rights (x as {left,key,right,...}) =
      case cmp (pkey,key) of
        LESS => treePart cmp pkey lefts (x :: rights) left
      | EQUAL => (mkLeft lefts left, SOME key, mkRight rights right)
      | GREATER => treePart cmp pkey (x :: lefts) rights right;
in
  fun nodePartition cmp x pkey = nodePart cmp pkey [] [] x;
end;

(* union first calls treeCombineRemove, to combine the values *)
(* for equal keys into the first map and remove them from the second map. *)
(* Note that the combined key is always the one from the second map. *)

local
  fun treeCombineRemove _ t1 E = (t1,E)
    | treeCombineRemove _ E t2 = (E,t2)
    | treeCombineRemove cmp (t1 as T x1) (t2 as T x2) =
      let
        val {priority = p1, left = l1, key = k1, right = r1, ...} = x1
        val (l2,k2,r2) = nodePartition cmp x2 k1
        val (l1,l2) = treeCombineRemove cmp l1 l2
        and (r1,r2) = treeCombineRemove cmp r1 r2
      in
        case k2 of
          NONE => if treeSize l2 + treeSize r2 = #size x2 then (t1,t2)
                  else (mkT p1 l1 k1 r1, treeAppend cmp l2 r2)
        | SOME k2 => (mkT p1 l1 k2 r1, treeAppend cmp l2 r2)
      end;

  fun treeUnionDisjoint _ t1 E = t1
    | treeUnionDisjoint _ E t2 = t2
    | treeUnionDisjoint cmp (T x1) (T x2) =
      case nodePriorityOrder cmp (x1,x2) of
        LESS => nodeUnionDisjoint cmp x2 x1
      | EQUAL => raise Bug "RandomSet.unionDisjoint: equal keys"
      | GREATER => nodeUnionDisjoint cmp x1 x2

  and nodeUnionDisjoint cmp x1 x2 =
      let
        val {priority = p1, left = l1, key = k1, right = r1, ...} = x1
        val (l2,_,r2) = nodePartition cmp x2 k1
        val l = treeUnionDisjoint cmp l1 l2
        and r = treeUnionDisjoint cmp r1 r2
      in
        mkT p1 l k1 r
      end;
in
  fun union (s1 as Set (cmp,t1)) (Set (_,t2)) =
      if pointerEqual (t1,t2) then s1
      else
        let
          val (t1,t2) = treeCombineRemove cmp t1 t2
        in
          Set (cmp, treeUnionDisjoint cmp t1 t2)
        end;
end;

(*DEBUG
val union = fn t1 => fn t2 =>
    checkWellformed "RandomSet.union: result"
      (union (checkWellformed "RandomSet.union: input 1" t1)
             (checkWellformed "RandomSet.union: input 2" t2));
*)

(* intersect is a simple case of the union algorithm. *)

local
  fun treeIntersect _ _ E = E
    | treeIntersect _ E _ = E
    | treeIntersect cmp (t1 as T x1) (t2 as T x2) =
      let
        val {priority = p1, left = l1, key = k1, right = r1, ...} = x1
        val (l2,k2,r2) = nodePartition cmp x2 k1
        val l = treeIntersect cmp l1 l2
        and r = treeIntersect cmp r1 r2
      in
        case k2 of
          NONE => treeAppend cmp l r
        | SOME k2 => mkT p1 l k2 r
      end;
in
  fun intersect (s1 as Set (cmp,t1)) (Set (_,t2)) =
      if pointerEqual (t1,t2) then s1
      else Set (cmp, treeIntersect cmp t1 t2);
end;

(*DEBUG
val intersect = fn t1 => fn t2 =>
    checkWellformed "RandomSet.intersect: result"
      (intersect (checkWellformed "RandomSet.intersect: input 1" t1)
                 (checkWellformed "RandomSet.intersect: input 2" t2));
*)

(* delete raises an exception if the supplied key is not found, which *)
(* makes it simpler to maximize sharing. *)

local
  fun treeDelete _ E _ = raise Error "RandomSet.delete: element not found"
    | treeDelete cmp (T {priority,left,key,right,...}) dkey =
      case cmp (dkey,key) of
        LESS => mkT priority (treeDelete cmp left dkey) key right
      | EQUAL => treeAppend cmp left right
      | GREATER => mkT priority left key (treeDelete cmp right dkey);
in
  fun delete (Set (cmp,tree)) key = Set (cmp, treeDelete cmp tree key);
end;
 
(*DEBUG
val delete = fn t => fn x =>
    checkWellformed "RandomSet.delete: result"
      (delete (checkWellformed "RandomSet.delete: input" t) x);
*)

(* Set difference *)

local
  fun treeDifference _ t1 E = t1
    | treeDifference _ E _ = E
    | treeDifference cmp (t1 as T x1) (T x2) =
      let
        val {size = s1, priority = p1, left = l1, key = k1, right = r1} = x1
        val (l2,k2,r2) = nodePartition cmp x2 k1
        val l = treeDifference cmp l1 l2
        and r = treeDifference cmp r1 r2
      in
        if Option.isSome k2 then treeAppend cmp l r
        else if treeSize l + treeSize r + 1 = s1 then t1
        else mkT p1 l k1 r
      end;
in
  fun difference (Set (cmp,tree1)) (Set (_,tree2)) =
      if pointerEqual (tree1,tree2) then Set (cmp,E)
      else Set (cmp, treeDifference cmp tree1 tree2);
end;

(*DEBUG
val difference = fn t1 => fn t2 =>
    checkWellformed "RandomSet.difference: result"
      (difference (checkWellformed "RandomSet.difference: input 1" t1)
                  (checkWellformed "RandomSet.difference: input 2" t2));
*)

(* Subsets *)

local
  fun treeSubset _ E _ = true
    | treeSubset _ _ E = false
    | treeSubset cmp (t1 as T x1) (T x2) =
      let
        val {size = s1, left = l1, key = k1, right = r1, ...} = x1
        and {size = s2, ...} = x2
      in
        s1 <= s2 andalso
        let
          val (l2,k2,r2) = nodePartition cmp x2 k1
        in
          Option.isSome k2 andalso
          treeSubset cmp l1 l2 andalso
          treeSubset cmp r1 r2
        end
      end;
in
  fun subset (Set (cmp,tree1)) (Set (_,tree2)) =
      pointerEqual (tree1,tree2) orelse
      treeSubset cmp tree1 tree2;
end;

(* Set equality *)

local
  fun treeEqual _ E E = true
    | treeEqual _ E _ = false
    | treeEqual _ _ E = false
    | treeEqual cmp (t1 as T x1) (T x2) =
      let
        val {size = s1, left = l1, key = k1, right = r1, ...} = x1
        and {size = s2, ...} = x2
      in
        s1 = s2 andalso
        let
          val (l2,k2,r2) = nodePartition cmp x2 k1
        in
          Option.isSome k2 andalso
          treeEqual cmp l1 l2 andalso
          treeEqual cmp r1 r2
        end
      end;
in
  fun equal (Set (cmp,tree1)) (Set (_,tree2)) =
      pointerEqual (tree1,tree2) orelse
      treeEqual cmp tree1 tree2;
end;

(* filter is the basic function for preserving the tree structure. *)

local
  fun treeFilter _ _ E = E
    | treeFilter cmp pred (T {priority,left,key,right,...}) =
      let
        val left = treeFilter cmp pred left
        and right = treeFilter cmp pred right
      in
        if pred key then mkT priority left key right
        else treeAppend cmp left right
      end;
in
  fun filter pred (Set (cmp,tree)) = Set (cmp, treeFilter cmp pred tree);
end;

(* nth picks the nth smallest key (counting from 0). *)

local
  fun treeNth E _ = raise Subscript
    | treeNth (T {left,key,right,...}) n =
      let
        val k = treeSize left
      in
        if n = k then key
        else if n < k then treeNth left n
        else treeNth right (n - (k + 1))
      end;
in
  fun nth (Set (_,tree)) n = treeNth tree n;
end;

(* ------------------------------------------------------------------------- *)
(* Iterators.                                                                *)
(* ------------------------------------------------------------------------- *)

fun leftSpine E acc = acc
  | leftSpine (t as T {left,...}) acc = leftSpine left (t :: acc);

fun rightSpine E acc = acc
  | rightSpine (t as T {right,...}) acc = rightSpine right (t :: acc);

datatype 'a iterator =
    LR of 'a * 'a tree * 'a tree list
  | RL of 'a * 'a tree * 'a tree list;

fun mkLR [] = NONE
  | mkLR (T {key,right,...} :: l) = SOME (LR (key,right,l))
  | mkLR (E :: _) = raise Bug "RandomSet.mkLR";

fun mkRL [] = NONE
  | mkRL (T {key,left,...} :: l) = SOME (RL (key,left,l))
  | mkRL (E :: _) = raise Bug "RandomSet.mkRL";

fun mkIterator (Set (_,tree)) = mkLR (leftSpine tree []);

fun mkRevIterator (Set (_,tree)) = mkRL (rightSpine tree []);

fun readIterator (LR (key,_,_)) = key
  | readIterator (RL (key,_,_)) = key;

fun advanceIterator (LR (_,next,l)) = mkLR (leftSpine next l)
  | advanceIterator (RL (_,next,l)) = mkRL (rightSpine next l);

(* ------------------------------------------------------------------------- *)
(* Derived operations.                                                       *)
(* ------------------------------------------------------------------------- *)

fun null s = size s = 0;

fun member x s = Option.isSome (peek s x);

fun add s x = union s (singleton (comparison s) x);

(*DEBUG
val add = fn s => fn x =>
    checkWellformed "RandomSet.add: result"
      (add (checkWellformed "RandomSet.add: input" s) x);
*)

local
  fun unionPairs ys [] = rev ys
    | unionPairs ys (xs as [_]) = List.revAppend (ys,xs)
    | unionPairs ys (x1 :: x2 :: xs) = unionPairs (union x1 x2 :: ys) xs;
in
  fun unionList [] = raise Error "RandomSet.unionList: no sets"
    | unionList [s] = s
    | unionList l = unionList (unionPairs [] l);
end;

local
  fun intersectPairs ys [] = rev ys
    | intersectPairs ys (xs as [_]) = List.revAppend (ys,xs)
    | intersectPairs ys (x1 :: x2 :: xs) =
      intersectPairs (intersect x1 x2 :: ys) xs;
in
  fun intersectList [] = raise Error "RandomSet.intersectList: no sets"
    | intersectList [s] = s
    | intersectList l = intersectList (intersectPairs [] l);
end;

fun symmetricDifference s1 s2 = union (difference s1 s2) (difference s2 s1);

fun disjoint s1 s2 = null (intersect s1 s2);

fun partition pred set = (filter pred set, filter (not o pred) set);

local
  fun fold _ NONE acc = acc
    | fold f (SOME iter) acc =
      let
        val key = readIterator iter
      in
        fold f (advanceIterator iter) (f (key,acc))
      end;
in
  fun foldl f b m = fold f (mkIterator m) b;

  fun foldr f b m = fold f (mkRevIterator m) b;
end;

local
  fun find _ NONE = NONE
    | find pred (SOME iter) =
      let
        val key = readIterator iter
      in
        if pred key then SOME key
        else find pred (advanceIterator iter)
      end;
in
  fun findl p m = find p (mkIterator m);

  fun findr p m = find p (mkRevIterator m);
end;

local
  fun first _ NONE = NONE
    | first f (SOME iter) =
      let
        val key = readIterator iter
      in
        case f key of
          NONE => first f (advanceIterator iter)
        | s => s
      end;
in
  fun firstl f m = first f (mkIterator m);

  fun firstr f m = first f (mkRevIterator m);
end;

fun count p = foldl (fn (x,n) => if p x then n + 1 else n) 0;

fun fromList cmp l = List.foldl (fn (k,s) => add s k) (empty cmp) l;

fun addList s l = union s (fromList (comparison s) l);

fun toList s = foldr op:: [] s;

fun map f s = rev (foldl (fn (x,l) => (x, f x) :: l) [] s);

fun transform f s = rev (foldl (fn (x,l) => f x :: l) [] s);

fun app f s = foldl (fn (x,()) => f x) () s;

fun exists p s = Option.isSome (findl p s);

fun all p s = not (exists (not o p) s);

local
  fun iterCompare _ NONE NONE = EQUAL
    | iterCompare _ NONE (SOME _) = LESS
    | iterCompare _ (SOME _) NONE = GREATER
    | iterCompare cmp (SOME i1) (SOME i2) =
      keyIterCompare cmp (readIterator i1) (readIterator i2) i1 i2

  and keyIterCompare cmp k1 k2 i1 i2 =
      case cmp (k1,k2) of
        LESS => LESS
      | EQUAL => iterCompare cmp (advanceIterator i1) (advanceIterator i2)
      | GREATER => GREATER;
in
  fun compare (s1,s2) =
      if pointerEqual (s1,s2) then EQUAL
      else
        case Int.compare (size s1, size s2) of
          LESS => LESS
        | EQUAL => iterCompare (comparison s1) (mkIterator s1) (mkIterator s2)
        | GREATER => GREATER;
end;

fun pick s =
    case findl (K true) s of
      SOME p => p
    | NONE => raise Error "RandomSet.pick: empty";

fun random s =
    case size s of
      0 => raise Error "RandomSet.random: empty"
    | n => nth s (randomInt n);

fun deletePick s = let val x = pick s in (x, delete s x) end;

fun deleteRandom s = let val x = random s in (x, delete s x) end;

fun close f s = let val s' = f s in if equal s s' then s else close f s' end;

fun toString s = "{" ^ (if null s then "" else Int.toString (size s)) ^ "}";

end