+------------------------------------------------------------------------------+
| file:      BSTree.lhs                                                        |
|                                                                              |
| created:   29/10/97 by Andrew Rock.                                          |
|                                                                              |
| purpose:   This implements a depth/height balanced (AVL) binary search       |
|            tree abstract data type.                                          |
|                                                                              |
| ref:       Bird and Wadler, Introduction to Functional Programming, Ch 9     |
+------------------------------------------------------------------------------+

> module BSTree (BSTree, emptyBST, isEmptyBST, updateBST, deleteBST, lookupBST,
>                memberBST, lookupGuard, flattenBST, domBST, ranBST, countBST,
>                leftBST, rightBST, mapBST, pairs2BST, list2BST) where

+-------------+
| BSTree Type |
+-------------+

A BSTree is either empty or Node containing key, an associated value and left
and right subtrees. Type key must be an instance of type class Ord, so that
< and == work. The slope of each node is stored at the node to avoid
recomputation.

> data Ord key =>
>    BSTree key value 
>       =   Empty
>         | Node key value (BSTree key value) (BSTree key value) Slope
>       deriving Show

All the functions in this module maintain the following invariant: The
depth of left and right subtrees differ by mo more than 1.

+------------------------+
| Slope of a BSTree Node |
+------------------------+

Associated with the left and right subtrees is a Slope value which
indicates the difference: depth left - depth right.

> type Slope = Int

This returns the slope of a node.

> slope :: Ord k => BSTree k v -> Slope
> slope (Node _ _ _ _ s) = s

+-------------------+
| Depth of a BSTree |
+-------------------+

> type Depth = Int

This returns the depth of a tree. Finding depth is O(depth) because of the
slope information.

> depth :: Ord k => BSTree k v -> Depth
> depth Empty
>    = 0
> depth (Node _ _ l r s)
>    | s <= 0    = 1 + depth l
>    | otherwise = 1 + depth r

+-------------+
| Empty Trees |
+-------------+

emptyBST is an empty BSTree.

> emptyBST :: Ord k => BSTree k v
> emptyBST = Empty

isEmptyBST is a predicate on BSTrees that tests for emptiness.

> isEmptyBST :: Ord k => BSTree k v -> Bool
> isEmptyBST Empty
>    = True
> isEmptyBST _
>    = False

+-------------+
| Rebalancing |
+-------------+

After inserting or deleting a node the slope at a node may be -2 or 2.
balance restores the -1 < slope < 1 invariant.

This next type helps with the rebalancing.

> type Change = Int

A value of 1 indicates a tree got deeper, 0 same, -1 shallower.

> balance :: Ord k => (BSTree k v, Change) -> (BSTree k v, Change)
> balance (Node k v l r s, c)
>    | 1 < s
>       = (shiftRight (Node k v l r s), c - 1)
>    | s < -1
>       = (shiftLeft (Node k v l r s), c - 1)
>    | otherwise
>       = (Node k v l r s, c)
>         where
>         shiftRight (Node k v l r s)
>            | slope l == -1
>               = rotateRight (Node k v (rotateLeft l) r s)
>            | otherwise
>               = rotateRight (Node k v l r s)
>         shiftLeft (Node k v l r s)
>            | slope r == 1
>               = rotateLeft (Node k v l (rotateRight r) s)
>            | otherwise
>               = rotateLeft (Node k v l r s)
>         rotateRight (Node k v (Node k' v' l' r' s') r s)
>            = let (ss, ss') 
>                     = case (s, s') of
>                          ( 2,  0) -> (-1,  1)
>                          ( 2,  1) -> ( 0,  0)
>                          ( 2,  2) -> ( 0, -1)
>                          ( 1,  1) -> (-1, -1)
>                          ( 1,  0) -> (-1,  0)
>                          ( 1, -1) -> (-2,  0)
>              in Node k' v' l' (Node k v r' r ss') ss
>         rotateLeft (Node k v l (Node k' v' l' r' s') s)
>            = let (ss, ss')
>                     = case (s, s') of
>                          (-2,  0) -> ( 1, -1)
>                          (-2, -1) -> ( 0,  0)
>                          (-2, -2) -> ( 0,  1)
>                          (-1, -1) -> ( 1,  1)
>                          (-1,  0) -> ( 1,  0)
>                          (-1,  1) -> ( 2,  0)
>              in Node k' v' (Node k v l l' ss') r' ss

+--------------------+
| Inserting/Updating |
+--------------------+

updateBST f key value bst returns the new tree obtained by updating
bst with the key and value. If the key already exists, f is used to 
combine the two values. Use \x _ -> x to merely replace.

> updateBST :: Ord k => (v -> v -> v) -> k -> v -> BSTree k v -> BSTree k v
> updateBST f k' v'
>    = fst . update
>      where
>      update Empty
>         = (Node k' v' Empty Empty 0, 1)
>      update (Node k v l r s)
>         | k' < k  
>            = let (l', c') = update l
>                  c = if s >= 0 && c' == 1 then 1 else 0
>              in balance (Node k v l' r (s + c'), c) 
>         | k' == k 
>            = (Node k (f v' v) l r s, 0)
>         | otherwise
>            = let (r', c') =  update r
>                  c = if s <= 0 && c' == 1 then 1 else 0
>              in balance (Node k v l r' (s - c'), c)

+-------------------+
| Deleting Elements |
+-------------------+

This deletes an element.

> deleteBST :: Ord k => k -> BSTree k v -> BSTree k v
> deleteBST k' 
>    = fst . delete
>      where
>      delete Empty
>         = (Empty, 0)
>      delete (Node k v l r s)
>         | k' < k  
>            = let (l', c') = delete l
>                  c = if s == 1 && c' == -1 then -1 else 0
>              in balance (Node k v l' r (s + c'), c) 
>         | k' == k 
>            = join l r s
>         | otherwise
>            = let (r', c') =  delete r
>                  c = if s == -1 && c' == -1 then -1 else 0
>              in balance (Node k v l r' (s - c'), c)
>      join Empty r _
>         = (r, -1)
>      join l r s
>         = let ((l', c'), k', v') = split l
>               s' = s + c'
>               c = if s == 1 && c' == -1 then -1 else 0
>           in balance (Node k' v' l' r s', c)
>      split (Node k v l Empty s)
>         = ((l, -1), k, v)
>      split (Node k v l r s)
>         = let ((r', c'), k', v') = split r
>               c = if s == -1 && c' == -1 then -1 else 0
>           in (balance (Node k v l r' (s - c'), c), k', v')

+-------------------+
| Looking Up Values |
+-------------------+

This returns the value associated with a key in a BST. The value
is returned in as either Nothing or Just value.

> lookupBST :: Ord k => k -> BSTree k v -> Maybe v
> lookupBST k Empty
>    = Nothing
> lookupBST k' (Node k v l r _)
>    | k' < k    = lookupBST k' l
>    | k' == k   = Just v
>    | otherwise = lookupBST k' r

This is a membership test.
   
> memberBST :: Ord k => k -> BSTree k v -> Bool
> memberBST k t
>    = case lookupBST k t of
>         Nothing -> False
>         Just _  -> True

lookupGuard bst keys handler process tries to look up the keys.
If any are missing the handler is applied to the missing key
otherwise the process is applied to the list of values
successfuly looked up.

> lookupGuard :: Ord a => BSTree a b -> [a] -> (a -> c) -> ([b] -> c) -> c
> lookupGuard bst keys handler process
>    = lookupGuard' keys []
>      where
>      lookupGuard' [] vals
>         = process vals
>      lookupGuard' (k:ks) vals
>         = case lookupBST k bst of
>              Nothing    -> handler k
>              Just stuff -> lookupGuard' ks (vals ++ [stuff]) 

+------------+
| Flattening |
+------------+

This returns the list of all keys and associated values in
ascending order of key.

> flattenBST :: Ord k => BSTree k v -> [(k,v)]
> flattenBST Empty
>    = []
> flattenBST (Node k v Empty r _)
>    = (k, v) : flattenBST r
> flattenBST (Node k v (Node k' v' l' r' _) r _)
>    = flattenBST (Node k' v' l' (Node k v r' r 0) 0)

These return just the keys and just the values.

> domBST :: Ord k => BSTree k v -> [k]
> domBST Empty
>    = []
> domBST (Node k _ Empty r _)
>    = k : domBST r
> domBST (Node k v (Node k' v' l' r' _) r _)
>    = domBST (Node k' v' l' (Node k v r' r 0) 0)

> ranBST :: Ord k => BSTree k v -> [v]
> ranBST Empty
>    = []
> ranBST (Node k v Empty r _)
>    = v : ranBST r
> ranBST (Node k v (Node k' v' l' r' _) r _)
>    = ranBST (Node k' v' l' (Node k v r' r 0) 0)


+-----------------------------+
| converting lists to BSTrees |
+-----------------------------+

pairs2BST converts an association list to a BSTree.
If there are duplicate v's for a k, only the first is retained.

> pairs2BST :: Ord k => [(k,v)] -> BSTree k v
> pairs2BST []
>    = Empty
> pairs2BST ((k,v):kvs)
>    = updateBST (\x _ -> x) k v (pairs2BST kvs)

list2BST converts a list of keys to a BSTree.
The values in the tree are filled up with ().

> list2BST :: Ord k => [k] -> BSTree k ()
> list2BST []
>    = Empty
> list2BST (k:ks)
>    = updateBST (\x _ -> x) k () (list2BST ks)


+----------+
| Sundries |
+----------+

Number of elements.

> countBST :: Ord k => BSTree k v -> Int
> countBST Empty
>    = 0
> countBST (Node _ _ l r _)
>    = countBST l + 1 + countBST r

Left-most and right-most elements.

> leftBST :: Ord k => BSTree k v -> Maybe (k, v)
> leftBST Empty
>    = Nothing
> leftBST (Node k v Empty _ _)
>    = Just (k, v)
> leftBST (Node k v l _ _)
>    = leftBST l

> rightBST :: Ord k => BSTree k v -> Maybe (k, v)
> rightBST Empty
>    = Nothing
> rightBST (Node k v _ Empty _)
>    = Just (k, v)
> rightBST (Node k v _ r _)
>    = rightBST r

Mapping the associated values.

> mapBST :: Ord k => (v -> v') -> BSTree k v -> BSTree k v'
> mapBST _ Empty
>    = Empty
> mapBST f (Node k v l r s)
>    = Node k (f v) (mapBST f l) (mapBST f r) s