{-# LANGUAGE BangPatterns #-}

module QC.Rechunked (
      rechunkBS
    , rechunkT
    ) where

import Control.Monad (forM, forM_)
import Control.Monad.ST (ST, runST)
import Data.List (unfoldr)
import Test.QuickCheck (Gen, choose)
import qualified Data.ByteString as B
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M

rechunkBS :: B.ByteString -> Gen [B.ByteString]
rechunkBS = fmap (map B.copy) . rechunk_ B.splitAt B.length

rechunkT :: T.Text -> Gen [T.Text]
rechunkT = fmap (map T.copy) . rechunk_ T.splitAt T.length

rechunk_ :: (Int -> a -> (a,a)) -> (a -> Int) -> a -> Gen [a]
rechunk_ split len xs = (unfoldr go . (,) xs) `fmap` rechunkSizes (len xs)
  where go (b,r:rs)   = Just (h, (t,rs))
          where (h,t) = split r b
        go (_,_)      = Nothing

rechunkSizes :: Int -> Gen [Int]
rechunkSizes n0 = shuffle =<< loop [] (0:repeat 1) n0
  where loop _ [] _ = error "it's 2014, where's my Stream type?"
        loop acc (lb:lbs) n
          | n <= 0 = shuffle (reverse acc)
          | otherwise = do
            !i <- choose (lb,n)
            loop (i:acc) lbs (n-i)

shuffle :: [Int] -> Gen [Int]
shuffle (0:xs) = (0:) `fmap` shuffle xs
shuffle xs = fisherYates xs

fisherYates :: [a] -> Gen [a]
fisherYates xs = (V.toList . V.backpermute v) `fmap` swapIndices (G.length v)
  where
    v = V.fromList xs
    swapIndices n0 = do
        swaps <- forM [0..n] $ \i -> ((,) i) `fmap` choose (i,n)
        return (runST (swapAll swaps))
      where
        n = n0 - 1
        swapAll :: [(Int,Int)] -> ST s (V.Vector Int)
        swapAll ijs = do
          mv <- G.unsafeThaw (G.enumFromTo 0 n :: V.Vector Int)
          forM_ ijs $ uncurry (M.swap mv)
          G.unsafeFreeze mv