diff --git a/Changelog.md b/Changelog.md new file mode 100644 --- /dev/null +++ b/Changelog.md @@ -0,0 +1,5 @@ +# 0.0.1 + +- Add `NFData SMGen` instance +- Fix a bug. http://www.pcg-random.org/posts/bugs-in-splitmix.html + The generated numbers will be different for the same seeds! diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -14,62 +14,63 @@ is therefore quite subtle. ``` -time dieharder-input splitmix | dieharder -a -g 200 +time $(cabal-plan list-bin splitmix-dieharder) splitmix ``` The test-suite takes around half-an-hour to complete. -All tests are PASSED (occasionally WEAK). +From 30 runs, 2.49% were weak (3247 passed, 83 weak, 0 failed). In comparison, built-in [Marsenne Twister](https://en.wikipedia.org/wiki/Mersenne_Twister) test takes around 15min. ``` -time dieharder-input -a +time dieharder -a ``` ## benchmarks ``` -benchmarking list/random -time 96.77 μs (96.28 μs .. 97.35 μs) - 1.000 R² (1.000 R² .. 1.000 R²) -mean 96.77 μs (96.35 μs .. 97.63 μs) -std dev 2.001 μs (1.028 μs .. 3.796 μs) -variance introduced by outliers: 15% (moderately inflated) +benchmarking list 64/random +time 1.317 ms (1.303 ms .. 1.335 ms) + 0.998 R² (0.998 R² .. 0.999 R²) +mean 1.380 ms (1.365 ms .. 1.411 ms) +std dev 70.83 μs (37.26 μs .. 131.8 μs) +variance introduced by outliers: 39% (moderately inflated) -benchmarking list/tf-random -time 60.43 μs (60.12 μs .. 60.79 μs) - 1.000 R² (1.000 R² .. 1.000 R²) -mean 60.52 μs (60.26 μs .. 60.91 μs) -std dev 1.120 μs (780.2 ns .. 1.513 μs) -variance introduced by outliers: 14% (moderately inflated) +benchmarking list 64/tf-random +time 141.1 μs (140.4 μs .. 142.1 μs) + 0.999 R² (0.998 R² .. 1.000 R²) +mean 145.9 μs (144.6 μs .. 150.4 μs) +std dev 7.131 μs (3.461 μs .. 14.75 μs) +variance introduced by outliers: 49% (moderately inflated) -benchmarking list/splitmix -time 16.38 μs (16.29 μs .. 16.47 μs) - 1.000 R² (0.999 R² .. 1.000 R²) -mean 16.34 μs (16.25 μs .. 16.51 μs) -std dev 386.8 ns (201.5 ns .. 669.8 ns) -variance introduced by outliers: 24% (moderately inflated) +benchmarking list 64/splitmix +time 17.86 μs (17.72 μs .. 18.01 μs) + 0.999 R² (0.998 R² .. 1.000 R²) +mean 17.95 μs (17.75 μs .. 18.47 μs) +std dev 1.000 μs (444.1 ns .. 1.887 μs) +variance introduced by outliers: 64% (severely inflated) -benchmarking tree/random -time 115.4 μs (108.6 μs .. 126.0 μs) - 0.942 R² (0.910 R² .. 0.977 R²) -mean 116.0 μs (110.3 μs .. 124.2 μs) -std dev 23.47 μs (16.60 μs .. 33.95 μs) -variance introduced by outliers: 95% (severely inflated) +benchmarking tree 64/random +time 800.3 μs (793.3 μs .. 806.5 μs) + 0.999 R² (0.998 R² .. 0.999 R²) +mean 803.2 μs (798.1 μs .. 811.2 μs) +std dev 22.09 μs (14.69 μs .. 35.47 μs) +variance introduced by outliers: 18% (moderately inflated) -benchmarking tree/tf-random -time 132.4 μs (132.2 μs .. 132.5 μs) - 1.000 R² (1.000 R² .. 1.000 R²) -mean 132.5 μs (132.3 μs .. 132.6 μs) -std dev 403.3 ns (323.0 ns .. 516.5 ns) +benchmarking tree 64/tf-random +time 179.0 μs (176.6 μs .. 180.7 μs) + 0.999 R² (0.998 R² .. 0.999 R²) +mean 172.7 μs (171.3 μs .. 174.6 μs) +std dev 5.590 μs (4.919 μs .. 6.382 μs) +variance introduced by outliers: 29% (moderately inflated) -benchmarking tree/splitmix -time 59.80 μs (59.42 μs .. 60.15 μs) - 1.000 R² (0.999 R² .. 1.000 R²) -mean 59.61 μs (59.26 μs .. 60.36 μs) -std dev 1.683 μs (864.5 ns .. 3.206 μs) -variance introduced by outliers: 27% (moderately inflated) +benchmarking tree 64/splitmix +time 51.54 μs (51.01 μs .. 52.15 μs) + 0.999 R² (0.998 R² .. 0.999 R²) +mean 52.50 μs (51.93 μs .. 53.55 μs) +std dev 2.603 μs (1.659 μs .. 4.338 μs) +variance introduced by outliers: 55% (severely inflated) ``` Note: the performance can be potentially further improved when GHC gets diff --git a/bench/Bench.hs b/bench/Bench.hs --- a/bench/Bench.hs +++ b/bench/Bench.hs @@ -7,6 +7,7 @@ import qualified Data.Tree as T import qualified System.Random as R import qualified System.Random.TF as TF +import qualified System.Random.TF.Instances as TF import qualified System.Random.SplitMix as SM ------------------------------------------------------------------------------- @@ -57,6 +58,54 @@ splitMixTree w64 = genTreeN $ SM.mkSMGen w64 ------------------------------------------------------------------------------- +-- List Word64 +------------------------------------------------------------------------------- + +-- infinite list +genList64 :: (g -> (Word64, g)) -> g -> [Word64] +genList64 r = unfoldr (Just . r) + +-- truncated +genListN64 :: (g -> (Word64, g)) -> g -> [Word64] +genListN64 r = take 2048 . genList64 r + +randomList64 :: Int -> [Word64] +randomList64 = genListN64 R.random . R.mkStdGen + +tfRandomList64 :: Word64 -> [Word64] +tfRandomList64 w64 = genListN64 TF.random $ TF.seedTFGen (w64, w64, w64, w64) + +splitMixList64 :: Word64 -> [Word64] +splitMixList64 w64 = genListN64 SM.nextWord64 $ SM.mkSMGen w64 + +------------------------------------------------------------------------------- +-- Tree Word64 +------------------------------------------------------------------------------- + +genTree64 :: R.RandomGen g => (g -> (Word64, g)) -> g -> T.Tree Word64 +genTree64 r = go where + go g = case r g of + ~(i, g') -> T.Node i $ case R.split g' of + (ga, gb) -> [go ga, go gb] + +genTreeN64 :: R.RandomGen g => (g -> (Word64, g)) -> g -> T.Tree Word64 +genTreeN64 r = cutTree 9 . genTree64 r + where + cutTree :: Word64 -> T.Tree a -> T.Tree a + cutTree n (T.Node x forest) + | n <= 0 = T.Node x [] + | otherwise = T.Node x (map (cutTree (n - 1)) forest) + +randomTree64 :: Int -> T.Tree Word64 +randomTree64 = genTreeN64 R.random . R.mkStdGen + +tfRandomTree64 :: Word64 -> T.Tree Word64 +tfRandomTree64 w64 = genTreeN64 TF.random $ TF.seedTFGen (w64, w64, w64, w64) + +splitMixTree64 :: Word64 -> T.Tree Word64 +splitMixTree64 w64 = genTreeN64 SM.nextWord64 $ SM.mkSMGen w64 + +------------------------------------------------------------------------------- -- Main ------------------------------------------------------------------------------- @@ -71,5 +120,15 @@ [ bench "random" $ nf randomTree 42 , bench "tf-random" $ nf tfRandomTree 42 , bench "splitmix" $ nf splitMixTree 42 + ] + , bgroup "list 64" + [ bench "random" $ nf randomList64 42 + , bench "tf-random" $ nf tfRandomList64 42 + , bench "splitmix" $ nf splitMixList64 42 + ] + , bgroup "tree 64" + [ bench "random" $ nf randomTree64 42 + , bench "tf-random" $ nf tfRandomTree64 42 + , bench "splitmix" $ nf splitMixTree64 42 ] ] diff --git a/splitmix.cabal b/splitmix.cabal --- a/splitmix.cabal +++ b/splitmix.cabal @@ -1,6 +1,8 @@ -name: splitmix -version: 0 -synopsis: Fast Splittable PRNG +cabal-version: >= 1.10 +name: splitmix +version: 0.0.1 + +synopsis: Fast Splittable PRNG description: Pure Haskell implementation of SplitMix described in . @@ -10,22 +12,38 @@ Programming Systems Languages & Applications (OOPSLA '14). ACM, New York, NY, USA, 453-472. DOI: -license: BSD3 -license-file: LICENSE -maintainer: Oleg Grenrus -bug-reports: https://github.com/phadej/splitmix#issues + . + The paper describes a new algorithm /SplitMix/ for /splittable/ + pseudorandom number generator that is quite fast: 9 64 bit arithmetic/logical + operations per 64 bits generated. + . + /SplitMix/ is tested with two standard statistical test suites (DieHarder and + TestU01, this implementation only using the former) and it appears to be + adequate for "everyday" use, such as Monte Carlo algorithms and randomized + data structures where speed is important. + . + In particular, it __should not be used for cryptographic or security applications__, + because generated sequences of pseudorandom values are too predictable + (the mixing functions are easily inverted, and two successive outputs + suffice to reconstruct the internal state). + +license: BSD3 +license-file: LICENSE +maintainer: Oleg Grenrus +bug-reports: https://github.com/phadej/splitmix#issues category: System build-type: Simple -cabal-version: >= 1.10 tested-with: GHC==7.6.3, GHC==7.8.4, GHC==7.10.3, GHC==8.0.2, - GHC==8.2.1 + GHC==8.2.2, + GHC==8.4.2 extra-source-files: README.md + Changelog.md library default-language: Haskell2010 @@ -34,9 +52,10 @@ exposed-modules: System.Random.SplitMix build-depends: - base >=4.5 && <4.11, - time >=1.4.0.1 && <1.9, - random >=1.1 && <1.2 + base >=4.5 && <4.12, + deepseq >=1.3.0.1 && <1.5, + time >=1.4.0.1 && <1.9, + random >=1.1 && <1.2 -- dump-core -- build-depends: dump-core @@ -58,7 +77,7 @@ random, containers >=0.5 && <0.6, tf-random >=0.5 && <0.6, - criterion >=1.1.0.0 && <1.3 + criterion >=1.1.0.0 && <1.5 test-suite montecarlo-pi type: exitcode-stdio-1.0 @@ -70,16 +89,20 @@ base, splitmix -test-suite dieharder-input +test-suite splitmix-dieharder type: exitcode-stdio-1.0 default-language: Haskell2010 - ghc-options: -Wall + ghc-options: -Wall -threaded -rtsopts hs-source-dirs: tests main-is: Dieharder.hs build-depends: base, splitmix, random, - base-compat >=0.9.3 && <0.10, - bytestring >=0.10.4.0 && <0.11, - tf-random >=0.5 && <0.6 + deepseq, + async >=2.2.1 && <2.3, + vector >=0.12.0.1 && <0.13, + base-compat-batteries >=0.10.1 && <0.11, + bytestring >=0.10.4.0 && <0.11, + process >=1.1.0.2 && <1.7, + tf-random >=0.5 && <0.6 diff --git a/src/System/Random/SplitMix.hs b/src/System/Random/SplitMix.hs --- a/src/System/Random/SplitMix.hs +++ b/src/System/Random/SplitMix.hs @@ -2,6 +2,27 @@ -- /SplitMix/ is a splittable pseudorandom number generator (PRNG) that is quite fast. -- -- Guy L. Steele, Jr., Doug Lea, and Christine H. Flood. 2014. Fast splittable pseudorandom number generators. /In Proceedings of the 2014 ACM International Conference on Object Oriented Programming Systems Languages & Applications/ (OOPSLA '13). ACM, New York, NY, USA, 453-472. DOI: +-- +-- Guy L. Steele, Jr., Doug Lea, and Christine H. Flood. 2014. +-- Fast splittable pseudorandom number generators. In Proceedings +-- of the 2014 ACM International Conference on Object Oriented +-- Programming Systems Languages & Applications (OOPSLA '14). ACM, +-- New York, NY, USA, 453-472. DOI: +-- +-- +-- The paper describes a new algorithm /SplitMix/ for /splittable/ +-- pseudorandom number generator that is quite fast: 9 64 bit arithmetic/logical +-- operations per 64 bits generated. +-- +-- /SplitMix/ is tested with two standard statistical test suites (DieHarder and +-- TestU01, this implementation only using the former) and it appears to be +-- adequate for "everyday" use, such as Monte Carlo algorithms and randomized +-- data structures where speed is important. +-- +-- In particular, it __should not be used for cryptographic or security applications__, +-- because generated sequences of pseudorandom values are too predictable +-- (the mixing functions are easily inverted, and two successive outputs +-- suffice to reconstruct the internal state). {-# LANGUAGE Trustworthy #-} module System.Random.SplitMix ( SMGen, @@ -18,12 +39,13 @@ unseedSMGen, ) where -import Data.Bits (popCount, shiftL, shiftR, xor, (.|.)) -import Data.IORef (IORef, atomicModifyIORef, newIORef) +import Control.DeepSeq (NFData (..)) +import Data.Bits (popCount, shiftL, shiftR, xor, (.|.)) +import Data.IORef (IORef, atomicModifyIORef, newIORef) import Data.Time.Clock.POSIX (getPOSIXTime) -import Data.Word (Word64, Word32) -import System.CPUTime (getCPUTime, cpuTimePrecision) -import System.IO.Unsafe (unsafePerformIO) +import Data.Word (Word32, Word64) +import System.CPUTime (cpuTimePrecision, getCPUTime) +import System.IO.Unsafe (unsafePerformIO) import qualified System.Random as R @@ -38,6 +60,9 @@ } deriving Show +instance NFData SMGen where + rnf (SMGen _ _) = () + ------------------------------------------------------------------------------- -- Operations ------------------------------------------------------------------------------- @@ -48,15 +73,6 @@ where seed' = seed + gamma --- Here, and in splitSMGen, "inlining" of nextSeed is worth doing. --- I looked at the Core. -{- -nextWord64 :: SMGen -> (Word64, SMGen) -nextWord64 g0 = (mix64 x, g1) - where - (x, g1) = nextSeed g0 --} - -- | Generate an 'Int'. nextInt :: SMGen -> (Int, SMGen) nextInt g = case nextWord64 g of @@ -75,14 +91,6 @@ seed' = seed + gamma seed'' = seed' + gamma -{- -splitSMGen :: SMGen -> (SMGen, SMGen) -splitSMGen g0 = (g2, SMGen (mix64 x) (mixGamma y)) - where - (x, g1) = nextSeed g0 - (y, g2) = nextSeed g1 --} - ------------------------------------------------------------------------------- -- Algorithm ------------------------------------------------------------------------------- @@ -93,34 +101,37 @@ doubleUlp :: Double doubleUlp = 1.0 / fromIntegral (1 `shiftL` 53 :: Word64) -{- -nextSeed :: SMGen -> (Word64, SMGen) -nextSeed (SMGen seed gamma) = (seed', SMGen seed' gamma) - where - seed' = seed + gamma --} - +-- Note: in JDK implementations the mix64 and mix64variant13 +-- (which is inlined into mixGamma) are swapped. +-- +-- I have no idea if swapping them affects statistical properties. mix64 :: Word64 -> Word64 mix64 z0 = - let z1 = shiftXorMultiply 33 0xc4ceb9fe1a85ec53 z0 - z2 = shiftXorMultiply 33 0xff51afd7ed558ccd z1 + -- MurmurHash3Mixer + let z1 = shiftXorMultiply 33 0xff51afd7ed558ccd z0 + z2 = shiftXorMultiply 33 0xc4ceb9fe1a85ec53 z1 z3 = shiftXor 33 z2 in z3 +-- used only in mixGamma mix64variant13 :: Word64 -> Word64 mix64variant13 z0 = - let z1 = shiftXorMultiply 30 0xbf58476d1ce4e5b9 z0 + -- Better Bit Mixing - Improving on MurmurHash3's 64-bit Finalizer + -- http://zimbry.blogspot.fi/2011/09/better-bit-mixing-improving-on.html + let z1 = shiftXorMultiply 30 0xbf58476d1ce4e5b9 z0 -- MurmurHash3 mix constants z2 = shiftXorMultiply 27 0x94d049bb133111eb z1 z3 = shiftXor 31 z2 in z3 mixGamma :: Word64 -> Word64 mixGamma z0 = - let z1 = mix64variant13 z0 .|. 1 + let z1 = mix64variant13 z0 .|. 1 -- force to be odd n = popCount (z1 `xor` (z1 `shiftR` 1)) + -- see: http://www.pcg-random.org/posts/bugs-in-splitmix.html + -- let's trust the text of the paper, not the code. in if n >= 24 - then z1 `xor` 0xaaaaaaaaaaaaaaaa - else z1 + then z1 + else z1 `xor` 0xaaaaaaaaaaaaaaaa shiftXor :: Int -> Word64 -> Word64 shiftXor n w = w `xor` (w `shiftR` n) diff --git a/tests/Dieharder.hs b/tests/Dieharder.hs --- a/tests/Dieharder.hs +++ b/tests/Dieharder.hs @@ -1,47 +1,240 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} module Main (main) where -import Prelude +import Prelude () import Prelude.Compat -import Data.List (unfoldr) -import System.Environment (getArgs) -import System.IO (stdout) +import Control.Concurrent.QSem +import Control.DeepSeq (force) +import Data.Bits (shiftL, (.|.)) +import Data.List (isInfixOf, unfoldr) +import Data.Maybe (fromMaybe) +import Data.Word (Word64) +import Foreign.C (Errno (..), ePIPE) +import Foreign.Ptr (castPtr) +import GHC.IO.Exception (IOErrorType (..), IOException (..)) +import System.Environment (getArgs) +import System.IO (Handle, hGetContents) +import Text.Read (readMaybe) +import Text.Printf (printf) -import qualified Data.ByteString.Builder as B -import qualified System.Random.SplitMix as SM -import qualified System.Random.TF as TF -import qualified System.Random.TF.Gen as TF -import qualified System.Random.TF.Init as TF +import qualified Control.Concurrent.Async as A +import qualified Control.Exception as E +import qualified Data.ByteString as BS +import qualified Data.ByteString.Unsafe as BS (unsafePackCStringLen) +import qualified Data.Vector.Storable.Mutable as MSV +import qualified System.Process as Proc +import qualified System.Random.SplitMix as SM +import qualified System.Random.TF as TF +import qualified System.Random.TF.Gen as TF +import qualified System.Random.TF.Init as TF main :: IO () main = do args <- getArgs - case args of - ["splitmix"] -> splitmix - ["tfrandom"] -> tfrandom - ["splitmix-tree"] -> splitmixTree - _ -> return () + if null args + then return () + else do + (cmd, runs, conc, seed, test, _help) <- parseArgsIO args $ (,,,,,) + <$> arg + <*> optDef "-n" 1 + <*> optDef "-j" 1 + <*> opt "-s" + <*> opt "-d" + <*> flag "-h" -splitmix :: IO () -splitmix = SM.initSMGen >>= B.hPutBuilder stdout . go - where - go :: SM.SMGen -> B.Builder - go g = case SM.nextWord64 g of - ~(w64, g') -> B.word64LE w64 `mappend` go g' + case cmd of + "splitmix" -> do + g <- maybe SM.initSMGen (return . SM.mkSMGen) seed + run test runs conc SM.splitSMGen SM.nextWord64 g + "tfrandom" -> do + g <- TF.initTFGen + run test runs conc TF.split tfNext64 g + _ -> return () -splitmixTree :: IO () -splitmixTree = SM.initSMGen >>= B.hPutBuilder stdout . go - where - go :: SM.SMGen -> B.Builder - go g = case SM.splitSMGen g of - ~(ga, gb) -> builder 8 ga `mappend` go gb +tfNext64 :: TF.TFGen -> (Word64, TF.TFGen) +tfNext64 g = + let (w, g') = TF.next g + (w', g'') = TF.next g' + in (fromIntegral w `shiftL` 32 .|. fromIntegral w', g'') - builder :: Int -> SM.SMGen -> B.Builder - builder n = mconcat . take n . map B.word64LE . unfoldr (Just . SM.nextWord64) +------------------------------------------------------------------------------- +-- Dieharder +------------------------------------------------------------------------------- -tfrandom :: IO () -tfrandom = TF.initTFGen >>= B.hPutBuilder stdout . go +run :: Maybe Int + -> Int + -> Int + -> (g -> (g, g)) + -> (g -> (Word64, g)) + -> g + -> IO () +run test runs conc split word gen = do + qsem <- newQSem conc + + rs <- A.forConcurrently (take runs $ unfoldr (Just . split) gen) $ \g -> + E.bracket_ (waitQSem qsem) (signalQSem qsem) $ + dieharder test (generate word g) + + case mconcat rs of + Result p w f -> do + let total = fromIntegral (p + w + f) :: Double + printf "PASSED %4d %6.02f%%\n" p (fromIntegral p / total * 100) + printf "WEAK %4d %6.02f%%\n" w (fromIntegral w / total * 100) + printf "FAILED %4d %6.02f%%\n" f (fromIntegral f / total * 100) +{-# INLINE run #-} + +dieharder :: Maybe Int -> (Handle -> IO ()) -> IO Result +dieharder test gen = do + let proc = Proc.proc "dieharder" $ ["-g", "200"] ++ maybe ["-a"] (\t -> ["-d", show t]) test + (Just hin, Just hout, _, ph) <- Proc.createProcess proc + { Proc.std_in = Proc.CreatePipe + , Proc.std_out = Proc.CreatePipe + } + + out <- hGetContents hout + waitOut <- A.async $ E.evaluate $ force out + + E.catch (gen hin) $ \e -> case e of + IOError { ioe_type = ResourceVanished , ioe_errno = Just ioe } + | Errno ioe == ePIPE -> return () + _ -> E.throwIO e + + res <- A.wait waitOut + _ <- Proc.waitForProcess ph + + return $ parseOutput res +{-# INLINE dieharder #-} + +parseOutput :: String -> Result +parseOutput = foldMap parseLine . lines where + parseLine l + | any (`isInfixOf` l) doNotUse = mempty + | "PASSED" `isInfixOf` l = Result 1 0 0 + | "WEAK" `isInfixOf` l = Result 0 1 0 + | "FAILED" `isInfixOf` l = Result 0 1 0 + | otherwise = mempty + + doNotUse = ["diehard_opso", "diehard_oqso", "diehard_dna", "diehard_weak"] + +------------------------------------------------------------------------------- +-- Results +------------------------------------------------------------------------------- + +data Result = Result + { _passed :: Int + , _weak :: Int + , _failed :: Int + } + deriving Show + +instance Semigroup Result where + Result p w f <> Result p' w' f' = Result (p + p') (w + w') (f + f') + +instance Monoid Result where + mempty = Result 0 0 0 + mappend = (<>) + +------------------------------------------------------------------------------- +-- Writer +------------------------------------------------------------------------------- + +size :: Int +size = 512 + +generate :: forall g. (g -> (Word64, g)) -> g -> Handle -> IO () +generate word gen0 h = do + vec <- MSV.new size + go gen0 vec where - go :: TF.TFGen -> B.Builder - go g = case TF.next g of - ~(w32, g') -> B.word32LE w32 `mappend` go g' + go :: g -> MSV.IOVector Word64 -> IO () + go gen vec = do + gen' <- write gen vec 0 + MSV.unsafeWith vec $ \ptr -> do + bs <- BS.unsafePackCStringLen (castPtr ptr, size * 8) + BS.hPutStr h bs + go gen' vec + + write :: g -> MSV.IOVector Word64 -> Int -> IO g + write !gen !vec !i = do + let (w64, gen') = word gen + MSV.unsafeWrite vec i w64 + if i < size + then write gen' vec (i + 1) + else return gen' +{-# INLINE generate #-} + +------------------------------------------------------------------------------- +-- Do it yourself command line parsing +------------------------------------------------------------------------------- + +-- | 'Parser' is not an 'Alternative', only a *commutative* 'Applicative'. +-- +-- Useful for quick cli parsers, like parametrising tests. +data Parser a where + Pure :: a -> Parser a + Ap :: Arg b -> Parser (b -> a) -> Parser a + +instance Functor Parser where + fmap f (Pure a) = Pure (f a) + fmap f (Ap x y) = Ap x (fmap (f .) y) + +instance Applicative Parser where + pure = Pure + + Pure f <*> z = fmap f z + Ap x y <*> z = Ap x (flip <$> y <*> z) + +data Arg a where + Flag :: String -> Arg Bool + Opt :: String -> (String -> Maybe a) -> Arg (Maybe a) + Arg :: Arg String + +arg :: Parser String +arg = Ap Arg (Pure id) + +flag :: String -> Parser Bool +flag n = Ap (Flag n) (Pure id) + +opt :: Read a => String -> Parser (Maybe a) +opt n = Ap (Opt n readMaybe) (Pure id) + +optDef :: Read a => String -> a -> Parser a +optDef n d = Ap (Opt n readMaybe) (Pure (fromMaybe d)) + +parseArgsIO :: [String] -> Parser a -> IO a +parseArgsIO args p = either fail pure (parseArgs args p) + +parseArgs :: [String] -> Parser a -> Either String a +parseArgs [] p = parserToEither p +parseArgs (x : xs) p = do + (xs', p') <- singleArg p x xs + parseArgs xs' p' + +singleArg :: Parser a -> String -> [String] -> Either String ([String], Parser a) +singleArg (Pure _) x _ = Left $ "Extra argument " ++ x +singleArg (Ap Arg p) x xs + | null x || head x /= '-' = Right (xs, fmap ($ x) p) + | otherwise = fmap2 (Ap Arg) (singleArg p x xs) +singleArg (Ap f@(Flag n) p) x xs + | x == n = Right (xs, fmap ($ True) p) + | otherwise = fmap2 (Ap f) (singleArg p x xs) +singleArg (Ap o@(Opt n r) p) x xs + | x == n = case xs of + [] -> Left $ "Expected an argument for " ++ n + (x' : xs') -> case r x' of + Nothing -> Left $ "Cannot read an argument of " ++ n ++ ": " ++ x' + Just y -> Right (xs', fmap ($ Just y) p) + | otherwise = fmap2 (Ap o) (singleArg p x xs) + +fmap2 :: (Functor f, Functor g) => (a -> b) -> f (g a) -> f (g b) +fmap2 = fmap . fmap + +-- | Convert parser to 'Right' if there are only defaultable pieces left. +parserToEither :: Parser a -> Either String a +parserToEither (Pure x) = pure x +parserToEither (Ap (Flag _) p) = parserToEither $ fmap ($ False) p +parserToEither (Ap (Opt _ _) p) = parserToEither $ fmap ($ Nothing) p +parserToEither (Ap Arg _) = Left "argument required"