-----------------------------------------------------------------------------
-- |
-- Module      :  Data.CappedList
-- Copyright   :  (c) 2010 John Millikin
-- License     :  BSD3
--
-- Maintainer  :  jmillikin@gmail.com
-- Portability :  portable
--
-- A list-like type for lazy sequences, with a user-defined termination value.
--
-- This module uses common names and so is designed to be imported qualified:
--
-- > import qualified Data.CappedList as CL
--
-----------------------------------------------------------------------------

module Data.CappedList
	( CappedList (..)
	, null
	, append
	, map
	, mapEither
	, concatMap
	, foldr
	, foldl
	, unfoldr
	, length
	) where
import Prelude hiding (null, map, concatMap, foldl, foldr, length)
import qualified Prelude as Prelude
import Data.Monoid (Monoid, mempty, mappend)
import Control.Monad (MonadPlus, mzero, mplus, ap)
import Control.Monad.Fix (MonadFix, fix, mfix)
import qualified Control.Applicative as A
import Control.Applicative ((<$>), (<*>))
import Data.Empty (Empty (empty))
import qualified Data.Foldable as F
import qualified Data.Traversable as T

-- | A list-like type for lazy sequences, with a user-defined termination value.
--
data CappedList cap a
	= Next a (CappedList cap a)
	| Cap cap
	deriving (Eq, Show)

-- TODO: deriving Data?

instance Empty cap => Empty (CappedList cap a) where
	empty = Cap empty

instance Functor (CappedList cap) where
	fmap = map

instance Empty cap => Monoid (CappedList cap a) where
	mempty = Cap empty
	mappend = append

instance Empty cap => Monad (CappedList cap) where
	Cap x       >>= _ = Cap x
	(Next x xs) >>= k = append (k x) (xs >>= k)
	return = flip Next (Cap empty)

instance Empty cap => MonadPlus (CappedList cap) where
	mzero = Cap empty
	mplus = append

instance Empty cap => A.Alternative (CappedList cap) where
	empty = Cap empty
	(<|>) = append

instance Empty cap => A.Applicative (CappedList cap) where
	pure = flip Next (Cap empty)
	(<*>) = ap

instance F.Foldable (CappedList cap) where
	foldMap = T.foldMapDefault

instance T.Traversable (CappedList cap) where
	sequenceA (Cap x)     = A.pure $ Cap x
	sequenceA (Next f fs) = Next <$> f <*> T.sequenceA fs

instance Empty cap => MonadFix (CappedList cap) where
	mfix f = case fix (f . head') of
		Cap x      -> Cap x
		(Next x _) -> append (return x) (mfix (tail' . f))

-- | Like the standard 'Prelude.null' function.
--
null :: CappedList cap a -> Bool
null (Next _ _) = False
null _          = True

-- | Like the standard '++' function.
--
append :: CappedList cap a -> CappedList cap a -> CappedList cap a
append x@(Cap _)   _ = x
append (Next x xs) y = Next x (append xs y)

-- | Like the standard 'Prelude.map' function.
--
map :: (a -> b) -> CappedList cap a -> CappedList cap b
map f = foldr (Next . f) Cap

{-# INLINE [1] map #-}

{-# RULES
  "map/map" forall f g t.
   map f (map g t) = map (f . g) t
  #-}

-- | Like the standard 'Prelude.map' function, but the mapping function may
-- return a capping value.
--
mapEither :: (a -> Either cap b) -> CappedList cap a -> CappedList cap b
mapEither f = foldr (\a acc -> either Cap (flip Next acc) (f a)) Cap

{-# INLINE [1] mapEither #-}

{-# RULES
  "mapEither/mapEither" forall f g t.
   mapEither f (mapEither g t) = mapEither (either Left g . f) t
  #-}

-- | Like the standard 'Prelude.concatMap' function.
--
concatMap :: (a -> CappedList cap b) -> CappedList cap a -> CappedList cap b
concatMap f = foldr (append . f) Cap

-- | Like the standard 'Prelude.foldr' function, but accepting an extra
-- parameter to handle 'Cap' values.
--
foldr :: (a -> b -> b) -> (cap -> b) -> CappedList cap a -> b
foldr f z = foldr' where
	foldr' (Cap x)     = z x
	foldr' (Next x xs) = f x (foldr' xs)

-- | Like the standard 'Prelude.foldl' function, but accepting an extra
-- parameter to handle 'Cap' values.
--
foldl :: (b -> a -> b) -> (cap -> b) -> CappedList cap a -> b
foldl f = foldl' where
	foldl' z (Cap x)     = z x
	foldl' z (Next x xs) = foldl' (\cap -> f (z cap) x) xs

-- | Like the standard 'Data.List.unfoldr' function, but the step function
-- must return a cap to terminate the unfold.
--
unfoldr :: (b -> Either cap (a, b)) -> b -> CappedList cap a
unfoldr f = unfoldr' where
	unfoldr' x = case f x of
		Left cap -> Cap cap
		Right (a, b) -> Next a (unfoldr' b)

-- | Like the standard 'Prelude.length' function; 'Cap' is considered
-- 0-length.
--
length :: CappedList cap a -> Int
length = foldr (const (1 +)) (const 0)

-- Partial functions, used only to implement mfix; not exported
head' :: CappedList cap a -> a
head' (Cap _)    = error "mfix CappedList: Cap"
head' (Next x _) = x

tail' :: CappedList cap a -> CappedList cap a
tail' (Cap _)     = error "mfix CappedList: Cap"
tail' (Next _ xs) = xs