-- Copyright 2013 Evan Laforge
-- This program is distributed under the terms of the GNU General Public
-- License 3.0, see COPYING or http://www.gnu.org/licenses/gpl-3.0.txt

-- | List functions with continuations.  This allows you to chain them and
-- easily express things like 'take until f then take one more'.
module Util.Then where
import Prelude hiding (break, span, take, takeWhile, map, mapM)
import qualified Data.List as List


takeWhile :: (a -> Bool) -> ([a] -> [a]) -> [a] -> [a]
takeWhile :: forall a. (a -> Bool) -> ([a] -> [a]) -> [a] -> [a]
takeWhile a -> Bool
_ [a] -> [a]
cont [] = [a] -> [a]
cont []
takeWhile a -> Bool
f [a] -> [a]
cont (a
x:[a]
xs)
    | a -> Bool
f a
x = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> Bool) -> ([a] -> [a]) -> [a] -> [a]
forall a. (a -> Bool) -> ([a] -> [a]) -> [a] -> [a]
takeWhile a -> Bool
f [a] -> [a]
cont [a]
xs
    | Bool
otherwise = [a] -> [a]
cont (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs)

-- | takeWhile plus one extra
takeWhile1 :: (a -> Bool) -> [a] -> [a]
takeWhile1 :: forall a. (a -> Bool) -> [a] -> [a]
takeWhile1 a -> Bool
f = (a -> Bool) -> ([a] -> [a]) -> [a] -> [a]
forall a. (a -> Bool) -> ([a] -> [a]) -> [a] -> [a]
takeWhile a -> Bool
f (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
List.take Int
1)

take :: Int -> ([a] -> [a]) -> [a] -> [a]
take :: forall a. Int -> ([a] -> [a]) -> [a] -> [a]
take Int
_ [a] -> [a]
cont [] = [a] -> [a]
cont []
take Int
n [a] -> [a]
cont (a
x:[a]
xs)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [a] -> [a]
cont (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs)
    | Bool
otherwise = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> ([a] -> [a]) -> [a] -> [a]
forall a. Int -> ([a] -> [a]) -> [a] -> [a]
take (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a] -> [a]
cont [a]
xs

filter :: (a -> Bool) -> (a -> Bool) -> ([a] -> [a]) -> [a] -> [a]
filter :: forall a. (a -> Bool) -> (a -> Bool) -> ([a] -> [a]) -> [a] -> [a]
filter a -> Bool
f a -> Bool
done [a] -> [a]
cont = [a] -> [a]
go
    where
    go :: [a] -> [a]
go [] = []
    go (a
x:[a]
xs)
        | a -> Bool
done a
x = [a] -> [a]
cont (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs)
        | a -> Bool
f a
x = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
go [a]
xs
        | Bool
otherwise = [a] -> [a]
go [a]
xs

map :: (a -> b) -> [b] -> [a] -> [b]
map :: forall a b. (a -> b) -> [b] -> [a] -> [b]
map a -> b
f [b]
bs = [a] -> [b]
go
    where
    go :: [a] -> [b]
go [] = [b]
bs
    go (a
a:[a]
as) = a -> b
f a
a b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [a] -> [b]
go [a]
as

-- | Like 'List.mapAccumL', except that you can do something with the final
-- state and append that to the list.
mapAccumL :: (acc -> x -> (acc, y)) -> acc -> (acc -> [y]) -> [x] -> [y]
mapAccumL :: forall acc x y.
(acc -> x -> (acc, y)) -> acc -> (acc -> [y]) -> [x] -> [y]
mapAccumL acc -> x -> (acc, y)
f acc
acc acc -> [y]
cont = acc -> [x] -> [y]
go acc
acc
    where
    go :: acc -> [x] -> [y]
go acc
acc [] = acc -> [y]
cont acc
acc
    go acc
acc (x
x:[x]
xs) = y
y y -> [y] -> [y]
forall a. a -> [a] -> [a]
: acc -> [x] -> [y]
go acc
acc2 [x]
xs
        where (acc
acc2, y
y) = acc -> x -> (acc, y)
f acc
acc x
x

mapM :: Monad m => (a -> m b) -> m [b] -> [a] -> m [b]
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> m [b] -> [a] -> m [b]
mapM a -> m b
f m [b]
cont = [a] -> m [b]
go
    where
    go :: [a] -> m [b]
go [] = m [b]
cont
    go (a
a:[a]
as) = do
        b
b <- a -> m b
f a
a
        (b
b:) ([b] -> [b]) -> m [b] -> m [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> m [b]
go [a]
as

break :: (a -> Bool) -> ([a] -> ([a], rest))
    -- ^ Given the list after the break, return (pre, post), where pre will
    -- be appended to the end of the first list.
    -> [a] -> ([a], rest)
break :: forall a rest.
(a -> Bool) -> ([a] -> ([a], rest)) -> [a] -> ([a], rest)
break a -> Bool
f [a] -> ([a], rest)
cont (a
x:[a]
xs)
    | a -> Bool
f a
x = [a] -> ([a], rest)
cont (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs)
    | Bool
otherwise = let ([a]
pre, rest
post) = (a -> Bool) -> ([a] -> ([a], rest)) -> [a] -> ([a], rest)
forall a rest.
(a -> Bool) -> ([a] -> ([a], rest)) -> [a] -> ([a], rest)
break a -> Bool
f [a] -> ([a], rest)
cont [a]
xs in (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
pre, rest
post)
break a -> Bool
_ [a] -> ([a], rest)
cont [] = [a] -> ([a], rest)
cont []

span :: (a -> Bool) -> ([a] -> ([a], rest)) -> [a] -> ([a], rest)
span :: forall a rest.
(a -> Bool) -> ([a] -> ([a], rest)) -> [a] -> ([a], rest)
span a -> Bool
f = (a -> Bool) -> ([a] -> ([a], rest)) -> [a] -> ([a], rest)
forall a rest.
(a -> Bool) -> ([a] -> ([a], rest)) -> [a] -> ([a], rest)
break (Bool -> Bool
not (Bool -> Bool) -> (a -> Bool) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

-- | Break right after the function returns True.
break1 :: (a -> Bool) -> [a] -> ([a], [a])
break1 :: forall a. (a -> Bool) -> [a] -> ([a], [a])
break1 a -> Bool
f = (a -> Bool) -> ([a] -> ([a], [a])) -> [a] -> ([a], [a])
forall a rest.
(a -> Bool) -> ([a] -> ([a], rest)) -> [a] -> ([a], rest)
break a -> Bool
f (Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
1)