{-# LANGUAGE ForeignFunctionInterface #-}
import System.Random.Mersenne.Pure64
import Control.Concurrent
import Control.Monad
import Data.Binary.Put
import qualified Data.ByteString.Lazy as BS
import qualified Data.ByteString as S
import Data.Word
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import System.Environment
import GHC.Conc
import Foreign.C.Types
import Debug.Trace
import System.IO
--
-- Parallel unfold and reduce.
--
-- -lcrypto
--
main :: IO ()
main = do
args <- getArgs
case args of
[file, size] -> checkReasonableAndGo (Just (show numCapabilities)) file size
[cpus, file, size] -> checkReasonableAndGo (Just cpus) file size
_ -> dumpUsageAndQuit
dumpUsageAndQuit :: IO ()
dumpUsageAndQuit = mapM_ putStrLn
["USAGE: genblock [cpus] filename size"
,"where 'size' is given in gigabytes"
,"and cpus defaults to the value of +RTS -N"]
checkReasonableAndGo :: Maybe String -> String -> String -> IO ()
checkReasonableAndGo mcpus file size =
case (maybe [(numCapabilities, "")] reads mcpus, reads size) of
([(cpus, "")], [(size', "")]) -> go cpus file size'
_ -> dumpUsageAndQuit
go :: Int -> String -> Int -> IO ()
go cpus filename size_in_gb = do
chan <- newChan
forM_ [1..cpus] $ \n -> forkIO {-n-} $ do
baseMT <- pureMT `fmap` getOpenSSLRand
writeList2Chan chan $
BS.toChunks $ runPut $ mersennePut baseMT
ls <- getChanContents chan -- TODO need a LBS chan
BS.writeFile filename . BS.take size . BS.fromChunks $ ls
where
size = fromIntegral size_in_gb * 1024 * 1024 * 1024
-- unroll the loop
mersennePut :: PureMT -> Put
mersennePut mt = do
let (k, mt1) = randomWord64 mt
putWord64host k
let (k, mt2) = randomWord64 mt1
putWord64host k
let (k, mt3) = randomWord64 mt2
putWord64host k
let (k, mt4) = randomWord64 mt3
putWord64host k
let (k, mt5) = randomWord64 mt4
putWord64host k
let (k, mt6) = randomWord64 mt5
putWord64host k
let (k, mt7) = randomWord64 mt6
putWord64host k
let (k, mt8) = randomWord64 mt7
putWord64host k
let (k, mt9) = randomWord64 mt8
putWord64host k
let (k, mt10) = randomWord64 mt9
putWord64host k
let (k, mt11) = randomWord64 mt10
putWord64host k
let (k, mt12) = randomWord64 mt11
putWord64host k
let (k, mt13) = randomWord64 mt12
putWord64host k
let (k, mt14) = randomWord64 mt13
putWord64host k
let (k, mt15) = randomWord64 mt14
putWord64host k
let (k, mt16) = randomWord64 mt15
putWord64host k
mersennePut mt16
------------------------------------------------------------------------
--
-- Get a decent random seed
--
getOpenSSLRand :: IO Word64
getOpenSSLRand = allocaBytes n $ \ ptr -> do
fill_RAND_bytes ptr (fromIntegral n)
peek $ castPtr ptr
where
n = sizeOf (undefined :: Word64)
foreign import ccall unsafe "openssl/rand.h RAND_bytes"
fill_RAND_bytes :: Ptr Word8 -> CInt -> IO ()