A Monad for taking values from a list in a safe way

ListConsumer.hs

{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.ListConsumer (
  ListConsumer, takeN, peekN, dropN, takeAll, takeAny, runL, runFoldableL
) where

import           Control.Applicative
import           Control.Arrow
import           Control.Monad

data MustHave = Zero | One

data ListConsumer a b where
  LError :: ListConsumer a b
  LValue :: b -> ListConsumer a b
  LTake  :: Int      -> ([a] -> ListConsumer a b) -> ListConsumer a b
  LAll   :: MustHave -> ([a] -> ListConsumer a b) -> ListConsumer a b
  LPeek  :: Int      -> ([a] -> ListConsumer a b) -> ListConsumer a b
  LSkip  :: Int      -> (()  -> ListConsumer a b) -> ListConsumer a b

go3 :: ListConsumer a t -> (ListConsumer a t -> ListConsumer a b) -> ListConsumer a b
go3 (LTake n fn) go = LTake n (go.fn)
go3 (LPeek n fn) go = LPeek n (go.fn)
go3 (LAll  n fn) go = LAll  n (go.fn)
go3 (LSkip n fn) go = LSkip n (go.fn)
go3 _ _             = error "applied to non 3 arg go"

instance Functor (ListConsumer a) where
  fmap f = go where
    go LError     = LError
    go (LValue b) = LValue (f b)
    go l          = go3 l go

instance Applicative (ListConsumer a) where
  pure = LValue
  fa <*> a = go fa where
    go LError      = LError
    go (LValue fv) = fmap fv a
    go l           = go3 l go

instance Monad (ListConsumer a) where
  return = pure
  a0 >>= f = go a0 where
    go LError     = LError
    go (LValue a) = f a
    go l          = go3 l go

instance Alternative (ListConsumer a) where
  empty   = LError
  a <|> b = go a where
    go LError = b
    go _      = a

instance MonadPlus (ListConsumer a)


takeN :: Int -> ListConsumer a [a]
takeN n = LTake n return

peekN :: Int -> ListConsumer a [a]
peekN n = LPeek n return

dropN :: Int -> ListConsumer a ()
dropN n = LSkip n return

takeAll :: ListConsumer a [a]
takeAll = LAll One return

takeAny :: ListConsumer a [a]
takeAny = LAll Zero return

tryTake :: MonadPlus m => Int -> [a] -> m ([a], [a])
tryTake 0 xs     = pure ([], xs)
tryTake n (x:xs) = first (x:) <$> tryTake (n-1) xs
tryTake _ _      = empty

tryDrop :: MonadPlus m => Int -> [a] -> m [a]
tryDrop 0 xs     = pure xs
tryDrop n (_:xs) = tryDrop (n-1) xs
tryDrop _ _      = empty

runL :: MonadPlus m => [a] -> ListConsumer a b -> m b
runL = go where
  go _   LError        = empty
  go _  (LValue b)     = pure b
  go xs (LAll Zero fn) = go [] (fn xs)
  go [] (LAll One _)   = empty
  go xs (LAll One fn)  = go [] (fn xs)
  go xs (LTake n fn)   = do
    (as, xs') <- tryTake n xs
    go xs' (fn as)
  go xs (LPeek n fn)   = do
    (as, _) <- tryTake n xs
    go xs (fn as)
  go xs (LSkip n fn)   = do
    xs' <- tryDrop n xs
    go xs' (fn ())

{- Foldable -}

data FoldN a b =
    FoldN    Int ([a] -> FoldN a b) ([a] -> [a])
  | FoldPeek Int ([a] -> FoldN a b) ([a] -> [a])
  | FoldAll      ([a] -> FoldN a b) ([a] -> [a])
  | FoldX    Int (()  -> FoldN a b)
  | FoldV b
  | FoldZ

fGet :: MonadPlus m => FoldN a b -> m b
fGet (FoldV b)      = return b
fGet (FoldAll f as) = fGet (f (as []))
fGet _              = mzero

runFoldableL :: (MonadPlus m, Foldable f) => f a -> ListConsumer a b -> m b
runFoldableL fs x = fGet (foldl foldN (go x) fs) where
  go  LError       = FoldZ
  go (LValue b)    = FoldV b
  go (LAll One fn) = FoldAll (go.fn') id where
    fn' [] = LError
    fn' xs = fn xs
  go (LAll _ fn)   = FoldAll (go.fn) id
  go (LTake n fn)  = FoldN n (go.fn) id
  go (LPeek n fn)  = FoldPeek n (go.fn) id
  go (LSkip n fn)  = FoldX n (go.fn)

  fold3 _    _ 0 _ _  _ = FoldZ
  fold3 onef _ 1 f as z = onef $ f (as z)
  fold3 _    c n f as _ = c (n-1) f as

  asList a as as' = as (a : as')
  foldx n f _ = FoldX n f

  foldN (FoldX    n f)    _ = fold3 id  foldx n f (const ()) undefined
  foldN (FoldN    n f as) a = fold3 id  FoldN n f (asList a as) [a]
  foldN (FoldPeek n f as) a = let one f' = foldl foldN f' (as [a]) in
                              fold3 one FoldN n f (asList a as) [a]
  foldN (FoldAll f as) a    = FoldAll f (asList a as)
  foldN v                 _ = v