{-| Module      : Data.List.Filter
    Description : Special takes and drops on lists
    Copyright   : (c) Preetham Gujjula, 2020
    License     : BSD3
    Maintainer  : pgujjula+list-utilities@protonmail.com
    Stability   : experimental

Special takes and drops on lists.
-}
module Data.List.Filter (
      takeEvery
    , dropEvery
    , takeUntil
    , dropUntil
    ) where

{-| @takeEvery n xs@ is a list of every @n@th element of @xs@.

    __Precondition:__ @n@ must be positive.

    >>> takeEvery 3 [1..10]
    [3, 6, 9]
    >>> takeEvery 1 [1..10] == [1..10]
    True
-}
takeEvery :: Int -> [a] -> [a]
takeEvery :: Int -> [a] -> [a]
takeEvery step :: Int
step xs :: [a]
xs = [a] -> [a]
forall a. [a] -> [a]
compute [a]
validated
  where
    compute :: [a] -> [a]
compute ys :: [a]
ys = case Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
step Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1) [a]
ys of
                     []    -> []
                     y :: a
y:ys' :: [a]
ys' -> a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a]
compute [a]
ys'
    validated :: [a]
validated
        | Int
step Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0  = [a]
xs
        | Bool
otherwise = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error ([Char] -> [a]) -> [Char] -> [a]
forall a b. (a -> b) -> a -> b
$ "Data.List.Transform.takeEvery: Step parameter "
                           [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ "must be positive."

{-| @dropEvery n xs@ is a list of every @n@th element of @xs@.

    __Precondition:__ @n@ must be positive.

    >>> dropEvery 3 [1..10]
    [1, 2, 4, 5, 7, 8, 10]
    >>> dropEvery 1 [1..10]
    []
-}
dropEvery :: Int -> [a] -> [a]
dropEvery :: Int -> [a] -> [a]
dropEvery step :: Int
step xs :: [a]
xs = [a] -> [a]
forall a. [a] -> [a]
compute [a]
validated
  where
    compute :: [a] -> [a]
compute ys :: [a]
ys = case Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt (Int
step Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1) [a]
ys of
                     (as :: [a]
as, [])   -> [a]
as
                     (as :: [a]
as, _:bs :: [a]
bs) -> [a]
as [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a] -> [a]
compute [a]
bs
    validated :: [a]
validated
        | Int
step Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> 0  = [a]
xs
        | Bool
otherwise = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error ([Char] -> [a]) -> [Char] -> [a]
forall a b. (a -> b) -> a -> b
$ "Data.List.Transform.dropEvery: Step parameter "
                           [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ "must be positive."

{-| Take a list until a predicate is satisfied, and include the element
    satisfying the predicate.

    >>> takeUntil (== 5) [1..]
    [1, 2, 3, 4, 5]
    >>> takeUntil (== 7) [3, 2, 1]
    [3, 2, 1]
    >>> takeUntil undefined []
    []

    Note that @takeUntil@ on a nonempty list must always yield the first
    element, and the implementation is lazy enough to take advantage of this
    fact.

    >>> head (takeUntil undefined [1..])
    1
-}
takeUntil :: (a -> Bool) -> [a] -> [a]
takeUntil :: (a -> Bool) -> [a] -> [a]
takeUntil _ []     = []
takeUntil _ [x :: a
x]    = [a
x]
takeUntil f :: a -> Bool
f (x :: a
x:xs :: [a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (if a -> Bool
f a
x then [] else (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
takeUntil a -> Bool
f [a]
xs)

{-| Drop a list until a predicate is satisfied, and include the element
    satisfying the predicate.

    >>> dropUntil (== 5) [1..10]
    [5, 6, 7, 8, 9, 10]
    >>> dropUntil (< 0) [1, 2, 3]
    []
    >>> dropUntil undefined []
    []
-}
dropUntil :: (a -> Bool) -> [a] -> [a]
dropUntil :: (a -> Bool) -> [a] -> [a]
dropUntil _ []     = []
dropUntil f :: a -> Bool
f (x :: a
x:xs :: [a]
xs) = if a -> Bool
f a
x then a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs else (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
dropUntil a -> Bool
f [a]
xs