module SkewRandList
  ( RList
  , cons
  , head
  , tail
  , lookup
  , update ) where

import           Prelude hiding (head, lookup, tail)

data BTree a = Leaf a
             | Node a (BTree a) (BTree a)
  deriving (Show, Eq)

data Digit a = Digit !Int (BTree a)
  deriving (Show, Eq)

type RList a = [Digit a]

cons :: a -> RList a -> RList a
cons x (Digit w1 t1 : Digit w2 t2 : ts')
   | w1 == w2  = Digit (1 + w1 + w2) (Node x t1 t2) : ts'
cons x ts      = Digit 1 (Leaf x) : ts

head :: RList a -> a
head (Digit 1 (Leaf x)     : _) = x
head (Digit _ (Node x _ _) : _) = x

tail :: RList a -> RList a
tail (Digit 1 (Leaf _)       : ts) = ts
tail (Digit w (Node _ t1 t2) : ts) =
  Digit w' t1 : Digit w' t2 : ts
 where w' = w `div` 2

lookup :: RList a -> Int -> a
lookup (Digit w t : ts) i
   | i < w = lookupTree w t i
   | otherwise = lookup ts (i - w)

update :: RList a -> Int -> a -> RList a
update (Digit w t : ts) i y
   | i < w = Digit w (updateTree w t i y) : ts
   | otherwise = Digit w t : update ts (i-w) y

lookupTree :: Int -> BTree a -> Int -> a
lookupTree 1 (Leaf x)     0 = x
lookupTree _ (Node x _ _) 0 = x
lookupTree w (Node _ t1 t2) i
   | i < w'    = lookupTree w' t1 (i-1)
   | otherwise = lookupTree w' t2 (i-1-w')
 where w' = w `div` 2

updateTree :: Int -> BTree a -> Int -> a -> BTree a
updateTree 1 (Leaf _) 0 y = Leaf y
updateTree _ (Node _ t1 t2) 0 y = Node y t1 t2
updateTree w (Node x t1 t2) i y
   | i < (w`div`2) =
     Node x (updateTree w' t1 (i-1) y) t2
   | otherwise     =
     Node x t1 (updateTree w' t2 (i-1-w') y)
 where w' = w `div` 2