{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}

-- | @futhark dataset@
module Futhark.CLI.Dataset (main) where

import Control.Monad
import Control.Monad.ST
import qualified Data.Binary as Bin
import qualified Data.ByteString.Lazy.Char8 as BS
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import Data.Vector.Generic (freeze)
import qualified Data.Vector.Storable as SVec
import qualified Data.Vector.Storable.Mutable as USVec
import Data.Word
import Futhark.Test.Values
import Futhark.Util.Options
import Language.Futhark.Parser
import Language.Futhark.Pretty ()
import Language.Futhark.Prop (UncheckedTypeExp, namesToPrimTypes)
import Language.Futhark.Syntax hiding
  ( FloatValue (..),
    IntValue (..),
    PrimValue (..),
    Value,
    ValueType,
  )
import System.Exit
import System.IO
import System.Random.PCG (Variate, initialize, uniformR)

-- | Run @futhark dataset@.
main :: String -> [String] -> IO ()
main :: String -> [String] -> IO ()
main = DataOptions
-> [FunOptDescr DataOptions]
-> String
-> ([String] -> DataOptions -> Maybe (IO ()))
-> String
-> [String]
-> IO ()
forall cfg.
cfg
-> [FunOptDescr cfg]
-> String
-> ([String] -> cfg -> Maybe (IO ()))
-> String
-> [String]
-> IO ()
mainWithOptions DataOptions
initialDataOptions [FunOptDescr DataOptions]
commandLineOptions String
"options..." [String] -> DataOptions -> Maybe (IO ())
forall {a}. [a] -> DataOptions -> Maybe (IO ())
f
  where
    f :: [a] -> DataOptions -> Maybe (IO ())
f [] DataOptions
config
      | [Word64 -> IO ()] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Word64 -> IO ()] -> Bool) -> [Word64 -> IO ()] -> Bool
forall a b. (a -> b) -> a -> b
$ DataOptions -> [Word64 -> IO ()]
optOrders DataOptions
config = IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just (IO () -> Maybe (IO ())) -> IO () -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$ do
        Maybe [Value]
maybe_vs <- ByteString -> Maybe [Value]
readValues (ByteString -> Maybe [Value])
-> IO ByteString -> IO (Maybe [Value])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ByteString
BS.getContents
        case Maybe [Value]
maybe_vs of
          Maybe [Value]
Nothing -> do
            Handle -> String -> IO ()
hPutStrLn Handle
stderr String
"Malformed data on standard input."
            IO ()
forall a. IO a
exitFailure
          Just [Value]
vs ->
            case DataOptions -> OutputFormat
format DataOptions
config of
              OutputFormat
Text -> (Value -> IO ()) -> [Value] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (String -> IO ()
putStrLn (String -> IO ()) -> (Value -> String) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> String
forall a. Pretty a => a -> String
pretty) [Value]
vs
              OutputFormat
Binary -> (Value -> IO ()) -> [Value] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (ByteString -> IO ()
BS.putStr (ByteString -> IO ()) -> (Value -> ByteString) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ByteString
forall a. Binary a => a -> ByteString
Bin.encode) [Value]
vs
              OutputFormat
Type -> (Value -> IO ()) -> [Value] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (String -> IO ()
putStrLn (String -> IO ()) -> (Value -> String) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueType -> String
forall a. Pretty a => a -> String
pretty (ValueType -> String) -> (Value -> ValueType) -> Value -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ValueType
valueType) [Value]
vs
      | Bool
otherwise =
        IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just (IO () -> Maybe (IO ())) -> IO () -> Maybe (IO ())
forall a b. (a -> b) -> a -> b
$
          ((Word64 -> IO ()) -> Word64 -> IO ())
-> [Word64 -> IO ()] -> [Word64] -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_
            (Word64 -> IO ()) -> Word64 -> IO ()
forall a b. (a -> b) -> a -> b
($)
            (DataOptions -> [Word64 -> IO ()]
optOrders DataOptions
config)
            [Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DataOptions -> Int
optSeed DataOptions
config) ..]
    f [a]
_ DataOptions
_ =
      Maybe (IO ())
forall a. Maybe a
Nothing

data OutputFormat
  = Text
  | Binary
  | Type
  deriving (OutputFormat -> OutputFormat -> Bool
(OutputFormat -> OutputFormat -> Bool)
-> (OutputFormat -> OutputFormat -> Bool) -> Eq OutputFormat
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: OutputFormat -> OutputFormat -> Bool
$c/= :: OutputFormat -> OutputFormat -> Bool
== :: OutputFormat -> OutputFormat -> Bool
$c== :: OutputFormat -> OutputFormat -> Bool
Eq, Eq OutputFormat
Eq OutputFormat
-> (OutputFormat -> OutputFormat -> Ordering)
-> (OutputFormat -> OutputFormat -> Bool)
-> (OutputFormat -> OutputFormat -> Bool)
-> (OutputFormat -> OutputFormat -> Bool)
-> (OutputFormat -> OutputFormat -> Bool)
-> (OutputFormat -> OutputFormat -> OutputFormat)
-> (OutputFormat -> OutputFormat -> OutputFormat)
-> Ord OutputFormat
OutputFormat -> OutputFormat -> Bool
OutputFormat -> OutputFormat -> Ordering
OutputFormat -> OutputFormat -> OutputFormat
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: OutputFormat -> OutputFormat -> OutputFormat
$cmin :: OutputFormat -> OutputFormat -> OutputFormat
max :: OutputFormat -> OutputFormat -> OutputFormat
$cmax :: OutputFormat -> OutputFormat -> OutputFormat
>= :: OutputFormat -> OutputFormat -> Bool
$c>= :: OutputFormat -> OutputFormat -> Bool
> :: OutputFormat -> OutputFormat -> Bool
$c> :: OutputFormat -> OutputFormat -> Bool
<= :: OutputFormat -> OutputFormat -> Bool
$c<= :: OutputFormat -> OutputFormat -> Bool
< :: OutputFormat -> OutputFormat -> Bool
$c< :: OutputFormat -> OutputFormat -> Bool
compare :: OutputFormat -> OutputFormat -> Ordering
$ccompare :: OutputFormat -> OutputFormat -> Ordering
Ord, Int -> OutputFormat -> ShowS
[OutputFormat] -> ShowS
OutputFormat -> String
(Int -> OutputFormat -> ShowS)
-> (OutputFormat -> String)
-> ([OutputFormat] -> ShowS)
-> Show OutputFormat
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [OutputFormat] -> ShowS
$cshowList :: [OutputFormat] -> ShowS
show :: OutputFormat -> String
$cshow :: OutputFormat -> String
showsPrec :: Int -> OutputFormat -> ShowS
$cshowsPrec :: Int -> OutputFormat -> ShowS
Show)

data DataOptions = DataOptions
  { DataOptions -> Int
optSeed :: Int,
    DataOptions -> RandomConfiguration
optRange :: RandomConfiguration,
    DataOptions -> [Word64 -> IO ()]
optOrders :: [Word64 -> IO ()],
    DataOptions -> OutputFormat
format :: OutputFormat
  }

initialDataOptions :: DataOptions
initialDataOptions :: DataOptions
initialDataOptions = Int
-> RandomConfiguration
-> [Word64 -> IO ()]
-> OutputFormat
-> DataOptions
DataOptions Int
1 RandomConfiguration
initialRandomConfiguration [] OutputFormat
Text

commandLineOptions :: [FunOptDescr DataOptions]
commandLineOptions :: [FunOptDescr DataOptions]
commandLineOptions =
  [ String
-> [String]
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
-> String
-> FunOptDescr DataOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"s"
      [String
"seed"]
      ( (String -> Either (IO ()) (DataOptions -> DataOptions))
-> String -> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          ( \String
n ->
              case ReadS Int
forall a. Read a => ReadS a
reads String
n of
                [(Int
n', String
"")] ->
                  (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. b -> Either a b
Right ((DataOptions -> DataOptions)
 -> Either (IO ()) (DataOptions -> DataOptions))
-> (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ \DataOptions
config -> DataOptions
config {optSeed :: Int
optSeed = Int
n'}
                [(Int, String)]
_ ->
                  IO () -> Either (IO ()) (DataOptions -> DataOptions)
forall a b. a -> Either a b
Left (IO () -> Either (IO ()) (DataOptions -> DataOptions))
-> IO () -> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ do
                    Handle -> String -> IO ()
hPutStrLn Handle
stderr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"'" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"' is not an integer."
                    IO ()
forall a. IO a
exitFailure
          )
          String
"SEED"
      )
      String
"The seed to use when initialising the RNG.",
    String
-> [String]
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
-> String
-> FunOptDescr DataOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"g"
      [String
"generate"]
      ( (String -> Either (IO ()) (DataOptions -> DataOptions))
-> String -> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          ( \String
t ->
              case String
-> Either
     String (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
tryMakeGenerator String
t of
                Right RandomConfiguration -> OutputFormat -> Word64 -> IO ()
g ->
                  (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. b -> Either a b
Right ((DataOptions -> DataOptions)
 -> Either (IO ()) (DataOptions -> DataOptions))
-> (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ \DataOptions
config ->
                    DataOptions
config
                      { optOrders :: [Word64 -> IO ()]
optOrders =
                          DataOptions -> [Word64 -> IO ()]
optOrders DataOptions
config
                            [Word64 -> IO ()] -> [Word64 -> IO ()] -> [Word64 -> IO ()]
forall a. [a] -> [a] -> [a]
++ [RandomConfiguration -> OutputFormat -> Word64 -> IO ()
g (DataOptions -> RandomConfiguration
optRange DataOptions
config) (DataOptions -> OutputFormat
format DataOptions
config)]
                      }
                Left String
err ->
                  IO () -> Either (IO ()) (DataOptions -> DataOptions)
forall a b. a -> Either a b
Left (IO () -> Either (IO ()) (DataOptions -> DataOptions))
-> IO () -> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ do
                    Handle -> String -> IO ()
hPutStrLn Handle
stderr String
err
                    IO ()
forall a. IO a
exitFailure
          )
          String
"TYPE"
      )
      String
"Generate a random value of this type.",
    String
-> [String]
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
-> String
-> FunOptDescr DataOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"text"]
      (Either (IO ()) (DataOptions -> DataOptions)
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a. a -> ArgDescr a
NoArg (Either (IO ()) (DataOptions -> DataOptions)
 -> ArgDescr (Either (IO ()) (DataOptions -> DataOptions)))
-> Either (IO ()) (DataOptions -> DataOptions)
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a b. (a -> b) -> a -> b
$ (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. b -> Either a b
Right ((DataOptions -> DataOptions)
 -> Either (IO ()) (DataOptions -> DataOptions))
-> (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ \DataOptions
opts -> DataOptions
opts {format :: OutputFormat
format = OutputFormat
Text})
      String
"Output data in text format (must precede --generate).",
    String
-> [String]
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
-> String
-> FunOptDescr DataOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"b"
      [String
"binary"]
      (Either (IO ()) (DataOptions -> DataOptions)
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a. a -> ArgDescr a
NoArg (Either (IO ()) (DataOptions -> DataOptions)
 -> ArgDescr (Either (IO ()) (DataOptions -> DataOptions)))
-> Either (IO ()) (DataOptions -> DataOptions)
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a b. (a -> b) -> a -> b
$ (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. b -> Either a b
Right ((DataOptions -> DataOptions)
 -> Either (IO ()) (DataOptions -> DataOptions))
-> (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ \DataOptions
opts -> DataOptions
opts {format :: OutputFormat
format = OutputFormat
Binary})
      String
"Output data in binary Futhark format (must precede --generate).",
    String
-> [String]
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
-> String
-> FunOptDescr DataOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"t"
      [String
"type"]
      (Either (IO ()) (DataOptions -> DataOptions)
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a. a -> ArgDescr a
NoArg (Either (IO ()) (DataOptions -> DataOptions)
 -> ArgDescr (Either (IO ()) (DataOptions -> DataOptions)))
-> Either (IO ()) (DataOptions -> DataOptions)
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a b. (a -> b) -> a -> b
$ (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. b -> Either a b
Right ((DataOptions -> DataOptions)
 -> Either (IO ()) (DataOptions -> DataOptions))
-> (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ \DataOptions
opts -> DataOptions
opts {format :: OutputFormat
format = OutputFormat
Type})
      String
"Output the type (textually) rather than the value (must precede --generate).",
    String
-> (Range Int8 -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"i8" Range Int8 -> RandomConfiguration -> RandomConfiguration
seti8Range,
    String
-> (Range Int16 -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"i16" Range Int16 -> RandomConfiguration -> RandomConfiguration
seti16Range,
    String
-> (Range Int32 -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"i32" Range Int32 -> RandomConfiguration -> RandomConfiguration
seti32Range,
    String
-> (Range Int64 -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"i64" Range Int64 -> RandomConfiguration -> RandomConfiguration
seti64Range,
    String
-> (Range Word8 -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"u8" Range Word8 -> RandomConfiguration -> RandomConfiguration
setu8Range,
    String
-> (Range Word16 -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"u16" Range Word16 -> RandomConfiguration -> RandomConfiguration
setu16Range,
    String
-> (Range Word32 -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"u32" Range Word32 -> RandomConfiguration -> RandomConfiguration
setu32Range,
    String
-> (Range Word64 -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"u64" Range Word64 -> RandomConfiguration -> RandomConfiguration
setu64Range,
    String
-> (Range Float -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"f32" Range Float -> RandomConfiguration -> RandomConfiguration
setf32Range,
    String
-> (Range Double -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"f64" Range Double -> RandomConfiguration -> RandomConfiguration
setf64Range
  ]

setRangeOption ::
  Read a =>
  String ->
  (Range a -> RandomConfiguration -> RandomConfiguration) ->
  FunOptDescr DataOptions
setRangeOption :: forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
tname Range a -> RandomConfiguration -> RandomConfiguration
set =
  String
-> [String]
-> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
-> String
-> FunOptDescr DataOptions
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
    String
""
    [String
name]
    ( (String -> Either (IO ()) (DataOptions -> DataOptions))
-> String -> ArgDescr (Either (IO ()) (DataOptions -> DataOptions))
forall a. (String -> a) -> String -> ArgDescr a
ReqArg
        ( \String
b ->
            let (String
lower, String
rest) = (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
':') String
b
                upper :: String
upper = Int -> ShowS
forall a. Int -> [a] -> [a]
drop Int
1 String
rest
             in case (ReadS a
forall a. Read a => ReadS a
reads String
lower, ReadS a
forall a. Read a => ReadS a
reads String
upper) of
                  ([(a
lower', String
"")], [(a
upper', String
"")]) ->
                    (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. b -> Either a b
Right ((DataOptions -> DataOptions)
 -> Either (IO ()) (DataOptions -> DataOptions))
-> (DataOptions -> DataOptions)
-> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ \DataOptions
config ->
                      DataOptions
config {optRange :: RandomConfiguration
optRange = Range a -> RandomConfiguration -> RandomConfiguration
set (a
lower', a
upper') (RandomConfiguration -> RandomConfiguration)
-> RandomConfiguration -> RandomConfiguration
forall a b. (a -> b) -> a -> b
$ DataOptions -> RandomConfiguration
optRange DataOptions
config}
                  ([(a, String)], [(a, String)])
_ ->
                    IO () -> Either (IO ()) (DataOptions -> DataOptions)
forall a b. a -> Either a b
Left (IO () -> Either (IO ()) (DataOptions -> DataOptions))
-> IO () -> Either (IO ()) (DataOptions -> DataOptions)
forall a b. (a -> b) -> a -> b
$ do
                      Handle -> String -> IO ()
hPutStrLn Handle
stderr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Invalid bounds for " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
tname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
": " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
b
                      IO ()
forall a. IO a
exitFailure
        )
        String
"MIN:MAX"
    )
    (String -> FunOptDescr DataOptions)
-> String -> FunOptDescr DataOptions
forall a b. (a -> b) -> a -> b
$ String
"Range of " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
tname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" values."
  where
    name :: String
name = String
tname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"-bounds"

tryMakeGenerator ::
  String ->
  Either String (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
tryMakeGenerator :: String
-> Either
     String (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
tryMakeGenerator String
t
  | Just [Value]
vs <- ByteString -> Maybe [Value]
readValues (ByteString -> Maybe [Value]) -> ByteString -> Maybe [Value]
forall a b. (a -> b) -> a -> b
$ String -> ByteString
BS.pack String
t =
    (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
-> Either
     String (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return ((RandomConfiguration -> OutputFormat -> Word64 -> IO ())
 -> Either
      String (RandomConfiguration -> OutputFormat -> Word64 -> IO ()))
-> (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
-> Either
     String (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
forall a b. (a -> b) -> a -> b
$ \RandomConfiguration
_ OutputFormat
fmt Word64
_ -> (Value -> IO ()) -> [Value] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (OutputFormat -> Value -> IO ()
outValue OutputFormat
fmt) [Value]
vs
  | Bool
otherwise = do
    ValueType
t' <- UncheckedTypeExp -> Either String ValueType
toValueType (UncheckedTypeExp -> Either String ValueType)
-> Either String UncheckedTypeExp -> Either String ValueType
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (ParseError -> Either String UncheckedTypeExp)
-> (UncheckedTypeExp -> Either String UncheckedTypeExp)
-> Either ParseError UncheckedTypeExp
-> Either String UncheckedTypeExp
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Either String UncheckedTypeExp
forall a b. a -> Either a b
Left (String -> Either String UncheckedTypeExp)
-> (ParseError -> String)
-> ParseError
-> Either String UncheckedTypeExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParseError -> String
forall a. Show a => a -> String
show) UncheckedTypeExp -> Either String UncheckedTypeExp
forall a b. b -> Either a b
Right (String -> Text -> Either ParseError UncheckedTypeExp
parseType String
name (String -> Text
T.pack String
t))
    (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
-> Either
     String (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return ((RandomConfiguration -> OutputFormat -> Word64 -> IO ())
 -> Either
      String (RandomConfiguration -> OutputFormat -> Word64 -> IO ()))
-> (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
-> Either
     String (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
forall a b. (a -> b) -> a -> b
$ \RandomConfiguration
conf OutputFormat
fmt Word64
seed -> do
      let v :: Value
v = RandomConfiguration -> ValueType -> Word64 -> Value
randomValue RandomConfiguration
conf ValueType
t' Word64
seed
      OutputFormat -> Value -> IO ()
outValue OutputFormat
fmt Value
v
  where
    name :: String
name = String
"option " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
t
    outValue :: OutputFormat -> Value -> IO ()
outValue OutputFormat
Text = String -> IO ()
putStrLn (String -> IO ()) -> (Value -> String) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> String
forall a. Pretty a => a -> String
pretty
    outValue OutputFormat
Binary = ByteString -> IO ()
BS.putStr (ByteString -> IO ()) -> (Value -> ByteString) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ByteString
forall a. Binary a => a -> ByteString
Bin.encode
    outValue OutputFormat
Type = String -> IO ()
putStrLn (String -> IO ()) -> (Value -> String) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueType -> String
forall a. Pretty a => a -> String
pretty (ValueType -> String) -> (Value -> ValueType) -> Value -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ValueType
valueType

toValueType :: UncheckedTypeExp -> Either String ValueType
toValueType :: UncheckedTypeExp -> Either String ValueType
toValueType TETuple {} = String -> Either String ValueType
forall a b. a -> Either a b
Left String
"Cannot handle tuples yet."
toValueType TERecord {} = String -> Either String ValueType
forall a b. a -> Either a b
Left String
"Cannot handle records yet."
toValueType TEApply {} = String -> Either String ValueType
forall a b. a -> Either a b
Left String
"Cannot handle type applications yet."
toValueType TEArrow {} = String -> Either String ValueType
forall a b. a -> Either a b
Left String
"Cannot generate functions."
toValueType TESum {} = String -> Either String ValueType
forall a b. a -> Either a b
Left String
"Cannot handle sumtypes yet."
toValueType (TEUnique UncheckedTypeExp
t SrcLoc
_) = UncheckedTypeExp -> Either String ValueType
toValueType UncheckedTypeExp
t
toValueType (TEArray UncheckedTypeExp
t DimExp Name
d SrcLoc
_) = do
  Int
d' <- DimExp Name -> Either String Int
forall {a} {vn}. IsString a => DimExp vn -> Either a Int
constantDim DimExp Name
d
  ValueType [Int]
ds PrimType
t' <- UncheckedTypeExp -> Either String ValueType
toValueType UncheckedTypeExp
t
  ValueType -> Either String ValueType
forall (m :: * -> *) a. Monad m => a -> m a
return (ValueType -> Either String ValueType)
-> ValueType -> Either String ValueType
forall a b. (a -> b) -> a -> b
$ [Int] -> PrimType -> ValueType
ValueType (Int
d' Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
ds) PrimType
t'
  where
    constantDim :: DimExp vn -> Either a Int
constantDim (DimExpConst Int
k SrcLoc
_) = Int -> Either a Int
forall a b. b -> Either a b
Right Int
k
    constantDim DimExp vn
_ = a -> Either a Int
forall a b. a -> Either a b
Left a
"Array has non-constant dimension declaration."
toValueType (TEVar (QualName [] Name
v) SrcLoc
_)
  | Just PrimType
t <- Name -> Map Name PrimType -> Maybe PrimType
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
v Map Name PrimType
namesToPrimTypes = ValueType -> Either String ValueType
forall a b. b -> Either a b
Right (ValueType -> Either String ValueType)
-> ValueType -> Either String ValueType
forall a b. (a -> b) -> a -> b
$ [Int] -> PrimType -> ValueType
ValueType [] PrimType
t
toValueType (TEVar QualName Name
v SrcLoc
_) =
  String -> Either String ValueType
forall a b. a -> Either a b
Left (String -> Either String ValueType)
-> String -> Either String ValueType
forall a b. (a -> b) -> a -> b
$ String
"Unknown type " String -> ShowS
forall a. [a] -> [a] -> [a]
++ QualName Name -> String
forall a. Pretty a => a -> String
pretty QualName Name
v

-- | Closed interval, as in @System.Random@.
type Range a = (a, a)

data RandomConfiguration = RandomConfiguration
  { RandomConfiguration -> Range Int8
i8Range :: Range Int8,
    RandomConfiguration -> Range Int16
i16Range :: Range Int16,
    RandomConfiguration -> Range Int32
i32Range :: Range Int32,
    RandomConfiguration -> Range Int64
i64Range :: Range Int64,
    RandomConfiguration -> Range Word8
u8Range :: Range Word8,
    RandomConfiguration -> Range Word16
u16Range :: Range Word16,
    RandomConfiguration -> Range Word32
u32Range :: Range Word32,
    RandomConfiguration -> Range Word64
u64Range :: Range Word64,
    RandomConfiguration -> Range Float
f32Range :: Range Float,
    RandomConfiguration -> Range Double
f64Range :: Range Double
  }

-- The following lines provide evidence about how Haskells record
-- system sucks.
seti8Range :: Range Int8 -> RandomConfiguration -> RandomConfiguration
seti8Range :: Range Int8 -> RandomConfiguration -> RandomConfiguration
seti8Range Range Int8
bounds RandomConfiguration
config = RandomConfiguration
config {i8Range :: Range Int8
i8Range = Range Int8
bounds}

seti16Range :: Range Int16 -> RandomConfiguration -> RandomConfiguration
seti16Range :: Range Int16 -> RandomConfiguration -> RandomConfiguration
seti16Range Range Int16
bounds RandomConfiguration
config = RandomConfiguration
config {i16Range :: Range Int16
i16Range = Range Int16
bounds}

seti32Range :: Range Int32 -> RandomConfiguration -> RandomConfiguration
seti32Range :: Range Int32 -> RandomConfiguration -> RandomConfiguration
seti32Range Range Int32
bounds RandomConfiguration
config = RandomConfiguration
config {i32Range :: Range Int32
i32Range = Range Int32
bounds}

seti64Range :: Range Int64 -> RandomConfiguration -> RandomConfiguration
seti64Range :: Range Int64 -> RandomConfiguration -> RandomConfiguration
seti64Range Range Int64
bounds RandomConfiguration
config = RandomConfiguration
config {i64Range :: Range Int64
i64Range = Range Int64
bounds}

setu8Range :: Range Word8 -> RandomConfiguration -> RandomConfiguration
setu8Range :: Range Word8 -> RandomConfiguration -> RandomConfiguration
setu8Range Range Word8
bounds RandomConfiguration
config = RandomConfiguration
config {u8Range :: Range Word8
u8Range = Range Word8
bounds}

setu16Range :: Range Word16 -> RandomConfiguration -> RandomConfiguration
setu16Range :: Range Word16 -> RandomConfiguration -> RandomConfiguration
setu16Range Range Word16
bounds RandomConfiguration
config = RandomConfiguration
config {u16Range :: Range Word16
u16Range = Range Word16
bounds}

setu32Range :: Range Word32 -> RandomConfiguration -> RandomConfiguration
setu32Range :: Range Word32 -> RandomConfiguration -> RandomConfiguration
setu32Range Range Word32
bounds RandomConfiguration
config = RandomConfiguration
config {u32Range :: Range Word32
u32Range = Range Word32
bounds}

setu64Range :: Range Word64 -> RandomConfiguration -> RandomConfiguration
setu64Range :: Range Word64 -> RandomConfiguration -> RandomConfiguration
setu64Range Range Word64
bounds RandomConfiguration
config = RandomConfiguration
config {u64Range :: Range Word64
u64Range = Range Word64
bounds}

setf32Range :: Range Float -> RandomConfiguration -> RandomConfiguration
setf32Range :: Range Float -> RandomConfiguration -> RandomConfiguration
setf32Range Range Float
bounds RandomConfiguration
config = RandomConfiguration
config {f32Range :: Range Float
f32Range = Range Float
bounds}

setf64Range :: Range Double -> RandomConfiguration -> RandomConfiguration
setf64Range :: Range Double -> RandomConfiguration -> RandomConfiguration
setf64Range Range Double
bounds RandomConfiguration
config = RandomConfiguration
config {f64Range :: Range Double
f64Range = Range Double
bounds}

initialRandomConfiguration :: RandomConfiguration
initialRandomConfiguration :: RandomConfiguration
initialRandomConfiguration =
  Range Int8
-> Range Int16
-> Range Int32
-> Range Int64
-> Range Word8
-> Range Word16
-> Range Word32
-> Range Word64
-> Range Float
-> Range Double
-> RandomConfiguration
RandomConfiguration
    (Int8
forall a. Bounded a => a
minBound, Int8
forall a. Bounded a => a
maxBound)
    (Int16
forall a. Bounded a => a
minBound, Int16
forall a. Bounded a => a
maxBound)
    (Int32
forall a. Bounded a => a
minBound, Int32
forall a. Bounded a => a
maxBound)
    (Int64
forall a. Bounded a => a
minBound, Int64
forall a. Bounded a => a
maxBound)
    (Word8
forall a. Bounded a => a
minBound, Word8
forall a. Bounded a => a
maxBound)
    (Word16
forall a. Bounded a => a
minBound, Word16
forall a. Bounded a => a
maxBound)
    (Word32
forall a. Bounded a => a
minBound, Word32
forall a. Bounded a => a
maxBound)
    (Word64
forall a. Bounded a => a
minBound, Word64
forall a. Bounded a => a
maxBound)
    (Float
0.0, Float
1.0)
    (Double
0.0, Double
1.0)

randomValue :: RandomConfiguration -> ValueType -> Word64 -> Value
randomValue :: RandomConfiguration -> ValueType -> Word64 -> Value
randomValue RandomConfiguration
conf (ValueType [Int]
ds PrimType
t) Word64
seed =
  case PrimType
t of
    Signed IntType
Int8 -> (RandomConfiguration -> Range Int8)
-> (Vector Int -> Vector Int8 -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Int8
i8Range Vector Int -> Vector Int8 -> Value
Int8Value
    Signed IntType
Int16 -> (RandomConfiguration -> Range Int16)
-> (Vector Int -> Vector Int16 -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Int16
i16Range Vector Int -> Vector Int16 -> Value
Int16Value
    Signed IntType
Int32 -> (RandomConfiguration -> Range Int32)
-> (Vector Int -> Vector Int32 -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Int32
i32Range Vector Int -> Vector Int32 -> Value
Int32Value
    Signed IntType
Int64 -> (RandomConfiguration -> Range Int64)
-> (Vector Int -> Vector Int64 -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Int64
i64Range Vector Int -> Vector Int64 -> Value
Int64Value
    Unsigned IntType
Int8 -> (RandomConfiguration -> Range Word8)
-> (Vector Int -> Vector Word8 -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Word8
u8Range Vector Int -> Vector Word8 -> Value
Word8Value
    Unsigned IntType
Int16 -> (RandomConfiguration -> Range Word16)
-> (Vector Int -> Vector Word16 -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Word16
u16Range Vector Int -> Vector Word16 -> Value
Word16Value
    Unsigned IntType
Int32 -> (RandomConfiguration -> Range Word32)
-> (Vector Int -> Vector Word32 -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Word32
u32Range Vector Int -> Vector Word32 -> Value
Word32Value
    Unsigned IntType
Int64 -> (RandomConfiguration -> Range Word64)
-> (Vector Int -> Vector Word64 -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Word64
u64Range Vector Int -> Vector Word64 -> Value
Word64Value
    FloatType FloatType
Float32 -> (RandomConfiguration -> Range Float)
-> (Vector Int -> Vector Float -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Float
f32Range Vector Int -> Vector Float -> Value
Float32Value
    FloatType FloatType
Float64 -> (RandomConfiguration -> Range Double)
-> (Vector Int -> Vector Double -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Double
f64Range Vector Int -> Vector Double -> Value
Float64Value
    PrimType
Bool -> (RandomConfiguration -> Range Bool)
-> (Vector Int -> Vector Bool -> Value) -> Value
forall {v}.
(Storable v, Variate v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen (Range Bool -> RandomConfiguration -> Range Bool
forall a b. a -> b -> a
const (Bool
False, Bool
True)) Vector Int -> Vector Bool -> Value
BoolValue
  where
    gen :: (RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range v
range Vector Int -> Vector v -> Value
final = Range v
-> (Vector Int -> Vector v -> Value) -> [Int] -> Word64 -> Value
forall v.
(Storable v, Variate v) =>
Range v
-> (Vector Int -> Vector v -> Value) -> [Int] -> Word64 -> Value
randomVector (RandomConfiguration -> Range v
range RandomConfiguration
conf) Vector Int -> Vector v -> Value
final [Int]
ds Word64
seed

randomVector ::
  (SVec.Storable v, Variate v) =>
  Range v ->
  (SVec.Vector Int -> SVec.Vector v -> Value) ->
  [Int] ->
  Word64 ->
  Value
randomVector :: forall v.
(Storable v, Variate v) =>
Range v
-> (Vector Int -> Vector v -> Value) -> [Int] -> Word64 -> Value
randomVector Range v
range Vector Int -> Vector v -> Value
final [Int]
ds Word64
seed = (forall s. ST s Value) -> Value
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s Value) -> Value)
-> (forall s. ST s Value) -> Value
forall a b. (a -> b) -> a -> b
$ do
  -- USe some nice impure computation where we can preallocate a
  -- vector of the desired size, populate it via the random number
  -- generator, and then finally reutrn a frozen binary vector.
  MVector s v
arr <- Int -> ST s (MVector (PrimState (ST s)) v)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
USVec.new Int
n
  Gen s
g <- Word64 -> Word64 -> ST s (Gen (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Word64 -> Word64 -> m (Gen (PrimState m))
initialize Word64
6364136223846793006 Word64
seed
  let fill :: Int -> ST s Value
fill Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = do
          v
v <- Range v -> Gen s -> ST s v
forall a g (m :: * -> *).
(Variate a, Generator g m) =>
(a, a) -> g -> m a
uniformR Range v
range Gen s
g
          MVector (PrimState (ST s)) v -> Int -> v -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
USVec.write MVector s v
MVector (PrimState (ST s)) v
arr Int
i v
v
          Int -> ST s Value
fill (Int -> ST s Value) -> Int -> ST s Value
forall a b. (a -> b) -> a -> b
$! Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        | Bool
otherwise =
          Vector Int -> Vector v -> Value
final ([Int] -> Vector Int
forall a. Storable a => [a] -> Vector a
SVec.fromList [Int]
ds) (Vector v -> Value) -> (Vector v -> Vector v) -> Vector v -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector v -> Vector v
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
SVec.convert (Vector v -> Value) -> ST s (Vector v) -> ST s Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mutable Vector (PrimState (ST s)) v -> ST s (Vector v)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
freeze MVector s v
Mutable Vector (PrimState (ST s)) v
arr
  Int -> ST s Value
fill Int
0
  where
    n :: Int
n = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
ds