-----------------------------------------------------------------------------
-- |
-- Module      :  Data.FMList
-- Copyright   :  (c) Sjoerd Visscher 2009
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  sjoerd
-- Stability   :  experimental
-- Portability :  portable
--
-- FoldMap lists: lists represented by their foldMap function.
--
-----------------------------------------------------------------------------
{-# LANGUAGE RankNTypes #-}

module Data.FMList ( 

    FMList(..)
  , transform
    
  -- * Construction  
  , empty
  , singleton
  , cons
  , snoc
  , append
  
  , toList
  , fromList
  , fromFoldable
  
  -- * Basic functions
  , null
  , length
  , genericLength
  
  , head
  , tail
  , last
  , init
  , reverse
  
  -- * Folding
  , flatten
  , filter
  , take
  , drop
  , takeWhile
  , dropWhile
  
  , zip
  , zipWith
  
  -- * Unfolding
  , iterate
  , repeat
  , unfoldr
  , unfoldl
  
  ) where

import Prelude 
  ( (.), ($), ($!), flip, const, id, error
  , Maybe(..), maybe
  , Bool(..), (||), not
  , Ord(..), Num(..), Int
  , Show(..), String, (++)
  )
import qualified Data.List as List
import Data.Monoid
import Data.Foldable
import Data.Traversable
import Control.Monad
import Control.Applicative


-- | FMList is a foldMap function wrapped up in a newtype.
-- Examples:
--
-- > -- A right-infinite list
-- > c = 1 `cons` c
--
-- > -- A left-infinite list
-- > d = d `snoc` 2
--
-- > -- A middle-infinite list ??
-- > e = c `append` d
--
-- > *> head e
-- > 1
-- > *> last e
-- > 2
-- 
newtype FMList a = FM { unFM :: forall b . Monoid b => (a -> b) -> b }

-- | Transform transforms a list by changing the map function that is passed to foldMap.
-- It has the following property:
--
-- @transform a . transform b = transform (b . a)@
--
-- For example:
--
--  * @  m >>= g@
--
--  * @= flatten (fmap g m)@
--
--  * @= flatten . fmap g $ m@
--
--  * @= transform foldMap . transform (. g) $ m@
--
--  * @= transform ((. g) . foldMap) m@
--
--  * @= transform (\\f -> foldMap f . g) m@
-- 
transform :: (forall b. Monoid b => (a -> b) -> (c -> b)) -> FMList c -> FMList a
transform t l = FM $ \f -> unFM l (t f)

-- nil is exported as empty from Applicative
nil          :: FMList a
nil          = FM $ \_ -> mempty

singleton    :: a -> FMList a
singleton x  = FM $ \f -> f x

cons         :: a -> FMList a -> FMList a
cons x l     = FM $ \f -> f x `mappend` unFM l f

snoc         :: FMList a -> a -> FMList a
snoc l x     = FM $ \f -> unFM l f `mappend` f x

append       :: FMList a -> FMList a -> FMList a
-- append l r  = FM $ \f -> unFM l f `mappend` unFM r f
append l r   = FM $ unFM l `mappend` unFM r


fromList     :: [a] -> FMList a
fromList     = fromFoldable

fromFoldable :: Foldable f => f a -> FMList a
fromFoldable l = FM $ flip foldMap l


null         :: FMList a -> Bool
null         = foldr (\_ _ -> False) True

length       :: FMList a -> Int
length       = genericLength

genericLength :: Num b => FMList a -> b
genericLength l = getSum $ unFM l (const $ Sum 1)


head         :: FMList a -> a
head l       = getFirst (unFM l (First . Just)) `fromJustOrError` "Data.FMList.head: empty list"

tail         :: FMList a -> FMList a
tail l       = if null l then error "Data.FMList.tail: empty list" else drop (1::Int) l

last         :: FMList a -> a
last l       = getLast (unFM l (Last . Just)) `fromJustOrError` "Data.FMList.last: empty list"

init         :: FMList a -> FMList a
init l       = if null l then error "Data.FMList.init: empty list" else reverse . drop (1::Int) . reverse $ l

reverse      :: FMList a -> FMList a
reverse l    = FM $ \f -> getDual $ unFM l (Dual . f)


flatten      :: FMList (FMList a) -> FMList a
flatten      = transform foldMap

filter       :: (a -> Bool) -> FMList a -> FMList a
filter p     = transform (\f e -> if p e then f e else mempty)

take         :: (Ord n, Num n) => n -> FMList a -> FMList a
take n l     = FM $ \f -> 
  foldr (\e r i -> if i > 0 then mappend (f e) (r (i-1)) else mempty) (const mempty) l n

drop         :: (Ord n, Num n) => n -> FMList a -> FMList a
drop n l     = FM $ \f -> 
  foldr (\e r i -> if i <= 0 then mappend (f e) (r i) else r (i-1)) (const mempty) l n

takeWhile    :: (a -> Bool) -> FMList a -> FMList a
takeWhile p l = FM $ \f -> 
  foldr (\e r -> if p e then mappend (f e) r else mempty) mempty l

dropWhile    :: (a -> Bool) -> FMList a -> FMList a
dropWhile p l = FM $ \f -> 
  foldr (\e r done -> if done || not (p e) then mappend (f e) (r True) else r False) (const mempty) l False


zipWith      :: (a -> b -> c) -> FMList a -> FMList b -> FMList c
zipWith t l1 l2 = FM $ \f -> 
  foldr (\e1 r r2 -> 
	  foldr (\e2 _ -> mappend (f (t e1 e2)) (r (drop (1::Int) r2))) mempty r2) (const mempty) l1 l2

zip          :: FMList a -> FMList b -> FMList (a,b)
zip          = zipWith (,)


iterate      :: (a -> a) -> a -> FMList a
iterate f x  = x `cons` iterate f (f x)

repeat       :: a -> FMList a
repeat x     = xs where xs = x `cons` xs

unfoldr      :: (b -> Maybe (a, b)) -> b -> FMList a
unfoldr pf b = FM $ \f -> u f mempty (pf b) where
  u _ acc Nothing = acc
  u f acc (Just (a, b')) = u f (acc `mappend` f a) (pf b')

unfoldl      :: (b -> Maybe (b, a)) -> b -> FMList a
unfoldl pf b = FM $ \f -> u f mempty (pf b) where
  u _ acc Nothing = acc
  u f acc (Just (b', a)) = u f (f a `mappend` acc) (pf b')


instance Functor FMList where
  fmap g     = transform (. g)
  
instance Foldable FMList where
  foldMap    = flip unFM
  
instance Traversable FMList where
  traverse f = foldr cons_f (pure empty) where cons_f x ys = cons <$> f x <*> ys

instance Monad FMList where
  return     = singleton
  m >>= g    = transform (\f -> foldMap f . g) m

instance Applicative FMList where
  pure       = return
  gs <*> xs  = transform (\f g -> unFM xs (f . g)) gs
    
instance Monoid (FMList a) where
  mempty     = nil
  mappend    = append
  
instance MonadPlus FMList where
  mzero      = nil
  mplus      = append
  
instance Alternative FMList where
  empty      = nil
  (<|>)      = append
  
instance Show a => Show (FMList a) where
  show l     = "fromList " ++ (show $! toList l)
  
  
fromJustOrError :: Maybe a -> String -> a
fromJustOrError ma e = maybe (error e) id ma