(* Copyright (C) 1990, Digital Equipment Corporation.         *)
(* All rights reserved.                                       *)
(* See the file COPYRIGHT for a full description.             *)

(* Last modified on Mon Jun 29 22:06:01 PDT 1992 by muller        *)
(*      modified on Thu Feb  8 16:52:28 1990 by mcjones       *)
(*      modified on Wed Feb  7 16:34:08 1990 by mhb           *)

MODULE STable EXPORTS STable, STableF;
IMPORT List;

  (* 2-3-4 trees, with top-down updates.
     References:
     Leo J. Guibas and Robert Sedgewick, "A dichromatic framework for
     balanced trees", in 19th Annual Symposium on Foundations of Computer
     Science, IEEE, 1978, pp 8-21.  Also in A Decade of Research 1970-1980,
     Xerox Palo Alto Research Center, Bowker, 1980.
     Robert Sedgewick, Algorithms, Addison-Wesley, 1983. *)

TYPE
  Link = REF Node;
  Node = RECORD
              key: Key;
              value: Value;
              red: BOOLEAN; (* color of link pointing to this node *)
              l, r: Link
            END;
REVEAL T = BRANDED REF
           RECORD
             h: Link;
             z: Link; (* shared external node *)
             y: Link; (* son of z *)
             compare: CompareProc;
             compareArg: REFANY
           END;

PROCEDURE Validate (table: T) RAISES {} =
    (* Check that no red node has a red parent, that every path from the root
       to an external node traverses the same number of black nodes, and that
       each node's key is greater than every key in its left subtree and less
       than every key in its right subtree. *)
  VAR z, y: Link; treeHeight: INTEGER;
  PROCEDURE V
    (x: Link; parentHeight: INTEGER; parentRed: BOOLEAN;
     lowerBound, upperBound: Key; checkLowerBound, checkUpperBound: BOOLEAN) =
    VAR height: INTEGER;
    BEGIN
      IF checkLowerBound THEN
        <* ASSERT (table.compare (table.compareArg, lowerBound, x.key) < 0) *>
      END;
      IF checkUpperBound THEN
        <* ASSERT (table.compare (table.compareArg, x.key, upperBound) < 0) *>
      END;
      IF x.red THEN
        <* ASSERT ( NOT parentRed) *>
        height := parentHeight;
      ELSE
        height := parentHeight + 1
      END;
      IF x.l = z THEN
        IF treeHeight =  -1 THEN
          treeHeight := height
        ELSE
          <* ASSERT (height = treeHeight) *>
        END
      ELSE
        V (x.l, height, x.red, lowerBound, x.key, checkLowerBound, TRUE)
      END;
      IF x.r = z THEN
        IF treeHeight =  -1 THEN
          treeHeight := height
        ELSE
          <* ASSERT (height = treeHeight) *>
        END
      ELSE
        V (x.r, height, x.red, x.key, upperBound, TRUE, checkUpperBound)
      END
    END V;
  BEGIN
    z := table.z;
    y := table.y;
    <* ASSERT NOT z.red AND (z.l = y) AND (z.r = y) AND (y.red)
              AND (table.h.l = z) *>
    treeHeight :=  -1;
    IF table.h.r # z THEN V (table.h.r,  -1, TRUE, NIL, NIL, FALSE, FALSE) END
  END Validate;

PROCEDURE New (compare: CompareProc; compareArg: REFANY := NIL): T RAISES {} =
  VAR table: T;
  BEGIN
    table :=  NEW (T);
    table.h :=  NEW (Link);
    table.z :=  NEW (Link);
    table.y :=  NEW (Link);
    table.compare := compare;
    table.compareArg := compareArg;
    WITH z_10 = table.h^ DO
      z_10.key := NIL;
      z_10.red := FALSE;
      z_10.l := table.z;
      z_10.r := table.z
    END;
    WITH z_11 = table.z^ DO
              (* key varies *)
      z_11.red := FALSE;
      z_11.l := table.y;
      z_11.r := table.y
    END;
    table.y.red := TRUE;
    RETURN table
  END New;

PROCEDURE Get (table: T; key: Key; VAR value: Value): BOOLEAN RAISES {} =
  VAR x, z: Link; cmp: INTEGER;
  BEGIN
    z := table.z;
    x := table.h.r;
    LOOP
      IF x = z THEN EXIT END;
      cmp := table.compare (table.compareArg, key, x.key);
      IF cmp = 0 THEN EXIT END;
      IF cmp < 0 THEN x := x.l ELSE x := x.r END
    END;
    IF x # z THEN value := x.value; RETURN TRUE ELSE RETURN FALSE END
  END Get;

PROCEDURE Put (table: T; key: Key; value: Value): BOOLEAN RAISES {} =
    (* Based on Sedgewick's rbtreeinsert. *)
  VAR
    x: Link;
    cmp: INTEGER; (* = compare(compareArg, key, x^.key) *)
    f: Link; (* father of x *)
    g: Link; (* grandfather of x *)
    gg: Link; (* great-grandfather of x *)

  PROCEDURE Split () =
    PROCEDURE Rotate (y: Link): Link =
      VAR s: Link; (* son of y *)gs: Link; (* grandson of y *)yLeft: BOOLEAN;
      BEGIN
        yLeft := table.compare (table.compareArg, key, y.key) < 0;
        IF yLeft THEN s := y.l ELSE s := y.r END;
        IF table.compare (table.compareArg, key, s.key) < 0 THEN
          gs := s.l;
          s.l := gs.r;
          gs.r := s
        ELSE
          gs := s.r;
          s.r := gs.l;
          gs.l := s
        END;
        IF yLeft THEN y.l := gs ELSE y.r := gs END;
        RETURN gs
      END Rotate;
    BEGIN
      x.red := TRUE;
      x.l.red := FALSE;
      x.r.red := FALSE;
      IF f.red THEN
        g.red := TRUE;
        IF (table.compare (table.compareArg, key, g.key) < 0)
          # (table.compare (table.compareArg, key, f.key) < 0) THEN
          f := Rotate (g)
        END;
        x := Rotate (gg);
        cmp := table.compare (table.compareArg, key, x.key);
        x.red := FALSE
      END;
      table.h.r.red := FALSE
    END Split;
  BEGIN
    x := table.h;
    cmp := 0; (* i.e., not < 0 *)
    table.z.key := key;
    f := x;
    g := x;
    LOOP
      gg := g;
      g := f;
      f := x;
      IF cmp < 0 THEN x := x.l ELSE x := x.r END;
      cmp := table.compare (table.compareArg, key, x.key);
      IF cmp = 0 THEN EXIT END;
      IF x.l.red AND x.r.red THEN Split () END
    END;
    IF x = table.z THEN
      x :=  NEW (Link);
      x.key := key;
      x.value := value;
      x.l := table.z;
      x.r := table.z;
      IF table.compare (table.compareArg, key, f.key) < 0 THEN
        f.l := x
      ELSE
        f.r := x
      END;
      Split ();
      RETURN FALSE
    ELSE
      x.value := value;
      RETURN TRUE
    END
  END Put;

PROCEDURE Delete (table: T; key: Key; VAR value: Value): BOOLEAN RAISES {} =
    (* Based on Program 8 of Guibas and Sedgewick. *)
  PROCEDURE Rotate (a, b, c: Link) =
        (* Assumes c is a son of b and b is a son of a.  Rotates b and c,
           updating a's son field. *)
    BEGIN
      IF table.compare (table.compareArg, c.key, b.key) < 0
        THEN
        b.l := c.r;
        c.r := b
      ELSE
        b.r := c.l;
        c.l := b
      END;
      IF table.compare (table.compareArg, b.key, a.key) < 0
        THEN
        a.l := c
      ELSE
        a.r := c
      END
    END Rotate;
  VAR
    x: Link;
    f: Link; (* father of x *)
    cmp: INTEGER; (* = compare(compareArg, key, f^.key) *)
    b: Link; (* brother of x *)
    g: Link; (* grandfather of x *)
    t: Link; (* node to be deleted *)
    n: Link; (* nephew of x *)
    z: Link;
  BEGIN
    z := table.z;
    t := NIL;
    f := table.h;
    x := f.r; (* x = root of tree *)
    b := z; (* b = brother of x = z *)
          (* Initially, the root of the tree is black.  If both of its sons are
             black, we can make the root red. *)
    IF  NOT x.l.red AND  NOT x.r.red THEN x.red := TRUE END;
          (* Now, either x or a son of x is red. *)
    LOOP
              (* INVARIANT: either x or f or b or a son of x is red. *)
      IF (t = NIL)
        AND (table.compare (table.compareArg, key, x.key) = 0) THEN
        t := x
      END;
      IF  NOT x.red AND  NOT f.red THEN
        (*
                  Since neither x nor f is red, either b or a son of x is
                  red.  In the former case, we rotate f and b to make f red;
                  x and its new brother will then be black:
                                   f|             g'|
                                    2               4
                                   / \             / \
                                 x/   \b   ==>   f/   \
                                 1     4#        2#    5
                                      / \      x/ \b'
                                     3   5     1   3
                                    # marks red nodes
        *)
        IF b.red THEN
          f.red := TRUE;
          b.red := FALSE;
          Rotate (g, f, b);
          g := b;
          IF cmp < 0 THEN b := f.r ELSE b := f.l END
        ELSE
                      (* A son of x is red. *)
        END
      END;
              (* Now, either x, f, or a son of x is red.  If x = z, we can
                 terminate: z is black, so f (the node to be removed) is red. *)
      IF x = z THEN EXIT END;
              (* To prepare for the next iteration, we must establish x or a son as
                 red. *)
      IF  NOT x.red AND  NOT x.l.red AND  NOT x.r.red THEN
                  (* Since x is black, f must be red, and b must be black.  If both
                     b's sons are black, we can simply flip the colors of f, x, and
                     b. But if one of b's sons is red, we must perform one or two
                     rotations.  In any case, we up with x red and f black. *)
        x.red := TRUE;
        f.red := FALSE;
                  (* Set n to the son of b farther from x. *)
        IF cmp < 0 THEN n := b.r ELSE n := b.l END;
        IF n.red THEN
          (*
                      A color-swap of b and its red far son, followed by a
                      rotation of f and b, will restore balance:
                                     f|             g'|
                                      2#              4#
                                     / \             / \
                                   x/   \b   ==>   f/   \
                                   1     4         2     5
                                        / \n     x/ \
                                       3?  5#    1#  3?
                                      # marks red nodes
                                      ? marks don't care nodes
          *)
          b.red := TRUE;
          n.red := FALSE;
          Rotate (g, f, b);
          g := b
        ELSE
                      (* Set n to son of b nearer to x. *)
          IF cmp < 0 THEN n := b.l ELSE n := b.r END;
          IF n.red THEN
            (*
                        Only the son of b near x is red.  A rotation of b and
                        its red son, followed by a rotation of f and its new red
                        son, will restore balance:
                                       f|             g'|
                                        2#              3#
                                       / \             / \
                                     x/   \b   ==>   f/   \
                                     1     4         2     4
                                         n/ \       / \   / \
                                         3#  5    1# 2.5 3.5 5
                                        / \
                                      2.5 3.5
                                        # marks red nodes
            *)
            Rotate (f, b, n);
            Rotate (g, f, n);
            g := n
          ELSE
            (*
                          Neither son of b is red; the color flip of f, x, and
                          b can be performed:
                                       f|              f|
                                        2#              2
                                       / \             / \
                                     x/   \b   ==>   f/   \b
                                     1     4         1#    4#
                                        # marks red nodes
            *)
            b.red := TRUE
          END
        END
      END;
              (* Now x or a son is red. *)
      g := f;
      f := x;
      cmp := table.compare (table.compareArg, key, f.key);
      IF cmp < 0 THEN
        x := f.l;
        b := f.r
      ELSE
        x := f.r;
        b := f.l
      END
    END; (* LOOP *)
    table.h.r.red := FALSE;
    IF t = NIL THEN (* not present *)
      RETURN FALSE
    ELSE (* present *)
      value := t.value;
      IF table.compare (table.compareArg, key, g.key) < 0 THEN
        g.l := z
      ELSE
        g.r := z
      END;
      t.key := f.key;
      t.value := f.value;
      RETURN TRUE
    END
  END Delete;

PROCEDURE Clear (table: T) RAISES {} =
  BEGIN
    table.h.r := table.z
  END Clear;

PROCEDURE Copy (tableOld: T): T RAISES {} =
  VAR table: T;
  PROCEDURE CopyNode (xOld: Link): Link =
    VAR x: Link;
    BEGIN
      IF xOld = tableOld.z THEN
        RETURN table.z
      ELSE
        x :=  NEW (Link);
        WITH z_12 = x^ DO
          z_12.key := xOld.key;
          z_12.value := xOld.value;
          z_12.red := xOld.red;
          z_12.l := CopyNode (xOld.l);
          z_12.r := CopyNode (xOld.r)
        END;
        RETURN x
      END
    END CopyNode;
  BEGIN
    table :=  NEW (T);
    WITH z_13 = table^ DO
      z_13.h :=  NEW (Link);
      z_13.z :=  NEW (Link);
      z_13.y :=  NEW (Link);
      WITH z_14 = z_13.h^ DO
        z_14.key := tableOld.h.key;
        z_14.red := FALSE;
        z_14.l := z_13.z;
        z_14.r := CopyNode (tableOld.h.r)
      END;
      WITH z_15 = z_13.z^ DO
                  (* key varies *)
        z_15.red := FALSE;
        z_15.l := z_13.y;
        z_15.r := z_13.y
      END;
      z_13.y.red := TRUE;
      z_13.compare := tableOld.compare;
      z_13.compareArg := tableOld.compareArg
    END;
    RETURN table
  END Copy;
CONST
  MaxStack = 50; (* allows tree with 2**(50/2) = 33,554,432 entries *)


REVEAL Stream = BRANDED REF RECORD
               z: Link;
               compare: CompareProc;
               compareArg: REFANY;
               start: Key;
               ascend: BOOLEAN;
               top: CARDINAL;
               stack: ARRAY [0..MaxStack] OF Link
             END;


PROCEDURE NewStream
  (table: T; ascending: BOOLEAN := TRUE; key0: Key := NIL): Stream RAISES {} =
  VAR s: Stream;
  BEGIN
    s :=  NEW (Stream);
    WITH z_16 = s^ DO
      z_16.z := table.z;
      z_16.compare := table.compare;
      z_16.compareArg := table.compareArg;
      z_16.start := key0;
      z_16.ascend := ascending;
      z_16.top := 0;
      IF ascending THEN
        z_16.stack[z_16.top] := table.h
      ELSE
                  (* Fake node at +infinity. *)
        z_16.stack[z_16.top] :=  NEW (Link);
        z_16.stack[z_16.top].l := table.h.r
      END
    END;
    RETURN s
  END NewStream;

PROCEDURE Next
  (s: Stream; VAR key: Key; VAR value: Value): BOOLEAN RAISES {} =
  VAR
    x: Link; (* = stack[top] *)
    cmp: INTEGER;
  BEGIN
    WITH z_17 = s^ DO
      x := z_17.stack[z_17.top];
      LOOP
        IF z_17.ascend THEN
          IF x.r = z_17.z THEN
            IF z_17.top = 0 THEN
              RETURN FALSE
            ELSE
              DEC (z_17.top);
              x := z_17.stack[z_17.top];
              key := x.key;
              value := x.value;
              RETURN TRUE
            END
          END;
          x := x.r;
          z_17.stack[z_17.top] := x;
          LOOP
            cmp := z_17.compare (z_17.compareArg, z_17.start, x.key);
            IF (x.l = z_17.z) OR (cmp >= 0) THEN EXIT END;
            x := x.l;
            INC (z_17.top);
            z_17.stack[z_17.top] := x
          END;
          IF cmp <= 0 THEN
            key := x.key;
            value := x.value;
            RETURN TRUE
          END
        ELSE (* NOT ascend *)
          IF x.l = z_17.z THEN
            IF z_17.top = 0 THEN
              RETURN FALSE
            ELSE
              DEC (z_17.top);
              x := z_17.stack[z_17.top];
              key := x.key;
              value := x.value;
              RETURN TRUE
            END
          END;
          x := x.l;
          z_17.stack[z_17.top] := x;
          LOOP
            IF z_17.start # NIL THEN
              cmp := z_17.compare (z_17.compareArg, z_17.start, x.key)
            ELSE (* treat as +Infinity *)
              cmp := 1
            END;
            IF (x.r = z_17.z) OR (cmp <=0) THEN EXIT END;
            x := x.r;
            INC (z_17.top);
            z_17.stack[z_17.top] := x
          END;
          IF cmp >= 0 THEN
            key := x.key;
            value := x.value;
            RETURN TRUE
          END
        END
      END
    END
  END Next;

PROCEDURE FindFirst
  (table: T; pred: Predicate; arg: REFANY; VAR key: Key; VAR value: Value;
   ascending: BOOLEAN := TRUE; key0: Key := NIL)
   : BOOLEAN(* RAISES {raises( pred )} *) =
  VAR k: Key; s: Stream; v: Value;
  BEGIN
    s := NewStream (table, ascending, key0);
    LOOP
      IF  NOT Next (s, k, v) THEN RETURN FALSE END;
      IF pred (arg, k, v) THEN key := k; value := v; RETURN TRUE END
    END
  END FindFirst;

PROCEDURE ToValuesList
  (table: T; ascending: BOOLEAN := TRUE): List.T RAISES {} =
  VAR k: Key; l: List.T; s: Stream; v: Value;
  BEGIN
    s := NewStream (table,  NOT ascending);
    l := NIL;
    WHILE Next (s, k, v) DO List.Push (l, v) END;
    RETURN l
  END ToValuesList;

PROCEDURE ToAssocList
  (table: T; ascending: BOOLEAN := TRUE): List.T RAISES {} =
  VAR k: Key; l: List.T; s: Stream; v: Value;
  BEGIN
    s := NewStream (table,  NOT ascending);
    l := NIL;
    WHILE Next (s, k, v) DO List.Push (l, List.List2 (k, v)) END;
    RETURN l
  END ToAssocList;

BEGIN
END STable.