-- Binary Random Access List

module BRAL where

import Prelude hiding (head, tail, lookup, drop)

data Tree a = Leaf a | Node Int (Tree a) (Tree a) deriving Show
data Digit a = Zero | One (Tree a) deriving Show
type RList a = [Digit a]


size (Leaf _)     = 1
size (Node w _ _) = w

sizeD Zero    = 0
sizeD (One t) = size t


-- t1 and t2 are assumed to be the same size
link t1 t2 = Node (2 * size t1) t1 t2


empty :: RList a
empty = []


isEmpty :: RList a -> Bool
isEmpty ts = null ts


cons :: a -> RList a -> RList a
cons x ts = consTree (Leaf x) ts


consTree :: Tree a -> RList a -> RList a
consTree t []            = [One t]
consTree t (Zero : ts)   = (One t : ts)
consTree t (One t' : ts) = Zero : consTree (link t t') ts


unconsTree :: RList a -> (Tree a, RList a)
unconsTree [One t]      = (t, [])
unconsTree (One t : ts) = (t, Zero : ts)
unconsTree (Zero  : ts) = (t1, One t2 : ts')
    where
        (Node _ t1 t2, ts') = unconsTree ts


head :: RList a -> a
head ts = x
    where
        (Leaf x, _) = unconsTree ts


tail :: RList a -> RList a
tail ts = ts' 
    where
        (_, ts') = unconsTree ts


lookup :: Int -> RList a -> a
lookup i (Zero : ts) = lookup i ts
lookup i (One t : ts) | i < s      = lookupTree i t
                      | otherwise  = lookup (i - s) ts
    where
        s = size t


lookupTree :: Int -> Tree a -> a
lookupTree _ (Leaf x) = x
lookupTree i (Node w t1 t2)
    | i < w `div` 2 = lookupTree i t1
    | otherwise     = lookupTree (i - w `div` 2) t2


update :: Int -> a -> RList a -> RList a
update i x (Zero : ts) = Zero : update i x ts
update i x (One t : ts)
    | i < s     = One (updateTree i x t) : ts
    | otherwise = One t : update (i - s) x ts
    where
        s = size t

updateTree :: Int -> a -> Tree a -> Tree a
updateTree _ x (Leaf _) = Leaf x
updateTree i x (Node w t1 t2)
    | i < w `div` 2 = Node w (updateTree i x t1) t2
    | otherwise     = Node w t1 (updateTree (i - w `div` 2) x t2)