-----------------------------------------------------------------------------
-- |
-- 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
--
-----------------------------------------------------------------------------

{-# LANGUAGE DeriveDataTypeable #-}
module Data.CappedList
	( CappedList (..)
	, toList
	, toList_
	, fromList
	, null
	, appendR
	, appendL
	, appendR_
	, appendL_
	, appendWith
	, map
	, mapEither
	, concatMapM
	, foldr
	, foldl
	, unfoldr
	, length
	) where
import Prelude hiding (null, map, 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 qualified Data.Foldable as F
import qualified Data.Traversable as T
import Data.Data (Data)
import Data.Typeable (Typeable)

-- | 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, Data, Typeable)

instance Functor (CappedList cap) where
	fmap = map

instance Monoid cap => Monoid (CappedList cap a) where
	mempty = Cap mempty
	mappend = appendWith mappend

instance Monoid cap => Monad (CappedList cap) where
	(>>=) = flip concatMapM
	return = flip Next (Cap mempty)

instance Monoid cap => MonadPlus (CappedList cap) where
	mzero = mempty
	mplus = mappend

instance Monoid cap => A.Applicative (CappedList cap) where
	pure = return
	(<*>) = ap

instance Monoid cap => A.Alternative (CappedList cap) where
	empty = mempty
	(<|>) = mappend

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

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

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

-- | Convert a capped list to a standard list.
--
-- The cap is returned in the first value of the result tuple.
--
toList :: CappedList cap a -> (cap, [a])
toList = foldl (\(cap, xs) x -> (cap, x:xs)) (\cap -> (cap, []))

-- | Convert a capped list to a standard list, discarding the cap.
--
toList_ :: CappedList cap a -> [a]
toList_ = foldr (:) (const [])

-- | Convert a standard list and cap to a capped list.
--
fromList :: [a] -> cap -> CappedList cap a
fromList []     = Cap
fromList (x:xs) = Next x . fromList xs

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

-- | Like the standard '++' function.
--
appendL :: CappedList cap1 a -> CappedList cap2 a -> (cap2, CappedList cap1 a)
appendL xs          (Cap y)     = (y, xs)
appendL x@(Cap _)   (Next y ys) = let (cap, z) = appendL x ys in (cap, Next y z)
appendL (Next x xs) ys          = let (cap, z) = appendL xs ys in (cap, Next x z)

-- | Like the standard '++' function.
--
appendR :: CappedList cap1 a -> CappedList cap2 a -> (cap1, CappedList cap2 a)
appendR (Next x xs) ys = let (cap, z) = appendR xs ys in (cap, Next x z)
appendR (Cap x)     ys = (x, ys)

-- | Like the standard '++' function.
--
-- The second list's \"cap\" will be discarded; to preserve the cap, use
-- 'appendL'.
--
appendL_ :: CappedList cap1 a -> CappedList cap2 a -> CappedList cap1 a
appendL_ xs          (Cap _)     = xs
appendL_ (Next x xs) ys          = Next x (appendL_ xs ys)
appendL_ xs          (Next y ys) = Next y (appendL_ xs ys)

-- | Like the standard '++' function.
--
-- The first list's \"cap\" will be discarded; to preserve the cap, use
-- 'appendR'.
--
appendR_ :: CappedList cap1 a -> CappedList cap2 a -> CappedList cap2 a
appendR_ (Next x xs) = Next x . appendR_ xs
appendR_ _           = id

-- | Append two capped lists, merging the caps together using a user-provided
-- function.
--
appendWith :: (c -> d -> e) -> CappedList c a -> CappedList d a
           -> CappedList e a
appendWith f (Cap x)     (Cap y)     = Cap (f x y)
appendWith f (Next x xs) ys          = Next x (appendWith f xs ys)
appendWith f xs          (Next y ys) = Next y (appendWith f xs ys)

-- | 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.
--
concatMapM :: Monoid cap => (a -> CappedList cap b) -> CappedList cap a
           -> CappedList cap b
concatMapM f = foldr (mappend . 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