File: Rechunked.hs

package info (click to toggle)
haskell-attoparsec 0.14.4-4
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 880 kB
  • sloc: haskell: 4,749; ansic: 170; makefile: 22
file content (56 lines) | stat: -rw-r--r-- 1,805 bytes parent folder | download | duplicates (5)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
{-# 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 2015, 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