@@ -0,0 +1,54 @@
1+{-# LANGUAGE BangPatterns #-}
2+
3+module QC.Rechunked (
4+ rechunkBS
5+ , rechunkT
6+ ) where
7+
8+import Control.Monad (forM, forM_)
9+import Control.Monad.ST (runST)
10+import Data.List (unfoldr)
11+import Test.QuickCheck (Gen, choose)
12+import qualified Data.ByteString as B
13+import qualified Data.Text as T
14+import qualified Data.Vector as V
15+import qualified Data.Vector.Generic as G
16+import qualified Data.Vector.Generic.Mutable as M
17+
18+rechunkBS :: B.ByteString -> Gen [B.ByteString]
19+rechunkBS = fmap (map B.copy) . rechunk_ B.splitAt B.length
20+
21+rechunkT :: T.Text -> Gen [T.Text]
22+rechunkT = fmap (map T.copy) . rechunk_ T.splitAt T.length
23+
24+rechunk_ :: (Int -> a -> (a,a)) -> (a -> Int) -> a -> Gen [a]
25+rechunk_ split len xs = (unfoldr go . (,) xs) `fmap` rechunkSizes (len xs)
26+ where go (b,r:rs) = Just (h, (t,rs))
27+ where (h,t) = split r b
28+ go (_,_) = Nothing
29+
30+rechunkSizes :: Int -> Gen [Int]
31+rechunkSizes n0 = shuffle =<< loop [] (0:repeat 1) n0
32+ where loop acc (lb:lbs) n
33+ | n <= 0 = shuffle (reverse acc)
34+ | otherwise = do
35+ !i <- choose (lb,n)
36+ loop (i:acc) lbs (n-i)
37+
38+shuffle :: [Int] -> Gen [Int]
39+shuffle (0:xs) = (0:) `fmap` shuffle xs
40+shuffle xs = fisherYates xs
41+
42+fisherYates :: [a] -> Gen [a]
43+fisherYates xs = (V.toList . V.backpermute v) `fmap` swapIndices (G.length v)
44+ where
45+ v = V.fromList xs
46+ swapIndices n0 = do
47+ swaps <- forM [0..n] $ \i -> ((,) i) `fmap` choose (i,n)
48+ return (runST (swapAll swaps))
49+ where
50+ n = n0 - 1
51+ swapAll ijs = do
52+ mv <- G.unsafeThaw (G.enumFromTo 0 n :: V.Vector Int)
53+ forM_ ijs $ uncurry (M.swap mv)
54+ G.unsafeFreeze mv