content: {-# LANGUAGE GADTs, DataKinds, KindSignatures, ExistentialQuantification, 
      TypeFamilies, StandaloneDeriving #-}
  
  
  -- This module is a Haskell transliteration of RBT.agda
  
  -- 49 lines of code for insertion (counting type defs & signatures)
  
  
  module RBT where
      
  -- a fixed element type for simplicity, could generalize the definitions here
  data A = A1 | A2 | A3 deriving (Eq, Ord, Show)
  
  data Color :: * where
      R :: Color
      B :: Color
   deriving Show
      
  data Nat = Z | S Nat
   deriving Show
  
  data Tree :: Color -> Nat -> * where
    E  :: Tree B Z
    TR :: Tree B n  -> A -> Tree B n  -> Tree R n 
    TB :: Tree c1 n -> A -> Tree c2 n -> Tree B (S n)
  deriving instance Show (Tree c n)
  
  -- a well-formed red-black tree. Root is black, but height 
  -- could be anything.
  data RBT :: * where
    Root :: Tree B m -> RBT
  deriving instance Show RBT
  
  -- example: reverse the tree
  rev :: Tree c n -> Tree c n 
  rev E = E
  rev (TR a x b) = TR (rev b) x (rev a)
  rev (TB a x b) = TB (rev b) x (rev a)
  
  -- example: find the maximum value stored in a non-empty tree
  maxB :: Tree B (S n) -> A
  maxB (TB _ x E)          = x
  maxB (TB _ _ (TR a x b)) = maxR (TR a x b)
  maxB (TB _ _ (TB a x b)) = maxB (TB a x b)
  
  maxR :: Tree R n -> A
  maxR (TR _ x E)          = x
  maxR (TR _ _ (TB a x b)) = maxB (TB a x b)
  
  -- type definitions (w/o Show instances): 11 LOC
  
  -- singleton type
  data Sing (c :: Color) :: * where    
      SR :: Sing R
      SB :: Sing B
  deriving instance Show (Sing c)
      
  -- for insertion, sometimes we need to break the invariant.
  type family Incr (c :: Color) (x :: Nat) :: Nat where
      Incr R x = x
      Incr B x = S x
  
  -- hide the color of a *non-empty* tree
  -- in Haskell, it also works to     
  data HiddenTree :: Nat -> * where
    HR :: Tree R n     -> HiddenTree n
    HB :: Tree B (S n) -> HiddenTree (S n)
  deriving instance Show (HiddenTree n)
  
  -- captures the height, but not the fact that red nodes have black children
  data AlmostTree :: Nat -> * where
    AT :: Sing c -> (Tree c1 n) -> A -> (Tree c2 n) -> AlmostTree (Incr c n)
  deriving instance Show (AlmostTree n)
  
  -- four balance functions (w/o catch-all): 18 LOC
  
  -- input color is implicitly black 
  balanceLB ::  AlmostTree n -> A -> Tree c n -> HiddenTree (S n)
  -- these are the two rotation cases
  balanceLB (AT SR (TR a x b) y c) z d = HR (TR (TB a x b) y (TB c z d))
  balanceLB (AT SR a x (TR b y c)) z d = HR (TR (TB a x b) y (TB c z d))
  -- need to expand the catch-all, because the *proof* is different in each case.  
  balanceLB (AT SB a  x b) kv r = HB (TB (TB a x b) kv r)
  balanceLB (AT SR E x E) kv r = HB (TB (TR E x E) kv r)
  balanceLB (AT SR (TB a1 x1 a2) x (TB b1 y1 b2)) y c = HB (TB (TR (TB a1 x1 a2) x (TB b1 y1 b2)) y c)
  balanceLB _ _ _ = error "balanceLB: this is impossible"
  
  -- input color is implicitly black 
  balanceRB :: Tree c n -> A -> AlmostTree n -> HiddenTree (S n)
  -- these are the two rotation cases
  balanceRB a x (AT SR (TR b y c)  z d) = HR (TR (TB a x b) y (TB c z d))
  balanceRB a x (AT SR b y (TR c z d)) = HR (TR (TB a x b) y (TB c z d))
  -- catch-all 
  balanceRB a x (AT SR E y E) = HB (TB a x (TR E y E))
  balanceRB a x (AT SR (TB l x0 r) y (TB l' x1 r')) = HB (TB a x (TR (TB l x0 r) y (TB l' x1 r')))
  balanceRB a x (AT SB l kv r) = HB (TB a x (TB l kv r))
  balanceRB _ _ _ = error "balanceRB: this is impossible"
  
  
  balanceLR :: HiddenTree n -> A -> Tree c n -> AlmostTree n
  balanceLR (HR l) x r = AT SR l x r
  balanceLR (HB l) x r = AT SR l x r
  
  balanceRR :: Tree c n -> A -> HiddenTree n -> AlmostTree n
  balanceRR l x (HR r) = AT SR l x r
  balanceRR l x (HB r) = AT SR l x r
  
  -- other functions: 20
  
  -- forget that the top node of the tree satisfies the color invariant
  forget :: HiddenTree n -> AlmostTree n
  forget (HR (TR l x r)) = AT SR l x r
  forget (HB (TB l x r)) = AT SB l x r
  
  insBlack :: Tree B n -> A -> HiddenTree n
  insBlack E x = HR (TR E x E)
  insBlack (TB l y r) x = case (compare x y) of
      LT -> balanceLB (insAny l x) y r
      GT -> balanceRB l y (insAny r x)
      EQ -> HB (TB l y r)
  
  insAny :: Tree c n -> A -> AlmostTree n
  insAny (TR l y r) x =  case compare x y of
     LT -> balanceLR (insBlack l x) y r
     GT -> balanceRR l y (insBlack r x) 
     EQ -> AT SR l y r
  insAny (TB l y r) x = forget (insBlack (TB l y r) x)  
  insAny E          x = forget (insBlack E x)
  
  blacken :: AlmostTree n -> RBT
  blacken (AT _ l x r) = (Root (TB l x r))
  
  insert :: RBT -> A -> RBT
  insert (Root t) x = blacken (insAny t x)
  
  --------------------------------------------------
  -- Tests:
  
  t0 :: RBT
  t0 = insert (Root E) A3
  
  t1 :: RBT
  t1 = insert t0 A1
  
  t2 :: RBT
  t2 = insert t1 A2
  
  