1{-# LANGUAGE BangPatterns        #-}
2{-# LANGUAGE GADTs               #-}
3{-# LANGUAGE ScopedTypeVariables #-}
4module Main (main) where
5
6import Prelude ()
7import Prelude.Compat
8
9import Control.Concurrent.QSem
10import Control.DeepSeq         (force)
11import Control.Monad           (when)
12import Data.Bits               (shiftL, (.|.))
13import Data.Char               (isSpace)
14import Data.List               (isInfixOf, unfoldr)
15import Data.Maybe              (fromMaybe)
16import Data.Word               (Word64)
17import Foreign.C               (Errno (..), ePIPE)
18import Foreign.Ptr             (castPtr)
19import GHC.IO.Exception        (IOErrorType (..), IOException (..))
20import System.Environment      (getArgs)
21import System.IO               (Handle, hGetContents, stdout)
22import Text.Printf             (printf)
23
24import qualified Control.Concurrent.Async     as A
25import qualified Control.Exception            as E
26import qualified Data.ByteString              as BS
27import qualified Data.ByteString.Unsafe       as BS (unsafePackCStringLen)
28import qualified Data.Vector.Storable.Mutable as MSV
29import qualified System.Process               as Proc
30import qualified System.Random.SplitMix       as SM
31import qualified System.Random.SplitMix32     as SM32
32import qualified System.Random.TF             as TF
33import qualified System.Random.TF.Gen         as TF
34import qualified System.Random.TF.Init        as TF
35
36main :: IO ()
37main = do
38    args <- getArgs
39    if null args
40    then return ()
41    else do
42        (cmd, runs, conc, seed, test, raw, _help) <- parseArgsIO args $ (,,,,,,)
43            <$> arg
44            <*> optDef "-n" 1
45            <*> optDef "-j" 1
46            <*> opt "-s"
47            <*> opt "-d"
48            <*> flag "-r"
49            <*> flag "-h"
50
51        let run :: RunType g
52            run | raw       = runRaw
53                | otherwise = runManaged
54
55        case cmd of
56              "splitmix"      -> do
57                  g <- maybe SM.initSMGen (return . SM.mkSMGen) seed
58                  run test runs conc SM.splitSMGen SM.nextWord64 g
59              "splitmix32"      -> do
60                  g <- maybe SM32.initSMGen (return . SM32.mkSMGen) (fmap fromIntegral seed)
61                  run test runs conc SM32.splitSMGen SM32.nextWord64 g
62              "tfrandom"      -> do
63                  g <- TF.initTFGen
64                  run test runs conc TF.split tfNext64 g
65              _               -> return ()
66
67tfNext64 :: TF.TFGen -> (Word64, TF.TFGen)
68tfNext64 g =
69    let (w, g')   = TF.next g
70        (w', g'') = TF.next g'
71    in (fromIntegral w `shiftL` 32 .|. fromIntegral w', g'')
72
73-------------------------------------------------------------------------------
74-- Dieharder
75-------------------------------------------------------------------------------
76
77type RunType g =
78       Maybe Int
79    -> Int
80    -> Int
81    -> (g -> (g, g))
82    -> (g -> (Word64, g))
83    -> g
84    -> IO ()
85
86runRaw :: RunType g
87runRaw _test _runs _conc split word gen =
88    generate word split gen stdout
89
90runManaged :: RunType g
91runManaged test runs conc split word gen = do
92    qsem <- newQSem conc
93
94    rs <- A.forConcurrently (take runs $ unfoldr (Just . split) gen) $ \g ->
95        E.bracket_ (waitQSem qsem) (signalQSem qsem) $
96            dieharder test (generate word split g)
97
98    case mconcat rs of
99        Result p w f -> do
100            let total = fromIntegral (p + w + f) :: Double
101            printf "PASSED %4d %6.02f%%\n" p (fromIntegral p / total * 100)
102            printf "WEAK   %4d %6.02f%%\n" w (fromIntegral w / total * 100)
103            printf "FAILED %4d %6.02f%%\n" f (fromIntegral f / total * 100)
104{-# INLINE runManaged #-}
105
106dieharder :: Maybe Int -> (Handle -> IO ()) -> IO Result
107dieharder test gen = do
108    let proc = Proc.proc "dieharder" $ ["-g", "200"] ++ maybe ["-a"] (\t -> ["-d", show t]) test
109    (Just hin, Just hout, _, ph) <- Proc.createProcess proc
110        { Proc.std_in  = Proc.CreatePipe
111        , Proc.std_out = Proc.CreatePipe
112        }
113
114    out <- hGetContents hout
115    waitOut <- A.async $ E.evaluate $ force out
116
117    E.catch (gen hin) $ \e -> case e of
118        IOError { ioe_type = ResourceVanished , ioe_errno = Just ioe }
119            | Errno ioe == ePIPE -> return ()
120        _ -> E.throwIO e
121
122    res <- A.wait waitOut
123    _ <- Proc.waitForProcess ph
124
125    return $ parseOutput res
126{-# INLINE dieharder #-}
127
128parseOutput :: String -> Result
129parseOutput = foldMap parseLine . lines where
130    parseLine l
131        | any (`isInfixOf` l) doNotUse = mempty
132        | "PASSED" `isInfixOf` l = Result 1 0 0
133        | "WEAK"   `isInfixOf` l = Result 0 1 0
134        | "FAILED" `isInfixOf` l = Result 0 1 0
135        | otherwise = mempty
136
137    doNotUse = ["diehard_opso", "diehard_oqso", "diehard_dna", "diehard_weak"]
138
139-------------------------------------------------------------------------------
140-- Results
141-------------------------------------------------------------------------------
142
143data Result = Result
144    { _passed :: Int
145    , _weak   :: Int
146    , _failed :: Int
147    }
148  deriving Show
149
150instance Semigroup Result where
151    Result p w f <> Result p' w' f' = Result (p + p') (w +  w') (f + f')
152
153instance Monoid Result where
154    mempty = Result 0 0 0
155    mappend = (<>)
156
157-------------------------------------------------------------------------------
158-- Writer
159-------------------------------------------------------------------------------
160
161size :: Int
162size = 512
163
164generate
165    :: forall g. (g -> (Word64, g))
166    -> (g -> (g, g))
167    -> g -> Handle -> IO ()
168generate word split gen0 h = do
169    vec <- MSV.new size
170    go gen0 vec
171  where
172    go :: g -> MSV.IOVector Word64 -> IO ()
173    go gen vec = do
174        let (g1, g2) = split gen
175        write g1 vec 0
176        MSV.unsafeWith vec $ \ptr -> do
177            bs <- BS.unsafePackCStringLen (castPtr ptr, size * 8)
178            BS.hPutStr h bs
179        go g2 vec
180
181    write :: g -> MSV.IOVector Word64 -> Int -> IO ()
182    write !gen !vec !i = do
183        let (w64, gen') = word gen
184        MSV.unsafeWrite vec i w64
185        when (i < size) $
186            write gen' vec (i + 1)
187{-# INLINE generate #-}
188
189-------------------------------------------------------------------------------
190-- readMaybe
191-------------------------------------------------------------------------------
192
193readEither :: Read a => String -> Either String a
194readEither s =
195  case [ x | (x,rest) <- reads s, all isSpace rest ] of
196    [x] -> Right x
197    []  -> Left "Prelude.read: no parse"
198    _   -> Left "Prelude.read: ambiguous parse"
199
200readMaybe :: Read a => String -> Maybe a
201readMaybe s = case readEither s of
202                Left _  -> Nothing
203                Right a -> Just a
204-------------------------------------------------------------------------------
205-- Do it yourself command line parsing
206-------------------------------------------------------------------------------
207
208-- | 'Parser' is not an 'Alternative', only a *commutative* 'Applicative'.
209--
210-- Useful for quick cli parsers, like parametrising tests.
211data Parser a where
212    Pure :: a -> Parser a
213    Ap :: Arg b -> Parser (b -> a) -> Parser a
214
215instance Functor Parser where
216    fmap f (Pure a) = Pure (f a)
217    fmap f (Ap x y) = Ap x (fmap (f .) y)
218
219instance  Applicative Parser where
220    pure = Pure
221
222    Pure f <*> z = fmap f z
223    Ap x y <*> z = Ap x (flip <$> y <*> z)
224
225data Arg a where
226    Flag :: String -> Arg Bool
227    Opt  :: String -> (String -> Maybe a) -> Arg (Maybe a)
228    Arg  :: Arg String
229
230arg :: Parser String
231arg = Ap Arg (Pure id)
232
233flag :: String -> Parser Bool
234flag n = Ap (Flag n) (Pure id)
235
236opt :: Read a => String -> Parser (Maybe a)
237opt n = Ap (Opt n readMaybe) (Pure id)
238
239optDef :: Read a => String -> a -> Parser a
240optDef n d = Ap (Opt n readMaybe) (Pure (fromMaybe d))
241
242parseArgsIO :: [String] -> Parser a -> IO a
243parseArgsIO args p = either fail pure (parseArgs args p)
244
245parseArgs :: [String] -> Parser a -> Either String a
246parseArgs []       p = parserToEither p
247parseArgs (x : xs) p = do
248    (xs', p') <- singleArg p x xs
249    parseArgs xs' p'
250
251singleArg :: Parser a -> String -> [String] -> Either String ([String], Parser a)
252singleArg (Pure _)           x _  = Left $ "Extra argument " ++ x
253singleArg (Ap Arg p)         x xs
254    | null x || head x /= '-'     = Right (xs, fmap ($ x) p)
255    | otherwise                   = fmap2 (Ap Arg) (singleArg p x xs)
256singleArg (Ap f@(Flag n) p)  x xs
257    | x == n                      = Right (xs, fmap ($ True) p)
258    | otherwise                   = fmap2 (Ap f) (singleArg p x xs)
259singleArg (Ap o@(Opt n r) p) x xs
260    | x == n                      = case xs of
261        [] -> Left $ "Expected an argument for " ++ n
262        (x' : xs') -> case r x' of
263            Nothing -> Left $ "Cannot read an argument of " ++ n ++ ": " ++ x'
264            Just y  -> Right (xs', fmap ($ Just y) p)
265    | otherwise                   = fmap2 (Ap o) (singleArg p x xs)
266
267fmap2 :: (Functor f, Functor g) => (a -> b) -> f (g a) -> f (g b)
268fmap2 = fmap . fmap
269
270-- | Convert parser to 'Right' if there are only defaultable pieces left.
271parserToEither :: Parser a -> Either String a
272parserToEither (Pure x)         = pure x
273parserToEither (Ap (Flag _) p)  = parserToEither $ fmap ($ False) p
274parserToEither (Ap (Opt _ _) p) = parserToEither $ fmap ($ Nothing) p
275parserToEither (Ap Arg _)       = Left "argument required"
276