-- Ada implementation of red-black trees as described in
--   Alternatives to Two Classic Data Structures
--   Chris Okasaki
--   SIGCSE 2005
--
-- Implements a set of integers.  Can easily be adapted to other
-- key types, or to other abstractions such as bags/dictionaries/etc.

PACKAGE BODY Red_Black IS

   FUNCTION Balance(X,Y,Z,B,C : Tree_Ptr) RETURN Tree_Ptr IS
   BEGIN
      -- Rearrange/recolor the tree as
      --       Y      <== red
      --      / \
      --     /   \
      --    X     Z   <== both black
      --   / \   / \
      --  A   B C   D
      --
      -- Note: A and D are not passed in because already in the right place

      X.Right := B;
      Y.Left := X;
      Y.Right := Z;
      Z.Left := C;
      X.Color := BLACK;
      Y.Color := RED;
      Z.Color := BLACK;
      return Y;
   END Balance;

   FUNCTION Is_Red(T : Tree_Ptr) RETURN Boolean IS
   BEGIN
      RETURN T /= NULL AND THEN T.Color = Red;
   END Is_Red;

   PROCEDURE Ins(Key : Integer; T : IN OUT Tree_Ptr) IS
   BEGIN
      IF T = NULL THEN
         T := NEW Tree_Node'(Key, Red, NULL, NULL);
      ELSIF Key < T.Key THEN
         Ins(Key, T.Left);
      ELSIF Key > T.Key THEN
         Ins(Key, T.Right);
      ELSE -- Key = T.Key
         RETURN;
      END IF;

      -- check for red child and red grandchild
      IF Is_Red(T.Left) AND THEN Is_Red(T.Left.Left) THEN
         --       Z           Y
         --      / \         / \
         --     Y   D  ==>  /   \
         --    / \         X     Z
         --   X   C       / \   / \
         --  / \         A   B C   D
         -- A   B
         T := Balance(T.Left.Left, T.Left, T,                 -- X,Y,Z
                      T.Left.Left.Right, T.Left.Right);       -- B,C
      ELSIF Is_Red(T.Left) AND THEN Is_Red(T.Left.Right) THEN
         --       Z           Y
         --      / \         / \
         --     X   D  ==>  /   \
         --    / \         X     Z
         --   A   Y       / \   / \
         --      / \     A   B C   D
         --     B   C
         T := Balance(T.Left, T.Left.Right, T,                -- X,Y,Z
                      T.Left.Right.Left, T.Left.Right.Right); -- B,C
      ELSIF Is_Red(T.Right) AND THEN Is_Red(T.Right.Left) THEN
         --     X             Y
         --    / \           / \
         --   A   Z    ==>  /   \
         --      / \       X     Z
         --     Y   D     / \   / \
         --    / \       A   B C   D
         --   B   C
         T := Balance(T, T.Right.Left, T.Right,               -- X,Y,Z
                      T.Right.Left.Left, T.Right.Left.Right); -- B,C
      ELSIF Is_Red(T.Right) AND THEN Is_Red(T.Right.Right) THEN
         --   X               Y
         --  / \             / \
         -- A   Y      ==>  /   \
         --    / \         X     Z
         --   B   Z       / \   / \
         --      / \     A   B C   D
         --     C   D
         T := Balance(T, T.Right, T.Right.Right,              -- X,Y,Z
                      T.Right.Left, T.Right.Right.Left);      -- B,C
      END IF;
   END Ins;

   PROCEDURE Insert(Key : Integer; Set : IN OUT Set_Type) IS
   BEGIN
      Ins(Key, Set.Root);
      Set.Root.Color := Black; -- always reColor root black
   END Insert;

   FUNCTION Member(Key : Integer; Set : Set_Type) RETURN Boolean IS
      T : Tree_Ptr := Set.Root;
   BEGIN
      WHILE T /= NULL LOOP
         IF Key < T.Key THEN
            T := T.Left;
         ELSIF Key > T.Key THEN
            T := T.Right;
         ELSE -- Key = T.Key
            RETURN True;
         END IF;
      END LOOP;
      RETURN False;
   END Member;

END Red_Black;