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

import Control.Monad
import Control.Monad.ST
import Data.Foldable (toList)
import Data.List.NonEmpty (NonEmpty(..))

import qualified Data.Vector         as V
import qualified Data.Vector.Mutable as VM

import Data.Falsify.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
_         = forall a. HasCallStack => [Char] -> a
error [Char]
"chunksOfNonEmpty: zero chunk size"
chunksOfNonEmpty Word
n (a
x :| [a]
xs) =
    let ([a]
chunk, [a]
rest) = forall a. Int -> [a] -> ([a], [a])
splitAt (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
n) (a
x forall a. a -> [a] -> [a]
: [a]
xs)
    in case ([a]
chunk, [a]
rest) of
         ([]   , [a]
_)    -> forall a. HasCallStack => [Char] -> a
error [Char]
"impossible"
         (a
c:[a]
cs , [])   -> (a
c forall a. a -> [a] -> NonEmpty a
:| [a]
cs) forall a. a -> [a] -> NonEmpty a
:| []
         (a
c:[a]
cs , a
r:[a]
rs) -> (a
c forall a. a -> [a] -> NonEmpty a
:| [a]
cs) forall a. a -> [a] -> NonEmpty a
:| forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (forall a. Word -> NonEmpty a -> NonEmpty (NonEmpty a)
chunksOfNonEmpty Word
n (a
r forall a. a -> [a] -> NonEmpty a
:| [a]
rs))

{-------------------------------------------------------------------------------
  Permutations
-------------------------------------------------------------------------------}

-- | Permutation is a sequence of swaps
type Permutation = [(Word, Word)]

applyPermutation :: Permutation -> [a] -> [a]
applyPermutation :: forall a. Permutation -> [a] -> [a]
applyPermutation Permutation
p [a]
xs =
    forall a. Vector a -> [a]
V.toList forall a b. (a -> b) -> a -> b
$ forall a.
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
V.modify (forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (forall a b. (a -> b) -> [a] -> [b]
map (Word, Word) -> (Int, Int)
conv Permutation
p) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. MVector s a -> (Int, Int) -> ST s ()
swap) (forall a. [a] -> Vector a
V.fromList [a]
xs)
  where
    swap :: V.MVector s a -> (Int, Int) -> ST s ()
    swap :: forall s a. MVector s a -> (Int, Int) -> ST s ()
swap MVector s a
vec (Int
i, Int
j) = do
        a
x <- forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
VM.read MVector s a
vec Int
i
        a
y <- forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> m a
VM.read MVector s a
vec Int
j
        forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector s a
vec Int
i a
y
        forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> a -> m ()
VM.write MVector s a
vec Int
j a
x

    conv :: (Word, Word) -> (Int, Int)
    conv :: (Word, Word) -> (Int, Int)
conv (Word
i, Word
j) = (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
i, forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
j)

{-------------------------------------------------------------------------------
  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 = forall (t :: * -> *) (f :: * -> *) a.
Foldable t =>
t (Marked f a) -> Word
countKept [Marked f a]
xs
    in if Word
kept forall a. Ord a => a -> a -> Bool
>= Word
n
         then [Marked f a]
xs
         else forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
go (Word
n 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) = forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x forall a. a -> [a] -> [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) = forall (f :: * -> *) a. Mark -> f a -> Marked f a
Marked Mark
Keep f a
x forall a. a -> [a] -> [a]
: forall (f :: * -> *) a. Word -> [Marked f a] -> [Marked f a]
go (Word
n forall a. Num a => a -> a -> a
- Word
1) [Marked f a]
xs