module Data.Falsify.Internal.List (
    -- * Splitting
    chunksOfNonEmpty
    -- * Dealing with marks
  , keepAtLeast
  ) where

import Data.Foldable (toList)
import Data.List.NonEmpty (NonEmpty(..))

import Test.Falsify.Marked (Mark(..), Marked(..))
import qualified Test.Falsify.Marked as Marked

{-------------------------------------------------------------------------------
  Splitting
-------------------------------------------------------------------------------}

-- | Take chunks of a non-empty list
--
-- This is lazy:
--
-- >    NE.take 4 $ chunksOfNonEmpty 3 (0 :| [1..])
-- > == [ 0 :| [1,2]
-- >    , 3 :| [4,5]
-- >    , 6 :| [7,8]
-- >    , 9 :| [10,11]
-- >    ]
chunksOfNonEmpty :: Word -> NonEmpty a -> NonEmpty (NonEmpty a)
chunksOfNonEmpty :: forall a. Word -> NonEmpty a -> NonEmpty (NonEmpty a)
chunksOfNonEmpty Word
0 NonEmpty a
_         = [Char] -> NonEmpty (NonEmpty a)
forall a. HasCallStack => [Char] -> a
error [Char]
"chunksOfNonEmpty: zero chunk size"
chunksOfNonEmpty Word
n (a
x :| [a]
xs) =
    let ([a]
chunk, [a]
rest) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n) (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs)
    in case ([a]
chunk, [a]
rest) of
         ([]   , [a]
_)    -> [Char] -> NonEmpty (NonEmpty a)
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"
         (a
c:[a]
cs , [])   -> (a
c a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
cs) NonEmpty a -> [NonEmpty a] -> NonEmpty (NonEmpty a)
forall a. a -> [a] -> NonEmpty a
:| []
         (a
c:[a]
cs , a
r:[a]
rs) -> (a
c a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
cs) NonEmpty a -> [NonEmpty a] -> NonEmpty (NonEmpty a)
forall a. a -> [a] -> NonEmpty a
:| NonEmpty (NonEmpty a) -> [NonEmpty a]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Word -> NonEmpty a -> NonEmpty (NonEmpty a)
forall a. Word -> NonEmpty a -> NonEmpty (NonEmpty a)
chunksOfNonEmpty Word
n (a
r a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
rs))

{-------------------------------------------------------------------------------
  Dealing with marks
-------------------------------------------------------------------------------}

keepAtLeast :: Word -> [Marked f a] -> [Marked f a]
keepAtLeast :: forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
keepAtLeast = \Word
n [Marked f a]
xs ->
    let kept :: Word
kept = [Marked f a] -> Word
forall (t :: * -> *) (f :: * -> *) a.
Foldable t =>
t (Marked f a) -> Word
Marked.countKept [Marked f a]
xs
    in if Word
kept Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Word
n
         then [Marked f a]
xs
         else Word -> [Marked f a] -> [Marked f a]
forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
go (Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
kept) [Marked f a]
xs
  where
    go :: Word -> [Marked f a] -> [Marked f a]
    go :: forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
go Word
_ []                 = []
    go Word
0 [Marked f a]
xs                 = [Marked f a]
xs
    go Word
n (Marked Mark
Keep f a
x:[Marked f a]
xs) = Mark -> f a -> Marked f a
forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x Marked f a -> [Marked f a] -> [Marked f a]
forall a. a -> [a] -> [a]
: Word -> [Marked f a] -> [Marked f a]
forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
go  Word
n      [Marked f a]
xs
    go Word
n (Marked Mark
Drop f a
x:[Marked f a]
xs) = Mark -> f a -> Marked f a
forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x Marked f a -> [Marked f a] -> [Marked f a]
forall a. a -> [a] -> [a]
: Word -> [Marked f a] -> [Marked f a]
forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
go (Word
n Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1) [Marked f a]
xs