{-# LANGUAGE CPP #-}

-- | 
-- A random-access list implementation based on Chris Okasaki's approach
-- on his book \"Purely Functional Data Structures\", Cambridge University
-- Press, 1998, chapter 9.3.
--
-- 'RandomAccessList's are finite lists providing random-access ('lookup',
-- 'update', etc.) in /O(log n)/ while the list functionality 'head',
-- 'tail' and 'cons' still works in /O(1)/.
--
-- A 'RandomAccessList' uses 'Int's for effective indexing. The valid index
-- range of a 'RandomAccessList' of size @n@ is @[0 .. n-1]@. If an index is
-- out of range, an 'error' is raised.

module Data.RandomAccessList
	(
	-- * The RandomAccessList type
	  RandomAccessList, View(..)
	-- * Construction
	, empty, singleton, replicate
	-- * Query
	, null, isEmpty, length, size, member, index
	-- * List specific operations
	, head, tail, uncons, view, cons, append
	-- * Random-access specific operations
	, lookup, update, adjust, adjustLookup
	-- * Miscellaneous
	, filter, partition, zip, zipWith, unzip
	-- * Conversion
	-- ** List
	, fromList, toList, toIndexedList
	-- ** Map
	, toMap, toIntMap
	-- ** Array
	, fromArray, toArray
	) where

import Data.Array.IArray (IArray, Ix, array, elems)
import Data.Foldable (Foldable(..), foldr')
import Data.IntMap as IMap (IntMap, fromDistinctAscList)
import Data.Map as Map (Map, fromDistinctAscList)
import Data.Maybe
import Data.Monoid
import qualified Prelude (replicate, zip, zipWith)
import Prelude hiding
	( filter, foldl, foldl1, foldr, foldr1
	, head, length, lookup, null
	, replicate, tail, unzip, zip, zipWith
	)
import Text.Read

-- |
-- Random-access lists allowing /O(1)/ list operations and /O(log n)/
-- indexed access.
data RandomAccessList a
	= RandomAccessList {-# UNPACK #-} !Int [(Int, CBTree a)]

instance (Show a) => Show (RandomAccessList a) where
	show list = "fromList " ++ show (toList list)

instance (Read a) => Read (RandomAccessList a) where
#ifdef __GLASGOW_HASKELL__
	readPrec = parens $ prec 10 $ do
		Ident "fromList" <- lexP
		xs <- readPrec
		return (fromList xs)
	readListPrec = readListPrecDefault
#else
	readsPrec p = readParen (p > 10) $ \r -> do
		("fromList", s) <- lex r
		(xs, t) <- reads s
		return (fromList xs, t)
#endif

instance (Eq a) => Eq (RandomAccessList a) where
	(RandomAccessList s1 l1) == (RandomAccessList s2 l2)
		= s1 == s2 && (map snd l1) == (map snd l2)

instance (Ord a) => Ord (RandomAccessList a) where
	compare xs ys = fold (Prelude.zipWith compare (toList xs) (toList ys))

instance Monoid (RandomAccessList a) where
	mempty  = empty
	mappend = append

instance Functor RandomAccessList where
	fmap f (RandomAccessList s list) = RandomAccessList s
		(fmap (\(weight, tree) -> (weight, fmap f tree)) list)

instance Foldable RandomAccessList where
	foldr f zero list = foldr f zero (view list)
	foldl f zero list = foldl f zero (view list)
	foldr1 f list     = foldr1 f $ toList list
	foldl1 f list     = foldl1 f (view list)

-- |
-- Complete binary trees. It is not possible to ensure completeness via
-- the type system, so make sure you do!
data CBTree a
	= Leaf a
	| Node a !(CBTree a) !(CBTree a)
#ifndef __HADDOCK
	deriving (Eq)
#else
instance (Eq a) => Eq (CBTree a)
#endif

instance Functor CBTree where
	fmap f (Leaf x)            = Leaf (f x)
	fmap f (Node x left right) = Node (f x) (fmap f left) (fmap f right)

-- |
-- View the end of a 'RandomAccessList' which is either empty or has
-- been constructed by 'head' and 'tail'.
data View a
	= Empty                       -- ^ An empty 'RandomAccessList'.
	| Cons a (RandomAccessList a) -- ^ 'head' and 'tail' of a non-empty 'RandomAccessList'.
#ifndef __HADDOCK__
	deriving (Show, Read, Eq, Ord)
#else
instance (Show a) => Show (View a)
instance (Read a) => Read (View a)
instance (Eq a)   => Eq   (View a)
instance (Ord a)  => Ord  (View a)
#endif

instance Functor View where
	fmap _ Empty       = Empty
	fmap f (Cons x xs) = Cons (f x) (fmap f xs)

instance Foldable View where
	foldr _ zero Empty       = zero
	foldr f zero (Cons x xs) = f x (foldr f zero xs)

	foldl _ zero Empty       = zero
	foldl f zero (Cons x xs) = foldl f (f zero x) xs

	foldr1 _ Empty           = error "foldr1: empty View"
	foldr1 f (Cons x xs)     = f x (foldr1 f xs)

	foldl1 _ Empty           = error "foldl1: empty View"
	foldl1 f (Cons x xs)     = foldl f x xs

-- |
-- Generate a readable error message for out-of-range index access.
indexErrorMessage :: Maybe (Int, Int) -- ^ Accessed index and list 'length' if available.
	-> String                     -- ^ The error message.
indexErrorMessage = maybe
	"RandomAccessList: index out of range"
	(\(i, s) -> "RandomAccessList: index " ++ show i
		++ " out of range [0.." ++ show (s - 1) ++ "]")

-- |
-- /O(1)/. Builds an empty 'RandomAccessList'.
empty :: RandomAccessList a
empty = RandomAccessList 0 []

-- |
-- /O(1)/. Builds a singleton 'RandomAccessList'.
singleton :: a -> RandomAccessList a
singleton x = RandomAccessList 1 [(1, Leaf x)]

-- |
-- /O(n)/. @'replicate' n x@ constructs a 'RandomAccessList' that
-- contains the same element @x@ @n@ times.
replicate :: Int -> a -> RandomAccessList a
replicate n = fromList . (Prelude.replicate n)

-- |
-- /O(1)/. Is the 'RandomAccessList' empty?
null :: RandomAccessList a -> Bool
null = isEmpty

-- |
-- /O(1)/. Is the 'RandomAccessList' empty?
isEmpty :: RandomAccessList a -> Bool
isEmpty (RandomAccessList 0 _) = True
isEmpty _                      = False

-- |
-- /O(1)/. The number of elements contained in a 'RandomAccessList'.
length :: RandomAccessList a -> Int
length = size

-- |
-- /O(1)/. The number of elements contained in a 'RandomAccessList'.
size :: RandomAccessList a -> Int
size (RandomAccessList s _) = s

-- |
-- /O(n)/. Is the given element a member of the 'RandomAccessList'?
member :: (Eq a) => a -> RandomAccessList a -> Bool
member x = isJust . (index x)

-- |
-- /O(n)/. Find the index of a given element.
index :: (Eq a) => a -> RandomAccessList a -> Maybe Int
index x (RandomAccessList _ list) = index' 0 list
	where
	index' _     []                     = Nothing
	index' accum ((weight, tree):trees) = maybe
		(index' (accum + weight) trees)
		(Just . (+accum))
		(indexTree weight x tree)

-- |
-- Find the index of a given element in a 'CBTree'. Returns @'Just'
-- index@ on success and 'Nothing' on failure.
indexTree :: (Eq a) => Int -> a -> CBTree a -> Maybe Int
indexTree _      x (Leaf e)            = if x == e then Just 0 else Nothing
indexTree weight x (Node e left right) = if x == e
	then Just 0
	else case indexTree halfWeight x left of
		Just i  -> Just (i + 1)
		Nothing -> case indexTree halfWeight x right of
			Just index' -> Just (index' + 1 + halfWeight)
			Nothing     -> Nothing
	where
	halfWeight = weight `div` 2

-- |
-- /O(1)/. Returns the head of a 'RandomAccessList'.
head :: RandomAccessList a -> a
head = fst . uncons

-- |
-- /O(1)/. Retrieve the tail of a 'RandomAccessList'.
tail :: RandomAccessList a -> RandomAccessList a
tail = snd . uncons

-- |
-- /O(1)/. Retrieve both, 'head' and 'tail' of a 'RandomAccessList'.
uncons :: RandomAccessList a -> (a, RandomAccessList a)
uncons list = case view list of
	Empty     -> error "RandomAccessList: empty list"
	Cons x xs -> (x, xs)

-- |
-- /O(1)/. Examine a 'RandomAccessList': Either it is 'Empty' or it has
-- a 'head' and a 'tail' (packed in 'Cons').
view :: RandomAccessList a -> View a
view (RandomAccessList s list) = case list of
	[]                                     -> Empty
	((_, Leaf x) : trees')                 -> Cons x (RandomAccessList (s - 1) trees')
	((weight, Node x left right) : trees') -> let
		halfWeight = weight `div` 2
		trees''    = (halfWeight, left) : (halfWeight, right) : trees'
		in
		Cons x (RandomAccessList (s - 1) trees'')

-- |
-- /O(1)/. Prepend an element to the 'RandomAccessList'.
cons :: a -> RandomAccessList a -> RandomAccessList a
cons x (RandomAccessList s list) = RandomAccessList (s + 1) (case list of
	xss@((size1, tree1) : (size2, tree2) : xs) -> if size1 == size2
		then (1 + 2 * size1, Node x tree1 tree2) : xs
		else (1,             Leaf x)             : xss
	xs                                         -> (1, Leaf x) : xs)

-- |
-- /O(n)/ where /n/ is the 'length' of the first list. Appends the second
-- list to the first list.
append :: RandomAccessList a -> RandomAccessList a -> RandomAccessList a
append = (flip (foldr cons))

-- |
-- /O(log n)/. Retrieve the /i/th element of the list. Unless
-- /0 <= i < n/, an 'error' is raised.
lookup :: Int -> RandomAccessList a -> a
lookup i = fst . (adjustLookup undefined i)

-- |
-- /O(log n)/. Set the /i/th element of the list. Unless
-- /0 <= i < n/, an 'error' is raised.
update :: Int -> a -> RandomAccessList a -> RandomAccessList a
update i x = snd . (adjustLookup (const x) i)

-- |
-- /O(log n)/. Adjust /i/th element of the list according to the
-- given function. Unless /0 <= i < n/, an 'error' is raised.
adjust :: (a -> a) -> Int -> RandomAccessList a -> RandomAccessList a
adjust f i = snd . (adjustLookup f i)

-- |
-- /O(log n)/. Find the /i/th element of the list and change it. This
-- function returns the element that is at index /i/ in the original
-- 'RandomAccessList' and a new 'RandomAccessList' with the /i/th
-- element replaced according to the given function:
--
-- @
--    lookup   index list === fst (adjustLookup undefined index list)
--    adjust f index list === snd (adjustLookup f         index list)
-- @
--
-- Unless /0 <= i < n/, an 'error' is raised.
adjustLookup :: (a -> a)           -- ^ Modifying element function.
	-> Int                     -- ^ Index of the affected element.
	-> RandomAccessList a      -- ^ 'RandomAccessList' to be modified.
	-> (a, RandomAccessList a) -- ^ Original element and modified 'RandomAccessList'.
adjustLookup f i (RandomAccessList s list) = let
	(x, list') = updateLookup' i list
	in
	(x, RandomAccessList s list')
	where
	updateLookup' _      []                    =
		error $ indexErrorMessage (Just (i, s))
	updateLookup' i' ((weight, tree) : xs) = if i' < weight
		then let
			(x, tree') = adjustLookupTree f weight i' tree
			in
			(x, (weight, tree') : xs)
		else let (x, trees) = updateLookup' (i' - weight) xs
			in
			(x, (weight, tree) : trees)

-- |
-- Performs 'adjustLookup' on a complete binary tree. Note that it
-- is crucial, that the 'CBTree' really is complete!
adjustLookupTree :: (a -> a) -- ^ Function to change a given element.
	-> Int               -- ^ Size of the complete binary tree.
	-> Int               -- ^ Index to be replaced.
	-> CBTree a          -- ^ The tree.
	-> (a, CBTree a)
adjustLookupTree f 1      0 (Leaf x) = (x, Leaf (f x))
adjustLookupTree _ _      _ (Leaf _) = error (indexErrorMessage Nothing)
adjustLookupTree f weight i (Node x left right)
	| i == 0          = (x, Node (f x) left right)
	| i <= halfWeight = let
		(x', left') = adjustLookupTree f halfWeight (i - 1) left
		in
		(x', Node x left' right)
	| otherwise       = let
		(x', right') = adjustLookupTree f halfWeight (i - 1 - halfWeight) right
		in
		(x', Node x left right')
	where
	halfWeight = weight `div` 2

-- |
-- /O(n)/. Remove all elements from a 'RandomAccessList' not fulfilling a
-- predicate.
filter :: (a -> Bool) -> RandomAccessList a -> RandomAccessList a
filter p = fst . (partition p)

-- |
-- /O(n)/. Split a 'RandomAccessList' into two: The elements in the first
-- fulfill the given prefix, the others don't.
partition :: (a -> Bool) -> RandomAccessList a -> (RandomAccessList a, RandomAccessList a)
partition p list = case view list of
	Empty     -> (empty, empty)
	Cons x xs -> let (yes, no) = partition p xs in if p x
		then (x `cons` yes,          no)
		else (         yes, x `cons` no)

-- |
-- /O(min(n, m))/. List-like 'Prelude.zip'. This function is slightly faster
-- when called with two 'RandomAccessList's of equal 'length'.
zip :: RandomAccessList a -> RandomAccessList b -> RandomAccessList (a, b)
zip = zipWith (,)

-- |
-- /O(min(n, m))/. List-like 'Prelude.zipWith'. This function is slightly faster
-- when called with two 'RandomAccessList's of equal 'length'.
zipWith :: (a -> b -> c) -> RandomAccessList a -> RandomAccessList b -> RandomAccessList c
zipWith f l1@(RandomAccessList size1 as) l2@(RandomAccessList size2 bs) = if size1 == size2
	then RandomAccessList size1 $ Prelude.zipWith
		(\(w, xs) (_, ys) -> (w, zipWithTree f xs ys)) as bs
	else case (view l1, view l2) of
		(Cons x xs, Cons y ys) -> f x y `cons` zipWith f xs ys
		_                      -> empty

-- |
-- Performs 'zipWith' for complete binary trees.
zipWithTree :: (a -> b -> c) -> CBTree a -> CBTree b -> CBTree c
zipWithTree f (Leaf x1)              (Leaf x2)              = Leaf (f x1 x2)
zipWithTree f (Node x1 left1 right1) (Node x2 left2 right2) = Node
	(f x1 x2)
	(zipWithTree f left1 left2)
	(zipWithTree f right1 right2)
zipWithTree _ _                      _                      = error
	"RandomAccessList: invalid skew binary number representation"

-- |
-- /O(n)/. List-like 'Prelude.unzip' for 'RandomAccessList's.
unzip :: RandomAccessList (a, b) -> (RandomAccessList a, RandomAccessList b)
unzip (RandomAccessList s xs) = let (as, bs) = unzip' xs
	in
	(RandomAccessList s as, RandomAccessList s bs)
	where
	unzip' []                       = ([], [])
	unzip' ((weight, tree) : trees) = let
		(treeA, treeB)   = unzipTree tree
		(treesA, treesB) = unzip' trees
		in
		((weight, treeA) : treesA, (weight, treeB) : treesB)

-- |
-- Performs 'unzip' on complete binary trees.
unzipTree :: CBTree (a, b) -> (CBTree a, CBTree b)
unzipTree (Leaf (x, y))            = (Leaf x, Leaf y)
unzipTree (Node (x, y) left right) = (Node x xLeft xRight, Node y yLeft yRight)
	where
	(xLeft,  yLeft)  = unzipTree left
	(xRight, yRight) = unzipTree right

-- |
-- /O(n)/. Build a 'RandomAccessList' from a list.
fromList :: [a] -> RandomAccessList a
fromList = foldr' cons empty -- just as fast as building the trees by hand!

-- |
-- /O(n)/. Convert a 'RandomAccessList' to a list.
toList :: RandomAccessList a -> [a]
toList = foldr (:) []

-- |
-- /O(n)/. Convert a 'RandomAccessList' to a list of tuples each holding
-- an element and its index. The list is ordered ascending regarding the
-- indices.
toIndexedList :: RandomAccessList a -> [(Int, a)]
toIndexedList list = Prelude.zip [0..] (toList list)

-- |
-- /O(n)/. Build a 'Map.Map' from a 'RandomAccessList'. The keys in the
-- 'Map.Map' are the indices of the elements in the 'RandomAccessList'.
toMap :: RandomAccessList a -> Map Int a
toMap = Map.fromDistinctAscList . toIndexedList

-- |
-- /O(n)/. Build an 'IMap.IntMap' from a 'RandomAccessList'. The keys in the
-- 'IMap.IntMap' are the indices of the elements in the 'RandomAccessList'.
toIntMap :: RandomAccessList a -> IntMap a
toIntMap = IMap.fromDistinctAscList . toIndexedList

-- |
-- /O(n)/. Given an 'IArray', generate a 'RandomAccessList'. The elements'
-- order will be preserved.
fromArray :: (IArray a e, Ix i) => a i e -> RandomAccessList e
fromArray = fromList . elems

-- |
-- /O(n)/. Build an 'IArray' from the 'RandomAccessList'. It will have
-- an index range from @[0 .. n-1]@, where @n@ is the 'RandomAccessList's
-- 'length'.
toArray :: (IArray a e) => RandomAccessList e -> a Int e
toArray list = array (0, length list - 1) (toIndexedList list)