Purely Functional AVL Tree

avl.hs

module AVLZipper (singleton, empty , fromList , maximum , minimum) where

import Prelude hiding (minimum, maximum)
import Control.Applicative hiding (empty)
import Control.Monad
import Control.Monad.Trans.State
import Control.Monad.Trans.Identity
import Debug.Trace

data Tree a = Empty
            | Node { key :: a, _height :: Int, left::Tree a, right::Tree a }
            deriving (Show, Eq)

data Balance = Balanced
             | LeftHeavy
             | RightHeavy
             deriving (Show, Eq)

data Move a = ToLeft  a Int (Tree a)
            | ToRight a Int (Tree a)
            deriving (Show, Eq)

type Path a = [Move a]
type Zipper a = (Tree a, Path a)
type Navigation a = State (Zipper a)

empty :: Tree a
empty = Empty

singleton :: a -> Tree a
singleton a = Node a 1 Empty Empty

height :: Tree a -> Int
height Empty = 0
height (Node _ n _ _) = n

checkBalance :: Tree a -> Balance
checkBalance t | hl - hr > 1 = LeftHeavy
               | hr - hl > 1 = RightHeavy
               | otherwise = Balanced
               where hl = height $ left t
                     hr = height $ right t

getNode :: Navigation a (Tree a)
getNode = fst <$> get

getPath :: Navigation a [Move a]
getPath = snd <$> get

setNode :: Tree a -> Navigation a ()
setNode node = getPath >>= (\path -> put (node, path))

modifyNode :: (Tree a -> Tree a) -> Navigation a ()
modifyNode f = getNode >>= setNode . f

heightify :: Navigation a ()
heightify = modifyNode (\node -> case node of
        Empty -> Empty
        (Node k n l r) -> Node k (1 + max (height l) (height r)) l r)

goLeft :: Navigation a ()
goLeft = modify goLeft' where
    goLeft' (Node x h l r, mv) = (l, ToLeft x h r : mv)
    goLeft' _ = error "Reached Leaf"

goRight :: Navigation a ()
goRight = modify goRight' where
    goRight' (Node x h l r, mv) = (r, ToRight x h l : mv)
    goRight' _ = error "Reached Leaf"

goBack :: Navigation a ()
goBack = modify goBack' where
    goBack' (cur, ToLeft y h sy : rest) = (Node y h cur sy, rest)
    goBack' (cur, ToRight y h sy : rest) = (Node y h sy cur, rest)
    goBack' _ = error "Reached Root"

goToRoot :: Navigation a () -> Navigation a ()
goToRoot action = do isRoot <- null <$> getPath
                     action
                     unless isRoot (goBack >> goToRoot action)

root :: Navigation a (Tree a)
root = goToRoot (return ()) >> getNode

insert :: Ord a => a -> Navigation a (Tree a)
insert key = do n <- getNode
                h <- getPath
                case n of
                  Node x _ _ _
                    | key == x  -> return n
                    | key < x   -> goLeft >> insert key
                    | otherwise -> goRight >> insert key
                  Empty -> do setNode (singleton key)
                              goToRoot (heightify >> rebalance)
                              getNode

goto :: Ord a => a -> Navigation a (Maybe (Tree a))
goto key = do n <- getNode
              case n of
                    Node x _ _ _
                        | key == x  -> return (Just n)
                        | key < x   -> goLeft >> goto key
                        | otherwise -> goRight >> goto key
                    Empty -> return Nothing

minimum :: Ord a => Navigation a (Tree a)
minimum = do n <- getNode
             if n == Empty || left n == Empty
                then return n
                else goLeft >> minimum

maximum :: Ord a => Navigation a (Tree a)
maximum = do n <- getNode
             if n == Empty || right n == Empty
                 then return n
                 else goRight >> maximum

rightRotate :: Ord a => Navigation a ()
rightRotate = do q <- getNode
                 let p = left q
                     a = left p
                     b = right p
                     c = right q
                 setNode (Node (key p) (height p) a (Node (key q) (height q) b c))
                 sequence_ [ goRight, heightify, goBack, heightify ]

leftRotate :: Ord a => Navigation a ()
leftRotate = do p <- getNode
                let a = left p
                    q = right p
                    b = left q
                    c = right q
                setNode (Node (key q) (height q) (Node (key p) (height p) a b) c)
                sequence_ [ goLeft, heightify, goBack, heightify ]

rebalance :: Ord a => Navigation a ()
rebalance = do node <- getNode
               case checkBalance node of
                    Balanced  -> return ()
                    LeftHeavy -> case checkBalance (left node) of
                                     Balanced  -> rightRotate
                                     LeftHeavy -> rightRotate
                                     RightHeavy -> do { goLeft; leftRotate; goBack;
                                                        heightify; rightRotate }
                    RightHeavy -> case checkBalance (right node) of
                                     Balanced   -> leftRotate
                                     RightHeavy -> leftRotate
                                     LeftHeavy  -> do { goRight; rightRotate; goBack;
                                                        heightify; leftRotate  }

fromList :: Ord a => [a] -> Tree a
fromList xs = fst (execState (mapM_ insert xs) (Empty, []))

navigate :: Navigation a b -> Tree a -> Tree a
navigate nav tree = fst (execState nav (tree, []))