-----------------------------------------------------------------------------
-- |
-- 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 Rank2Types #-}

module Data.FMList ( 

    FMList(..)
    
  , empty
  , singleton
  , cons
  , snoc
  , append
  , flatten
  
  , toList
  , fromList
  
  , null
  , head
  , tail
  , last
  , init
  , reverse
  
  , filter
  , take
  , drop
  , takeWhile
  , dropWhile
  
  , zip
  , zipWith
  
  , iterate
  , repeat
  , unfoldr
  , unfoldl
  
  ) where

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

newtype FMList a = FM { unFM :: forall b . Monoid b => (a -> b) -> b }

-- Moved to Alternative instance
-- empty       :: FMList a
-- empty       = FM $ \f -> 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

flatten     :: FMList (FMList a) -> FMList a
flatten l   = FM $ \f -> unFM l (foldMap f)

fromList    :: [a] -> FMList a
fromList    = List.foldr cons empty

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

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 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 . reverse $ l

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

filter :: (a -> Bool) -> FMList a -> FMList a
filter p l = FM $ \f -> unFM l $ \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 n > 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 n <= 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 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 c = FM $ \f -> unFM c (f . 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  = FM $ \f -> unFM m (foldMap f . g)

instance Applicative FMList where
  pure     = return
  (<*>)    = ap
    
instance Monoid (FMList a) where
  mempty   = empty
  mappend  = append
  
instance MonadPlus FMList where
  mzero    = empty
  mplus    = append
  
instance Alternative FMList where
  empty    = FM $ \_ -> mempty
  (<|>)    = append
  
instance Show a => Show (FMList a) where
  show l = case show (toList l) of s@(_:_) -> "fromList " ++ s
  
  
fromJustOrError :: Maybe a -> String -> a
fromJustOrError ma e = maybe (error e) id ma