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

-- | This module provides an efficient value representation as well as
-- parsing and comparison functions.
module Futhark.Test.Values
  ( module Futhark.Data,
    module Futhark.Data.Compare,
    module Futhark.Data.Reader,

    -- * Random value generation
    Range,
    RandomConfiguration (..),
    initialRandomConfiguration,
    randomValue,
  )
where

import Control.Monad.ST
import Data.Int
import Data.Vector.Storable qualified as SVec
import Data.Vector.Storable.Mutable qualified as USVec
import Data.Word
import Futhark.Data
import Futhark.Data.Compare
import Futhark.Data.Reader
import Futhark.Util (convFloat)
import Futhark.Util.Pretty (Pretty (..))
import Numeric.Half
import System.Random.Stateful (UniformRange (..), mkStdGen, uniformR)

instance Pretty Value where
  pretty :: forall ann. Value -> Doc ann
pretty = Text -> Doc ann
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Text -> Doc ann) -> (Value -> Text) -> Value -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Text
valueText

instance Pretty ValueType where
  pretty :: forall ann. ValueType -> Doc ann
pretty = Text -> Doc ann
forall ann. Text -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Text -> Doc ann) -> (ValueType -> Text) -> ValueType -> Doc ann
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ValueType -> Text
valueTypeText

randomVector ::
  (SVec.Storable v, UniformRange v) =>
  Range v ->
  (SVec.Vector Int -> SVec.Vector v -> Value) ->
  [Int] ->
  Word64 ->
  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 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.
  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
  let fill t
g Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = do
            let (v
v, t
g') = Range v -> t -> (v, t)
forall a g. (UniformRange a, RandomGen g) => (a, a) -> g -> (a, g)
uniformR Range v
range t
g
            MVector (PrimState m) v -> Int -> v -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Int -> a -> m ()
USVec.write MVector s v
MVector (PrimState m) v
arr Int
i v
v
            t
g' t -> m () -> m ()
forall a b. a -> b -> b
`seq` t -> Int -> m ()
fill t
g' (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$! Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        | Bool
otherwise =
            () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  fill (mkStdGen $ fromIntegral seed) 0
  final (SVec.fromList ds) . SVec.convert <$> SVec.freeze arr
  where
    n :: Int
n = [Int] -> Int
forall a. Num a => [a] -> a
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 =>
(Half, Half) -> g -> m Half
uniformRM (Half
a, Half
b) g
g =
    (Float -> Half
forall from to. (RealFloat from, RealFloat to) => from -> to
convFloat :: Float -> Half) (Float -> Half) -> m Float -> m Half
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Float, Float) -> g -> m Float
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *).
StatefulGen g m =>
(Float, Float) -> g -> m Float
uniformRM (Half -> Float
forall from to. (RealFloat from, RealFloat to) => from -> to
convFloat Half
a, Half -> Float
forall from to. (RealFloat from, RealFloat to) => from -> to
convFloat Half
b) g
g

-- | 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 -> (Half, Half)
f16Range :: Range Half,
    RandomConfiguration -> (Float, Float)
f32Range :: Range Float,
    RandomConfiguration -> Range Double
f64Range :: Range Double
  }

initialRandomConfiguration :: RandomConfiguration
initialRandomConfiguration :: RandomConfiguration
initialRandomConfiguration =
  Range Int8
-> Range Int16
-> Range Int32
-> Range Int64
-> Range Word8
-> Range Word16
-> Range Word32
-> Range Word64
-> (Half, Half)
-> (Float, 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)
    (Half
0.0, Half
1.0)
    (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
    PrimType
I8 -> (RandomConfiguration -> Range Int8)
-> (Vector Int -> Vector Int8 -> Value) -> Value
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
I8Value
    PrimType
I16 -> (RandomConfiguration -> Range Int16)
-> (Vector Int -> Vector Int16 -> Value) -> Value
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
I16Value
    PrimType
I32 -> (RandomConfiguration -> Range Int32)
-> (Vector Int -> Vector Int32 -> Value) -> Value
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
I32Value
    PrimType
I64 -> (RandomConfiguration -> Range Int64)
-> (Vector Int -> Vector Int64 -> Value) -> Value
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
I64Value
    PrimType
U8 -> (RandomConfiguration -> Range Word8)
-> (Vector Int -> Vector Word8 -> Value) -> Value
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
U8Value
    PrimType
U16 -> (RandomConfiguration -> Range Word16)
-> (Vector Int -> Vector Word16 -> Value) -> Value
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
U16Value
    PrimType
U32 -> (RandomConfiguration -> Range Word32)
-> (Vector Int -> Vector Word32 -> Value) -> Value
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
U32Value
    PrimType
U64 -> (RandomConfiguration -> Range Word64)
-> (Vector Int -> Vector Word64 -> Value) -> Value
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
U64Value
    PrimType
F16 -> (RandomConfiguration -> (Half, Half))
-> (Vector Int -> Vector Half -> Value) -> Value
forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> (Half, Half)
f16Range Vector Int -> Vector Half -> Value
F16Value
    PrimType
F32 -> (RandomConfiguration -> (Float, Float))
-> (Vector Int -> Vector Float -> Value) -> Value
forall {v}.
(Storable v, UniformRange v) =>
(RandomConfiguration -> Range v)
-> (Vector Int -> Vector v -> Value) -> Value
gen RandomConfiguration -> (Float, Float)
f32Range Vector Int -> Vector Float -> Value
F32Value
    PrimType
F64 -> (RandomConfiguration -> Range Double)
-> (Vector Int -> Vector Double -> Value) -> Value
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
F64Value
    PrimType
Bool -> (RandomConfiguration -> Range Bool)
-> (Vector Int -> Vector Bool -> Value) -> Value
forall {v}.
(Storable v, UniformRange 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, 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