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

(* Last modified on Tue Apr 24 23:06:07 1990 by muller        *)
(*      modified on Fri Mar  9 10:13:45 1990 by mcjones       *)

MODULE SIntTable EXPORTS SIntTable, SIntTableF;

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: REFANY;
           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 *)
        END;

PROCEDURE Validate (table: T) RAISES {} =
    (* Checks 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; min, max: Key) =
    VAR height: INTEGER;
    BEGIN
      <* ASSERT ((min <= x.key) AND (x.key <= max)) *>
      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, min, x.key)
      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, max)
      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, FIRST (Key), LAST (Key))
    END
  END Validate;

PROCEDURE New (): T RAISES {} =
  VAR table: T;
  BEGIN
    table :=  NEW (T);
    WITH z_0 = table^ DO
      z_0.y :=  NEW (Link);
      WITH z_1 = z_0.y^ DO z_1.red := TRUE END;
      z_0.z :=  NEW (Link);
      WITH z_2 = z_0.z^ DO
                  (* key only used by Put *)
        z_2.red := FALSE;
        z_2.l := z_0.y;
        z_2.r := z_0.y
      END;
      z_0.h :=  NEW (Link);
      WITH z_3 = z_0.h^ DO
        z_3.key := FIRST (Key);
        z_3.red := FALSE;
        z_3.l := z_0.z;
        z_3.r := z_0.z
      END
    END;
    RETURN table
  END New;

PROCEDURE Get (table: T; key: Key; VAR value: REFANY): BOOLEAN RAISES {} =
  VAR x, z: Link;
  BEGIN
    z := table.z;
    x := table.h.r;
    WHILE (x # z) AND (key # x.key) DO
      IF key < x.key 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: REFANY): BOOLEAN RAISES {} =
    (* Based on Sedgewick's rbtreeinsert. *)
  VAR
    x: Link;
    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 *)

      BEGIN
        IF key < y.key THEN s := y.l ELSE s := y.r END;
        IF key < s.key THEN
          gs := s.l;
          s.l := gs.r;
          gs.r := s
        ELSE
          gs := s.r;
          s.r := gs.l;
          gs.l := s
        END;
        IF key < y.key 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 (key < g.key) # (key < f.key) THEN f := Rotate (g) END;
        x := Rotate (gg);
        x.red := FALSE
      END;
      table.h.r.red := FALSE
    END Split;
  BEGIN
    x := table.h;
    table.z.key := key;
    f := x;
    g := x;
    LOOP
      gg := g;
      g := f;
      f := x;
      IF key < x.key THEN x := x.l ELSE x := x.r END;
      IF key = x.key 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 key < f.key 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: REFANY): 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 c.key < b.key THEN
        b.l := c.r;
        c.r := b
      ELSE
        b.r := c.l;
        c.l := b
      END;
      IF b.key < a.key THEN a.l := c ELSE a.r := c END
    END Rotate;
  VAR
    x: Link;
    f: Link; (* father of x *)
    left: BOOLEAN; (* TRUE iff x = f^.l *)
    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 (key = x.key) 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 left 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 left 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 left 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;
      left := key < f.key;
      IF left 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 key < g.key 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_4 = x^ DO
          z_4.key := xOld.key;
          z_4.value := xOld.value;
          z_4.red := xOld.red;
          z_4.l := CopyNode (xOld.l);
          z_4.r := CopyNode (xOld.r)
        END;
        RETURN x
      END
    END CopyNode;
  BEGIN
    table :=  NEW (T);
    WITH z_5 = table^ DO
      z_5.h :=  NEW (Link);
      z_5.z :=  NEW (Link);
      z_5.y :=  NEW (Link);
      WITH z_6 = z_5.h^ DO
        z_6.key := tableOld.h.key;
        z_6.red := FALSE;
        z_6.l := z_5.z;
        z_6.r := CopyNode (tableOld.h.r)
      END;
      WITH z_7 = z_5.z^ DO
                  (* key varies *)
        z_7.red := FALSE;
        z_7.l := z_5.y;
        z_7.r := z_5.y
      END;
      z_5.y.red := TRUE
    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;
               start: Key;
               ascend: BOOLEAN;
               top: CARDINAL;
               stack: ARRAY [0..MaxStack] OF Link
             END;


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

PROCEDURE Next
  (s: Stream; VAR key: Key; VAR value: REFANY): BOOLEAN RAISES {} =
  VAR
    x: Link; (* = stack[top] *)

  BEGIN
    WITH z_9 = s^ DO
      x := z_9.stack[z_9.top];
      LOOP
        IF z_9.ascend THEN
          IF x.r = z_9.z THEN
            IF z_9.top = 0 THEN
              RETURN FALSE
            ELSE
              DEC (z_9.top);
              x := z_9.stack[z_9.top];
              key := x.key;
              value := x.value;
              RETURN TRUE
            END
          END;
          x := x.r;
          z_9.stack[z_9.top] := x;
          LOOP
            IF (x.l = z_9.z) OR (x.key <= z_9.start) THEN EXIT END;
            x := x.l;
            INC (z_9.top);
            z_9.stack[z_9.top] := x
          END;
          IF z_9.start <= x.key THEN
            key := x.key;
            value := x.value;
            RETURN TRUE
          END
        ELSE (* NOT ascend *)
          IF x.l = z_9.z THEN
            IF z_9.top = 0 THEN
              RETURN FALSE
            ELSE
              DEC (z_9.top);
              x := z_9.stack[z_9.top];
              key := x.key;
              value := x.value;
              RETURN TRUE
            END
          END;
          x := x.l;
          z_9.stack[z_9.top] := x;
          LOOP
                          (* Treat start = FIRST(Key) as start = LAST(Key). *)
            IF (x.r = z_9.z) OR (z_9.start # FIRST (Key))
              AND (z_9.start <= x.key) THEN
              EXIT
            END;
            x := x.r;
            INC (z_9.top);
            z_9.stack[z_9.top] := x
          END;
          IF (x.key <= z_9.start) OR (z_9.start = FIRST (Key)) 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 := FIRST (Key))
   : BOOLEAN(* RAISES {raises( pred )} *) =
  VAR k: Key; s: Stream; v: REFANY;
  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, k0: Key; l: List.T; s: Stream; v: REFANY;
  BEGIN
    IF ascending THEN k0 := LAST (Key) ELSE k0 := FIRST (Key) END;
    s := NewStream (table,  NOT ascending, k0);
    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, k0: Key; l: List.T; s: Stream; v: REFANY; r: REF INTEGER;
  BEGIN
    IF ascending THEN k0 := LAST (Key) ELSE k0 := FIRST (Key) END;
    s := NewStream (table,  NOT ascending, k0);
    l := NIL;
    WHILE Next (s, k, v) DO
      r := NEW (REF INTEGER);
      r^ := k;
      List.Push (l, List.List2 (r, v))
    END;
    RETURN l
  END ToAssocList;

BEGIN
END SIntTable.