module Data.Falsify.List (
chunksOfNonEmpty
, Permutation
, applyPermutation
, 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
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))
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)
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