(* Lecture 7: Representation Invariants *)
(* Author: Frank Pfenning *)

(* Workaround for Harlequin "restricted beta" problem *)
(* This make 'a option, SOME, NONE available at the top-level *)
open General;

(* Signature for dictionaries *)
(*
   For simplicity, we assume keys are strings, while stored entries
   are of arbitrary type.  This is prescribed in the signature.

   Existing entries may be "updated" by inserting a new entry
   with the same key.
*)

signature DICT =
sig
  type key = string
  type 'a entry = string * 'a

  type 'a dict

  val empty : 'a dict
  val lookup : 'a dict -> key -> 'a option
  val insert : 'a dict * 'a entry -> 'a dict
end;  (* signature DICT *)

(* Association list implementation of dictionaries popular in Lisp *)
structure AssocList :> DICT =
struct
  type key = string
  type 'a entry = string * 'a

  type 'a dict = ('a entry) list

  (* Representation Invariant *)
  (* Only leftmost entry in ('a entry) list matching key is valid *)

  val empty = nil

  fun lookup l key =
    (* lk l >=> leftmost entry in l matching key or NONE *)
    let fun lk nil = NONE
	  | lk ((key1,datum1)::l) = if key = key1 then SOME(datum1) else lk l
    in 
      lk l
    end

  (* Inserting shadows previous entries with same key *)
  fun insert (l, entry) = entry::l
end;  (* structure AssocList *)

(* Straight binary search tree implementation of dictionaries *)

structure BinarySearchTree :> DICT =
struct
  type key = string
  type 'a entry = string * 'a

  datatype 'a tree =
    Empty
  | Node of 'a entry * 'a tree * 'a tree

  type 'a dict = 'a tree

  (* Representation Invariant *)
  (*
     The tree is ordered:
       For every node  Node((key1,datum1), left, right),
       every key in left is LESS than key1, and
       every key in right is GREATER than key1
  *)

  val empty = Empty

  fun lookup tree key =
    let fun lk (Empty) = NONE
	  | lk (Node((key1,datum1), left, right)) =
              (case String.compare(key,key1)
		 of EQUAL => SOME(datum1)
	          | LESS => lk left
	          | GREATER => lk right)
    in
      lk tree
    end

  fun insert (tree, entry as (key,datum)) =
    (* val ins : 'a dict -> 'a dict *)
    (* ins (tree) inserts entry into tree *)
    let fun ins (Empty) = Node(entry, Empty, Empty)
	  | ins (Node(entry1 as (key1,datum1), left, right)) =
              (case String.compare(key,key1)
		 of EQUAL => Node(entry, left, right)
	          | LESS => Node(entry1, ins left, right)
                  | GREATER => Node(entry1, left, ins right))
    in 
      ins tree
    end
end;  (* structure BinarySearchTree *)

(* Some Examples
structure BST = BinarySearchTree;
val d1 = BST.empty;
val d2 = BST.insert (d1, ("a",1));
val d3 = BST.insert (d2, ("b",2));
val d4 = BST.insert (d3, ("c",3));
val d5 = BST.insert (d4, ("d",4));
*)

(* Red/Black tree refinement of binary search trees *)

structure RedBlackTree :> DICT =
struct
  type key = string
  type 'a entry = string * 'a

  datatype 'a dict =
    Empty				(* considered black *)
  | Red of 'a entry * 'a dict * 'a dict
  | Black of 'a entry * 'a dict * 'a dict

  (* Representation Invariants *)
  (*
     1. The tree is ordered: for every node Red((key1,datum1), left, right) or
        Black ((key1,datum1), left, right), every key in left is less than
        key1 and every key in right is greater than key1.

     2. The children of a red node are black (color invariant).

     3. Every path from the root to a leaf has the same number of
        black nodes, called the black height of the tree.
  *)

  val empty = Empty

  fun lookup dict key =
    let
      fun lk (Empty) = NONE
	| lk (Red tree) = lk' tree
        | lk (Black tree) = lk' tree
      and lk' ((key1, datum1), left, right) =
	    (case String.compare(key,key1)
	       of EQUAL => SOME(datum1)
	        | LESS => lk left
		| GREATER => lk right)
      in
	lk dict
      end

  (* val restore_right : 'a dict -> 'a dict *)
  (*
     restore_right (Black(e,l,r)) >=> dict
     where (1) Black(e,l,r) is ordered,
           (2) Black(e,l,r) has black height n,
	   (3) color invariant may be violated at the root of r:
               one of its children might be red.
     and dict is a re-balanced red/black tree (satisfying all invariants)
     and same black height n.
  *)
  fun restore_right (Black(e, Red lt, Red (rt as (_,Red _,_)))) =
         Red(e, Black lt, Black rt)	(* re-color *)
    | restore_right (Black(e, Red lt, Red (rt as (_,_,Red _)))) =
         Red(e, Black lt, Black rt)	(* re-color *)
    | restore_right (Black(e, l, Red(re, Red(rle, rll, rlr), rr))) =
	 (* l is black, deep rotate *)
	 Black(rle, Red(e, l, rll), Red(re, rlr, rr))
    | restore_right (Black(e, l, Red(re, rl, rr as Red _))) =
	 (* l is black, shallow rotate *)
	 Black(re, Red(e, l, rl), rr)
    | restore_right dict = dict

  (* restore_left is like restore_right, except *)
  (* the color invariant may be violated only at the root of left child *)
  fun restore_left (Black(e, Red (lt as (_,Red _,_)), Red rt)) =
	 Red(e, Black lt, Black rt)	(* re-color *)
    | restore_left (Black(e, Red (lt as (_,_,Red _)), Red rt)) =
	 Red(e, Black lt, Black rt)	(* re-color *)
    | restore_left (Black(e, Red(le, ll as Red _, lr), r)) =
	 (* r is black, shallow rotate *)
	 Black(le, ll, Red(e, lr, r))
    | restore_left (Black(e, Red(le, ll, Red(lre, lrl, lrr)), r)) =
	 (* r is black, deep rotate *)
	 Black(lre, Red(le, ll, lrl), Red(e, lrr, r))
    | restore_left dict = dict

  fun insert (dict, entry as (key,datum)) =
    let
      (* val ins : 'a dict -> 'a dict  inserts entry *)
      (* ins (Red _) may violate color invariant at root *)
      (* ins (Black _) or ins (Empty) will be red/black tree *)
      (* ins preserves black height *)
      fun ins (Empty) = Red(entry, Empty, Empty)
	| ins (Red(entry1 as (key1, datum1), left, right)) =
	  (case String.compare(key,key1)
	     of EQUAL => Red(entry, left, right)
	      | LESS => Red(entry1, ins left, right)
	      | GREATER => Red(entry1, left, ins right))
	| ins (Black(entry1 as (key1, datum1), left, right)) =
	  (case String.compare(key,key1)
	     of EQUAL => Black(entry, left, right)
	      | LESS => restore_left (Black(entry1, ins left, right))
	      | GREATER => restore_right (Black(entry1, left, ins right)))
    in
      case ins dict
	of Red (t as (_, Red _, _)) => Black t (* re-color *)
	 | Red (t as (_, _, Red _)) => Black t (* re-color *)
	 | dict => dict
    end

end;  (* structure RedBlackTree *)

(* Some examples:
structure RBT = RedBlackTree;
val r0 = RBT.empty;
val r1 = RBT.insert (r0, ("a",1));
val r2 = RBT.insert (r1, ("b",2));
val r3 = RBT.insert (r2, ("c",3));
val r4 = RBT.insert (r3, ("d",4));
val r5 = RBT.insert (r4, ("e",5));
*)