{-# LANGUAGE Strict #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

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

import Control.Monad
import Control.Monad.ST
import Data.Binary qualified as Bin
import Data.ByteString.Lazy.Char8 qualified as BS
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Data.Vector.Generic (freeze)
import Data.Vector.Storable qualified as SVec
import Data.Vector.Storable.Mutable qualified as USVec
import Data.Word
import Futhark.Data qualified as V
import Futhark.Data.Reader (readValues)
import Futhark.Util (convFloat)
import Futhark.Util.Options
import Language.Futhark.Parser
import Language.Futhark.Pretty ()
import Language.Futhark.Prop (UncheckedTypeExp)
import Language.Futhark.Syntax hiding
  ( FloatValue (..),
    IntValue (..),
    PrimValue (..),
    ValueType,
  )
import System.Exit
import System.IO
import System.Random (mkStdGen, uniformR)
import System.Random.Stateful (UniformRange (..))

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

data OutputFormat
  = Text
  | Binary
  | Type
  deriving (OutputFormat -> OutputFormat -> Bool
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
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
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 =
  [ forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"s"
      [String
"seed"]
      ( forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          ( \String
n ->
              case forall a. Read a => ReadS a
reads String
n of
                [(Int
n', String
"")] ->
                  forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ \DataOptions
config -> DataOptions
config {optSeed :: Int
optSeed = Int
n'}
                [(Int, String)]
_ ->
                  forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ do
                    Handle -> String -> IO ()
hPutStrLn Handle
stderr forall a b. (a -> b) -> a -> b
$ String
"'" forall a. [a] -> [a] -> [a]
++ String
n forall a. [a] -> [a] -> [a]
++ String
"' is not an integer."
                    forall a. IO a
exitFailure
          )
          String
"SEED"
      )
      String
"The seed to use when initialising the RNG.",
    forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"g"
      [String
"generate"]
      ( forall a. (String -> a) -> String -> ArgDescr a
ReqArg
          ( \String
t ->
              case String
-> Either
     Text (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
tryMakeGenerator String
t of
                Right RandomConfiguration -> OutputFormat -> Word64 -> IO ()
g ->
                  forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ \DataOptions
config ->
                    DataOptions
config
                      { optOrders :: [Word64 -> IO ()]
optOrders =
                          DataOptions -> [Word64 -> IO ()]
optOrders DataOptions
config
                            forall a. [a] -> [a] -> [a]
++ [RandomConfiguration -> OutputFormat -> Word64 -> IO ()
g (DataOptions -> RandomConfiguration
optRange DataOptions
config) (DataOptions -> OutputFormat
format DataOptions
config)]
                      }
                Left Text
err ->
                  forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ do
                    Handle -> Text -> IO ()
T.hPutStrLn Handle
stderr Text
err
                    forall a. IO a
exitFailure
          )
          String
"TYPE"
      )
      String
"Generate a random value of this type.",
    forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      []
      [String
"text"]
      (forall a. a -> ArgDescr a
NoArg forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ \DataOptions
opts -> DataOptions
opts {format :: OutputFormat
format = OutputFormat
Text})
      String
"Output data in text format (default; must precede --generate).",
    forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"b"
      [String
"binary"]
      (forall a. a -> ArgDescr a
NoArg forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right 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).",
    forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
"t"
      [String
"type"]
      (forall a. a -> ArgDescr a
NoArg forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right 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).",
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"i8" Range Int8 -> RandomConfiguration -> RandomConfiguration
seti8Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"i16" Range Int16 -> RandomConfiguration -> RandomConfiguration
seti16Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"i32" Range Int32 -> RandomConfiguration -> RandomConfiguration
seti32Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"i64" Range Int64 -> RandomConfiguration -> RandomConfiguration
seti64Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"u8" Range Word8 -> RandomConfiguration -> RandomConfiguration
setu8Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"u16" Range Word16 -> RandomConfiguration -> RandomConfiguration
setu16Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"u32" Range Word32 -> RandomConfiguration -> RandomConfiguration
setu32Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"u64" Range Word64 -> RandomConfiguration -> RandomConfiguration
setu64Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"f16" Range Half -> RandomConfiguration -> RandomConfiguration
setf16Range,
    forall a.
Read a =>
String
-> (Range a -> RandomConfiguration -> RandomConfiguration)
-> FunOptDescr DataOptions
setRangeOption String
"f32" Range Float -> RandomConfiguration -> RandomConfiguration
setf32Range,
    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 =
  forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
    String
""
    [String
name]
    ( forall a. (String -> a) -> String -> ArgDescr a
ReqArg
        ( \String
b ->
            let (String
lower, String
rest) = forall a. (a -> Bool) -> [a] -> ([a], [a])
span (forall a. Eq a => a -> a -> Bool
/= Char
':') String
b
                upper :: String
upper = forall a. Int -> [a] -> [a]
drop Int
1 String
rest
             in case (forall a. Read a => ReadS a
reads String
lower, forall a. Read a => ReadS a
reads String
upper) of
                  ([(a
lower', String
"")], [(a
upper', String
"")]) ->
                    forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ \DataOptions
config ->
                      DataOptions
config {optRange :: RandomConfiguration
optRange = Range a -> RandomConfiguration -> RandomConfiguration
set (a
lower', a
upper') forall a b. (a -> b) -> a -> b
$ DataOptions -> RandomConfiguration
optRange DataOptions
config}
                  ([(a, String)], [(a, String)])
_ ->
                    forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ do
                      Handle -> String -> IO ()
hPutStrLn Handle
stderr forall a b. (a -> b) -> a -> b
$ String
"Invalid bounds for " forall a. [a] -> [a] -> [a]
++ String
tname forall a. [a] -> [a] -> [a]
++ String
": " forall a. [a] -> [a] -> [a]
++ String
b
                      forall a. IO a
exitFailure
        )
        String
"MIN:MAX"
    )
    forall a b. (a -> b) -> a -> b
$ String
"Range of " forall a. [a] -> [a] -> [a]
++ String
tname forall a. [a] -> [a] -> [a]
++ String
" values."
  where
    name :: String
name = String
tname forall a. [a] -> [a] -> [a]
++ String
"-bounds"

tryMakeGenerator ::
  String ->
  Either T.Text (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
tryMakeGenerator :: String
-> Either
     Text (RandomConfiguration -> OutputFormat -> Word64 -> IO ())
tryMakeGenerator String
t
  | Just [Value]
vs <- ByteString -> Maybe [Value]
readValues forall a b. (a -> b) -> a -> b
$ String -> ByteString
BS.pack String
t =
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ \RandomConfiguration
_ OutputFormat
fmt Word64
_ -> 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 Text ValueType
toValueType forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. SyntaxError -> Text
syntaxErrorMsg) forall a b. b -> Either a b
Right (String -> Text -> Either SyntaxError UncheckedTypeExp
parseType String
name (String -> Text
T.pack String
t))
      forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 " forall a. [a] -> [a] -> [a]
++ String
t
    outValue :: OutputFormat -> Value -> IO ()
outValue OutputFormat
Text = Text -> IO ()
T.putStrLn forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Text
V.valueText
    outValue OutputFormat
Binary = ByteString -> IO ()
BS.putStr forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Binary a => a -> ByteString
Bin.encode
    outValue OutputFormat
Type = Text -> IO ()
T.putStrLn forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueType -> Text
V.valueTypeText forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> ValueType
V.valueType

toValueType :: UncheckedTypeExp -> Either T.Text V.ValueType
toValueType :: UncheckedTypeExp -> Either Text ValueType
toValueType TETuple {} = forall a b. a -> Either a b
Left Text
"Cannot handle tuples yet."
toValueType TERecord {} = forall a b. a -> Either a b
Left Text
"Cannot handle records yet."
toValueType TEApply {} = forall a b. a -> Either a b
Left Text
"Cannot handle type applications yet."
toValueType TEArrow {} = forall a b. a -> Either a b
Left Text
"Cannot generate functions."
toValueType TESum {} = forall a b. a -> Either a b
Left Text
"Cannot handle sumtypes yet."
toValueType TEDim {} = forall a b. a -> Either a b
Left Text
"Cannot handle existential sizes."
toValueType (TEParens UncheckedTypeExp
t SrcLoc
_) = UncheckedTypeExp -> Either Text ValueType
toValueType UncheckedTypeExp
t
toValueType (TEUnique UncheckedTypeExp
t SrcLoc
_) = UncheckedTypeExp -> Either Text ValueType
toValueType UncheckedTypeExp
t
toValueType (TEArray SizeExp NoInfo Name
d UncheckedTypeExp
t SrcLoc
_) = do
  Int
d' <- forall {b} {a} {f :: * -> *} {vn}.
(Num b, IsString a) =>
SizeExp f vn -> Either a b
constantDim SizeExp NoInfo Name
d
  V.ValueType [Int]
ds PrimType
t' <- UncheckedTypeExp -> Either Text ValueType
toValueType UncheckedTypeExp
t
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Int] -> PrimType -> ValueType
V.ValueType (Int
d' forall a. a -> [a] -> [a]
: [Int]
ds) PrimType
t'
  where
    constantDim :: SizeExp f vn -> Either a b
constantDim (SizeExp (IntLit Integer
k f PatType
_ SrcLoc
_) SrcLoc
_) = forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ forall a. Num a => Integer -> a
fromInteger Integer
k
    constantDim SizeExp f vn
_ = forall a b. a -> Either a b
Left a
"Array has non-constant dimension declaration."
toValueType (TEVar (QualName [] Name
v) SrcLoc
_)
  | Just PrimType
t <- forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
v [(Name, PrimType)]
m = forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ [Int] -> PrimType -> ValueType
V.ValueType [] PrimType
t
  where
    m :: [(Name, PrimType)]
m = forall a b. (a -> b) -> [a] -> [b]
map PrimType -> (Name, PrimType)
f [forall a. Bounded a => a
minBound .. forall a. Bounded a => a
maxBound]
    f :: PrimType -> (Name, PrimType)
f PrimType
t = (Text -> Name
nameFromText (PrimType -> Text
V.primTypeText PrimType
t), PrimType
t)
toValueType (TEVar QualName Name
v SrcLoc
_) =
  forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ Text
"Unknown type " forall a. Semigroup a => a -> a -> a
<> forall a. Pretty a => a -> Text
prettyText 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 Half
f16Range :: Range Half,
    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}

setf16Range :: Range Half -> RandomConfiguration -> RandomConfiguration
setf16Range :: Range Half -> RandomConfiguration -> RandomConfiguration
setf16Range Range Half
bounds RandomConfiguration
config = RandomConfiguration
config {f16Range :: Range Half
f16Range = Range Half
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 Half
-> Range Float
-> Range Double
-> RandomConfiguration
RandomConfiguration
    (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound)
    (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound)
    (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound)
    (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound)
    (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound)
    (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound)
    (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound)
    (forall a. Bounded a => a
minBound, forall a. Bounded a => a
maxBound)
    (Half
0.0, Half
1.0)
    (Float
0.0, Float
1.0)
    (Double
0.0, Double
1.0)

randomValue :: RandomConfiguration -> V.ValueType -> Word64 -> V.Value
randomValue :: RandomConfiguration -> ValueType -> Word64 -> Value
randomValue RandomConfiguration
conf (V.ValueType [Int]
ds PrimType
t) Word64
seed =
  case PrimType
t of
    PrimType
V.I8 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Int8
i8Range Vector Int -> Vector Int8 -> Value
V.I8Value
    PrimType
V.I16 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Int16
i16Range Vector Int -> Vector Int16 -> Value
V.I16Value
    PrimType
V.I32 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Int32
i32Range Vector Int -> Vector Int32 -> Value
V.I32Value
    PrimType
V.I64 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Int64
i64Range Vector Int -> Vector Int64 -> Value
V.I64Value
    PrimType
V.U8 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Word8
u8Range Vector Int -> Vector Word8 -> Value
V.U8Value
    PrimType
V.U16 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Word16
u16Range Vector Int -> Vector Word16 -> Value
V.U16Value
    PrimType
V.U32 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Word32
u32Range Vector Int -> Vector Word32 -> Value
V.U32Value
    PrimType
V.U64 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Word64
u64Range Vector Int -> Vector Word64 -> Value
V.U64Value
    PrimType
V.F16 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Half
f16Range Vector Int -> Vector Half -> Value
V.F16Value
    PrimType
V.F32 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Float
f32Range Vector Int -> Vector Float -> Value
V.F32Value
    PrimType
V.F64 -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range Double
f64Range Vector Int -> Vector Double -> Value
V.F64Value
    PrimType
V.Bool -> forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen (forall a b. a -> b -> a
const (Bool
False, Bool
True)) Vector Int -> Vector Bool -> Value
V.BoolValue
  where
    gen :: (RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> Range v
range Vector Int -> Vector v -> Value
final = forall v.
(Storable v, UniformRange 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, UniformRange v) =>
  Range v ->
  (SVec.Vector Int -> SVec.Vector v -> V.Value) ->
  [Int] ->
  Word64 ->
  V.Value
randomVector :: forall v.
(Storable v, UniformRange 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 a. (forall s. ST s a) -> a
runST 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 (PrimState (ST s)) v
arr <- forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
Int -> m (MVector (PrimState m) a)
USVec.new Int
n
  let fill :: t -> Int -> ST s ()
fill t
g Int
i
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
n = do
            let (v
v, t
g') = forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR Range v
range t
g
            forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
USVec.write MVector (PrimState (ST s)) v
arr Int
i v
v
            t
g' seq :: forall a b. a -> b -> b
`seq` t -> Int -> ST s ()
fill t
g' forall a b. (a -> b) -> a -> b
$! Int
i forall a. Num a => a -> a -> a
+ Int
1
        | Bool
otherwise =
            forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  forall {t}. RandomGen t => t -> Int -> ST s ()
fill (Int -> StdGen
mkStdGen forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
seed) Int
0
  Vector Int -> Vector v -> Value
final (forall a. Storable a => [a] -> Vector a
SVec.fromList [Int]
ds) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
SVec.convert forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
freeze MVector (PrimState (ST s)) v
arr
  where
    n :: Int
n = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
ds

-- XXX: The following instance is an orphan.  Maybe it could be
-- avoided with some newtype trickery or refactoring, but it's so
-- convenient this way.
instance UniformRange Half where
  uniformRM :: forall g (m :: * -> *).
StatefulGen g m =>
Range Half -> g -> m Half
uniformRM (Half
a, Half
b) g
g =
    (forall from to. (RealFloat from, RealFloat to) => from -> to
convFloat :: Float -> Half) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
uniformRM (forall from to. (RealFloat from, RealFloat to) => from -> to
convFloat Half
a, forall from to. (RealFloat from, RealFloat to) => from -> to
convFloat Half
b) g
g