{- -----------------------------------------------------
  Binary search trees (BST) for key-value associations
  Pedro Vasconcelos, 2020
  Based on "How to specify it" talk by John Hughes
  https://youtu.be/G0NUOst-53U

 -- build trees
 empty  :: BST k v
 insert :: Ord k => k -> v  -> BST k v -> BST k v
 delete :: Ord k => k       -> BST k v -> BST k v
 union  :: Ord k => BST k v -> BST k v -> BST k v
 fromList :: Ord k => [(k,v)] -> BST k v

 -- find value by key
 find :: Ord k => k -> BST k v -> Maybe v

 -- auxiliary: list keys and values
 toList   :: BST k v -> [(k,v)]
 keys     :: BST k v -> [k]
 -------------------------------------------------------
-}

module BST where

data BST k v = Branch k v (BST k v) (BST k v)
             | Leaf
             deriving (Eq, Show)


empty :: BST k v
empty = Leaf


insert :: Ord k => k -> v -> BST k v -> BST k v
insert k v Leaf
  = Branch k v Leaf Leaf
insert k v (Branch k' v' left right)
  = case compare k k' of
      LT -> Branch k' v' (insert k v left) right 
      GT -> Branch k' v' left (insert k v right)
      EQ -> Branch k v left right


find :: Ord k => k -> BST k v -> Maybe v
find k Leaf
  = Nothing
find k (Branch k' v' left right)
  = case compare k k' of
      EQ -> Just v'
      LT -> find k left
      GT -> find k right


delete :: Ord k => k -> BST k v -> BST k v
delete k Leaf
  = Leaf
delete k (Branch  k' v' left right)
  = case compare k k' of
      LT -> Branch k' v' (delete k left)  right
      GT -> Branch k' v' left (delete k right)
      EQ -> union left right


union :: Ord k => BST k v -> BST k v -> BST k v
union Leaf t = t
union t Leaf = t
union (Branch k1 v1 l1 r1) t2
  = case split k1 t2 of
      (l2, r2) -> Branch k1 v1 (union l1 l2) (union r1 r2)

split :: Ord k => k -> BST k v -> (BST k v, BST k v)
split k Leaf = (Leaf, Leaf)
split k (Branch k' v' left right)
  = case compare k k' of
      EQ -> (left, right)
      LT -> let (left1, left2) = split k left
            in (left1, Branch k' v' left2 right)
      GT -> let (right1, right2) = split k right
            in (Branch k' v' left  right1,  right2)


toList :: BST k v -> [(k,v)]
toList Leaf = []
toList (Branch k v left right) = toList left ++ [(k,v)] ++ toList right


fromList :: Ord k => [(k,v)] -> BST k v
fromList [] = empty
fromList ((k,v):kvs) = insert k v (fromList kvs)

keys :: BST k v -> [k]
keys t = map fst (toList t)