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

-- | This module defines an efficient value representation as well as
-- parsing and comparison functions.  This is because the standard
-- Futhark parser is not able to cope with large values (like arrays
-- that are tens of megabytes in size).  The representation defined
-- here does not support tuples, so don't use those as input/output
-- for your test programs.
module Futhark.Test.Values
  ( Value (..),
    Vector,

    -- * Reading Values
    readValues,

    -- * Types of values
    ValueType (..),
    valueType,

    -- * Comparing Values
    compareValues,
    compareValues1,
    Mismatch,
  )
where

import Control.Monad
import Control.Monad.ST
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString.Lazy.Char8 as BS
import Data.Char (chr, isSpace, ord)
import Data.Int (Int16, Int32, Int64, Int8)
import Data.Vector.Binary
import Data.Vector.Generic (freeze)
import qualified Data.Vector.Unboxed as UVec
import qualified Data.Vector.Unboxed.Mutable as UMVec
import Futhark.IR.Pretty ()
import Futhark.IR.Primitive (PrimValue)
import Futhark.IR.Prop.Constants (IsValue (..))
import Futhark.Util (maybeHead)
import Futhark.Util.Loc (Pos (..))
import Futhark.Util.Pretty
import qualified Futhark.Util.Pretty as PP
import Language.Futhark.Parser.Lexer
import Language.Futhark.Pretty ()
import qualified Language.Futhark.Syntax as F

type STVector s = UMVec.STVector s

-- | An Unboxed vector.
type Vector = UVec.Vector

-- | An efficiently represented Futhark value.  Use 'pretty' to get a
-- human-readable representation, and v'put' to obtain binary a
-- representation.
data Value
  = Int8Value (Vector Int) (Vector Int8)
  | Int16Value (Vector Int) (Vector Int16)
  | Int32Value (Vector Int) (Vector Int32)
  | Int64Value (Vector Int) (Vector Int64)
  | Word8Value (Vector Int) (Vector Word8)
  | Word16Value (Vector Int) (Vector Word16)
  | Word32Value (Vector Int) (Vector Word32)
  | Word64Value (Vector Int) (Vector Word64)
  | Float32Value (Vector Int) (Vector Float)
  | Float64Value (Vector Int) (Vector Double)
  | BoolValue (Vector Int) (Vector Bool)
  deriving (Int -> Value -> ShowS
[Value] -> ShowS
Value -> String
(Int -> Value -> ShowS)
-> (Value -> String) -> ([Value] -> ShowS) -> Show Value
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Value] -> ShowS
$cshowList :: [Value] -> ShowS
show :: Value -> String
$cshow :: Value -> String
showsPrec :: Int -> Value -> ShowS
$cshowsPrec :: Int -> Value -> ShowS
Show)

binaryFormatVersion :: Word8
binaryFormatVersion :: Word8
binaryFormatVersion = Word8
2

instance Binary Value where
  put :: Value -> Put
put (Int8Value Vector Int
shape Vector Int8
vs) = String -> Vector Int -> Vector Int8 -> (Int8 -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
"  i8" Vector Int
shape Vector Int8
vs Int8 -> Put
putInt8
  put (Int16Value Vector Int
shape Vector Int16
vs) = String -> Vector Int -> Vector Int16 -> (Int16 -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
" i16" Vector Int
shape Vector Int16
vs Int16 -> Put
putInt16le
  put (Int32Value Vector Int
shape Vector Int32
vs) = String -> Vector Int -> Vector Int32 -> (Int32 -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
" i32" Vector Int
shape Vector Int32
vs Int32 -> Put
putInt32le
  put (Int64Value Vector Int
shape Vector Int64
vs) = String -> Vector Int -> Vector Int64 -> (Int64 -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
" i64" Vector Int
shape Vector Int64
vs Int64 -> Put
putInt64le
  put (Word8Value Vector Int
shape Vector Word8
vs) = String -> Vector Int -> Vector Word8 -> (Word8 -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
"  u8" Vector Int
shape Vector Word8
vs Word8 -> Put
putWord8
  put (Word16Value Vector Int
shape Vector Word16
vs) = String -> Vector Int -> Vector Word16 -> (Word16 -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
" u16" Vector Int
shape Vector Word16
vs Word16 -> Put
putWord16le
  put (Word32Value Vector Int
shape Vector Word32
vs) = String -> Vector Int -> Vector Word32 -> (Word32 -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
" u32" Vector Int
shape Vector Word32
vs Word32 -> Put
putWord32le
  put (Word64Value Vector Int
shape Vector Word64
vs) = String -> Vector Int -> Vector Word64 -> (Word64 -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
" u64" Vector Int
shape Vector Word64
vs Word64 -> Put
putWord64le
  put (Float32Value Vector Int
shape Vector Float
vs) = String -> Vector Int -> Vector Float -> (Float -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
" f32" Vector Int
shape Vector Float
vs Float -> Put
putFloatle
  put (Float64Value Vector Int
shape Vector Double
vs) = String -> Vector Int -> Vector Double -> (Double -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
" f64" Vector Int
shape Vector Double
vs Double -> Put
putDoublele
  put (BoolValue Vector Int
shape Vector Bool
vs) = String -> Vector Int -> Vector Bool -> (Bool -> Put) -> Put
forall a.
Unbox a =>
String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
"bool" Vector Int
shape Vector Bool
vs ((Bool -> Put) -> Put) -> (Bool -> Put) -> Put
forall a b. (a -> b) -> a -> b
$ Int8 -> Put
putInt8 (Int8 -> Put) -> (Bool -> Int8) -> Bool -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Int8
forall p. Num p => Bool -> p
boolToInt
    where
      boolToInt :: Bool -> p
boolToInt Bool
True = p
1
      boolToInt Bool
False = p
0

  get :: Get Value
get = do
    Int8
first <- Get Int8
getInt8
    Word8
version <- Get Word8
getWord8
    Int8
rank <- Get Int8
getInt8

    Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int -> Char
chr (Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
first) Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'b') (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$
      String -> Get ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Input does not begin with ASCII 'b'."
    Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word8
version Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
binaryFormatVersion) (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$
      String -> Get ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get ()) -> String -> Get ()
forall a b. (a -> b) -> a -> b
$ String
"Expecting binary format version 1; found version: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word8 -> String
forall a. Show a => a -> String
show Word8
version
    Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int8
rank Int8 -> Int8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Int8
0) (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$
      String -> Get ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get ()) -> String -> Get ()
forall a b. (a -> b) -> a -> b
$ String
"Rank must be non-negative, but is: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int8 -> String
forall a. Show a => a -> String
show Int8
rank

    ByteString
type_f <- Int64 -> Get ByteString
getLazyByteString Int64
4

    [Int]
shape <- Int -> Get Int -> Get [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
rank) (Get Int -> Get [Int]) -> Get Int -> Get [Int]
forall a b. (a -> b) -> a -> b
$ Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> Get Int64 -> Get Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Int64
getInt64le
    let num_elems :: Int
num_elems = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
shape
        shape' :: Vector Int
shape' = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape

    case ByteString -> String
BS.unpack ByteString
type_f of
      String
"  i8" -> (Vector Int8 -> Value) -> Get Int8 -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Int8 -> Value
Int8Value Vector Int
shape') Get Int8
getInt8 Int
num_elems
      String
" i16" -> (Vector Int16 -> Value) -> Get Int16 -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Int16 -> Value
Int16Value Vector Int
shape') Get Int16
getInt16le Int
num_elems
      String
" i32" -> (Vector Int32 -> Value) -> Get Int32 -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Int32 -> Value
Int32Value Vector Int
shape') Get Int32
getInt32le Int
num_elems
      String
" i64" -> (Vector Int64 -> Value) -> Get Int64 -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Int64 -> Value
Int64Value Vector Int
shape') Get Int64
getInt64le Int
num_elems
      String
"  u8" -> (Vector Word8 -> Value) -> Get Word8 -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Word8 -> Value
Word8Value Vector Int
shape') Get Word8
getWord8 Int
num_elems
      String
" u16" -> (Vector Word16 -> Value) -> Get Word16 -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Word16 -> Value
Word16Value Vector Int
shape') Get Word16
getWord16le Int
num_elems
      String
" u32" -> (Vector Word32 -> Value) -> Get Word32 -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Word32 -> Value
Word32Value Vector Int
shape') Get Word32
getWord32le Int
num_elems
      String
" u64" -> (Vector Word64 -> Value) -> Get Word64 -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Word64 -> Value
Word64Value Vector Int
shape') Get Word64
getWord64le Int
num_elems
      String
" f32" -> (Vector Float -> Value) -> Get Float -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Float -> Value
Float32Value Vector Int
shape') Get Float
getFloatle Int
num_elems
      String
" f64" -> (Vector Double -> Value) -> Get Double -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Double -> Value
Float64Value Vector Int
shape') Get Double
getDoublele Int
num_elems
      String
"bool" -> (Vector Bool -> Value) -> Get Bool -> Int -> Get Value
forall (v :: * -> *) a b.
Vector v a =>
(v a -> b) -> Get a -> Int -> Get b
get' (Vector Int -> Vector Bool -> Value
BoolValue Vector Int
shape') Get Bool
getBool Int
num_elems
      String
s -> String -> Get Value
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get Value) -> String -> Get Value
forall a b. (a -> b) -> a -> b
$ String
"Cannot parse binary values of type " String -> ShowS
forall a. [a] -> [a] -> [a]
++ ShowS
forall a. Show a => a -> String
show String
s
    where
      getBool :: Get Bool
getBool = (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
0) (Word8 -> Bool) -> Get Word8 -> Get Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8

      get' :: (v a -> b) -> Get a -> Int -> Get b
get' v a -> b
mk Get a
get_elem Int
num_elems =
        v a -> b
mk (v a -> b) -> Get (v a) -> Get b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Int -> Get a -> Get (v a)
forall (v :: * -> *) a. Vector v a => Get Int -> Get a -> Get (v a)
genericGetVectorWith (Int -> Get Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
num_elems) Get a
get_elem

putBinaryValue ::
  UVec.Unbox a =>
  String ->
  Vector Int ->
  Vector a ->
  (a -> Put) ->
  Put
putBinaryValue :: String -> Vector Int -> Vector a -> (a -> Put) -> Put
putBinaryValue String
tstr Vector Int
shape Vector a
vs a -> Put
putv = do
  Int8 -> Put
putInt8 (Int8 -> Put) -> Int8 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int8) -> Int -> Int8
forall a b. (a -> b) -> a -> b
$ Char -> Int
ord Char
'b'
  Word8 -> Put
putWord8 Word8
binaryFormatVersion
  Word8 -> Put
putWord8 (Word8 -> Put) -> Word8 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int
forall a. Unbox a => Vector a -> Int
UVec.length Vector Int
shape
  (Char -> Put) -> String -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Int8 -> Put
putInt8 (Int8 -> Put) -> (Char -> Int8) -> Char -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int8) -> (Char -> Int) -> Char -> Int8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Int
ord) String
tstr
  (Int -> Put) -> [Int] -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Int64 -> Put
putInt64le (Int64 -> Put) -> (Int -> Int64) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral) ([Int] -> Put) -> [Int] -> Put
forall a b. (a -> b) -> a -> b
$ Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
  (a -> Put) -> [a] -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ a -> Put
putv ([a] -> Put) -> [a] -> Put
forall a b. (a -> b) -> a -> b
$ Vector a -> [a]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector a
vs

instance PP.Pretty Value where
  ppr :: Value -> Doc
ppr Value
v
    | [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (Value -> [Int]
valueShape Value
v) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 =
      String -> Doc
text String
"empty"
        Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc -> Doc
parens (Doc
dims Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr (Value -> PrimType
valueElemType Value
v))
    where
      dims :: Doc
dims = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ (Int -> Doc) -> [Int] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> Doc
brackets (Doc -> Doc) -> (Int -> Doc) -> Int -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Doc
forall a. Pretty a => a -> Doc
ppr) ([Int] -> [Doc]) -> [Int] -> [Doc]
forall a b. (a -> b) -> a -> b
$ Value -> [Int]
valueShape Value
v
  ppr (Int8Value Vector Int
shape Vector Int8
vs) = [Int] -> Vector Int8 -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Int8
vs
  ppr (Int16Value Vector Int
shape Vector Int16
vs) = [Int] -> Vector Int16 -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Int16
vs
  ppr (Int32Value Vector Int
shape Vector Int32
vs) = [Int] -> Vector Int32 -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Int32
vs
  ppr (Int64Value Vector Int
shape Vector Int64
vs) = [Int] -> Vector Int64 -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Int64
vs
  ppr (Word8Value Vector Int
shape Vector Word8
vs) = [Int] -> Vector Word8 -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Word8
vs
  ppr (Word16Value Vector Int
shape Vector Word16
vs) = [Int] -> Vector Word16 -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Word16
vs
  ppr (Word32Value Vector Int
shape Vector Word32
vs) = [Int] -> Vector Word32 -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Word32
vs
  ppr (Word64Value Vector Int
shape Vector Word64
vs) = [Int] -> Vector Word64 -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Word64
vs
  ppr (Float32Value Vector Int
shape Vector Float
vs) = [Int] -> Vector Float -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Float
vs
  ppr (Float64Value Vector Int
shape Vector Double
vs) = [Int] -> Vector Double -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Double
vs
  ppr (BoolValue Vector Int
shape Vector Bool
vs) = [Int] -> Vector Bool -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape) Vector Bool
vs

pprArray :: (UVec.Unbox a, F.IsPrimValue a) => [Int] -> UVec.Vector a -> Doc
pprArray :: [Int] -> Vector a -> Doc
pprArray [] Vector a
vs =
  PrimValue -> Doc
forall a. Pretty a => a -> Doc
ppr (PrimValue -> Doc) -> PrimValue -> Doc
forall a b. (a -> b) -> a -> b
$ a -> PrimValue
forall v. IsPrimValue v => v -> PrimValue
F.primValue (a -> PrimValue) -> a -> PrimValue
forall a b. (a -> b) -> a -> b
$ Vector a -> a
forall a. Unbox a => Vector a -> a
UVec.head Vector a
vs
pprArray (Int
d : [Int]
ds) Vector a
vs =
  Doc -> Doc
brackets (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ [Doc] -> Doc
cat ([Doc] -> Doc) -> [Doc] -> Doc
forall a b. (a -> b) -> a -> b
$ Doc -> [Doc] -> [Doc]
punctuate Doc
separator ([Doc] -> [Doc]) -> [Doc] -> [Doc]
forall a b. (a -> b) -> a -> b
$ (Int -> Doc) -> [Int] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map ([Int] -> Vector a -> Doc
forall a. (Unbox a, IsPrimValue a) => [Int] -> Vector a -> Doc
pprArray [Int]
ds (Vector a -> Doc) -> (Int -> Vector a) -> Int -> Doc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector a
slice) [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
  where
    slice_size :: Int
slice_size = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
ds
    slice :: Int -> Vector a
slice Int
i = Int -> Int -> Vector a -> Vector a
forall a. Unbox a => Int -> Int -> Vector a -> Vector a
UVec.slice (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
slice_size) Int
slice_size Vector a
vs
    separator :: Doc
separator
      | [Int] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
ds = Doc
comma Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
space
      | Bool
otherwise = Doc
comma Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
line

-- | A representation of the simple values we represent in this module.
data ValueType = ValueType [Int] F.PrimType
  deriving (Int -> ValueType -> ShowS
[ValueType] -> ShowS
ValueType -> String
(Int -> ValueType -> ShowS)
-> (ValueType -> String)
-> ([ValueType] -> ShowS)
-> Show ValueType
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ValueType] -> ShowS
$cshowList :: [ValueType] -> ShowS
show :: ValueType -> String
$cshow :: ValueType -> String
showsPrec :: Int -> ValueType -> ShowS
$cshowsPrec :: Int -> ValueType -> ShowS
Show)

instance PP.Pretty ValueType where
  ppr :: ValueType -> Doc
ppr (ValueType [Int]
ds PrimType
t) = [Doc] -> Doc
forall a. Monoid a => [a] -> a
mconcat ((Int -> Doc) -> [Int] -> [Doc]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Doc
forall a. Pretty a => a -> Doc
pprDim [Int]
ds) Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> PrimType -> Doc
forall a. Pretty a => a -> Doc
ppr PrimType
t
    where
      pprDim :: a -> Doc
pprDim a
d = Doc -> Doc
brackets (Doc -> Doc) -> Doc -> Doc
forall a b. (a -> b) -> a -> b
$ a -> Doc
forall a. Pretty a => a -> Doc
ppr a
d

-- | A textual description of the type of a value.  Follows Futhark
-- type notation, and contains the exact dimension sizes if an array.
valueType :: Value -> ValueType
valueType :: Value -> ValueType
valueType Value
v = [Int] -> PrimType -> ValueType
ValueType (Value -> [Int]
valueShape Value
v) (PrimType -> ValueType) -> PrimType -> ValueType
forall a b. (a -> b) -> a -> b
$ Value -> PrimType
valueElemType Value
v

valueElemType :: Value -> F.PrimType
valueElemType :: Value -> PrimType
valueElemType Int8Value {} = IntType -> PrimType
F.Signed IntType
F.Int8
valueElemType Int16Value {} = IntType -> PrimType
F.Signed IntType
F.Int16
valueElemType Int32Value {} = IntType -> PrimType
F.Signed IntType
F.Int32
valueElemType Int64Value {} = IntType -> PrimType
F.Signed IntType
F.Int64
valueElemType Word8Value {} = IntType -> PrimType
F.Unsigned IntType
F.Int8
valueElemType Word16Value {} = IntType -> PrimType
F.Unsigned IntType
F.Int16
valueElemType Word32Value {} = IntType -> PrimType
F.Unsigned IntType
F.Int32
valueElemType Word64Value {} = IntType -> PrimType
F.Unsigned IntType
F.Int64
valueElemType Float32Value {} = FloatType -> PrimType
F.FloatType FloatType
F.Float32
valueElemType Float64Value {} = FloatType -> PrimType
F.FloatType FloatType
F.Float64
valueElemType BoolValue {} = PrimType
F.Bool

valueShape :: Value -> [Int]
valueShape :: Value -> [Int]
valueShape (Int8Value Vector Int
shape Vector Int8
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Int16Value Vector Int
shape Vector Int16
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Int32Value Vector Int
shape Vector Int32
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Int64Value Vector Int
shape Vector Int64
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Word8Value Vector Int
shape Vector Word8
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Word16Value Vector Int
shape Vector Word16
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Word32Value Vector Int
shape Vector Word32
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Word64Value Vector Int
shape Vector Word64
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Float32Value Vector Int
shape Vector Float
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (Float64Value Vector Int
shape Vector Double
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape
valueShape (BoolValue Vector Int
shape Vector Bool
_) = Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector Int
shape

-- The parser

dropRestOfLine, dropSpaces :: BS.ByteString -> BS.ByteString
dropRestOfLine :: ByteString -> ByteString
dropRestOfLine = Int64 -> ByteString -> ByteString
BS.drop Int64
1 (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ByteString -> ByteString
BS.dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
'\n')
dropSpaces :: ByteString -> ByteString
dropSpaces ByteString
t = case (Char -> Bool) -> ByteString -> ByteString
BS.dropWhile Char -> Bool
isSpace ByteString
t of
  ByteString
t'
    | ByteString
"--" ByteString -> ByteString -> Bool
`BS.isPrefixOf` ByteString
t' -> ByteString -> ByteString
dropSpaces (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
dropRestOfLine ByteString
t'
    | Bool
otherwise -> ByteString
t'

type ReadValue v = BS.ByteString -> Maybe (v, BS.ByteString)

symbol :: Char -> BS.ByteString -> Maybe BS.ByteString
symbol :: Char -> ByteString -> Maybe ByteString
symbol Char
c ByteString
t
  | Just (Char
c', ByteString
t') <- ByteString -> Maybe (Char, ByteString)
BS.uncons ByteString
t, Char
c' Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
c = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString) -> ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
dropSpaces ByteString
t'
  | Bool
otherwise = Maybe ByteString
forall a. Maybe a
Nothing

lexeme :: BS.ByteString -> BS.ByteString -> Maybe BS.ByteString
lexeme :: ByteString -> ByteString -> Maybe ByteString
lexeme ByteString
l ByteString
t
  | ByteString
l ByteString -> ByteString -> Bool
`BS.isPrefixOf` ByteString
t = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString) -> ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
dropSpaces (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Int64 -> ByteString -> ByteString
BS.drop (ByteString -> Int64
BS.length ByteString
l) ByteString
t
  | Bool
otherwise = Maybe ByteString
forall a. Maybe a
Nothing

-- (Used elements, shape, elements, remaining input)
type State s v = (Int, Vector Int, STVector s v, BS.ByteString)

readArrayElemsST ::
  UMVec.Unbox v =>
  Int ->
  Int ->
  ReadValue v ->
  State s v ->
  ST s (Maybe (Int, State s v))
readArrayElemsST :: Int
-> Int -> ReadValue v -> State s v -> ST s (Maybe (Int, State s v))
readArrayElemsST Int
j Int
r ReadValue v
rv State s v
s = do
  Maybe (State s v)
ms <- Int -> ReadValue v -> State s v -> ST s (Maybe (State s v))
forall v s.
Unbox v =>
Int -> ReadValue v -> State s v -> ST s (Maybe (State s v))
readRankedArrayOfST Int
r ReadValue v
rv State s v
s
  case Maybe (State s v)
ms of
    Just (Int
i, Vector Int
shape, STVector s v
arr, ByteString
t)
      | Just ByteString
t' <- Char -> ByteString -> Maybe ByteString
symbol Char
',' ByteString
t -> do
        Maybe (Int, State s v)
next <- Int
-> Int -> ReadValue v -> State s v -> ST s (Maybe (Int, State s v))
forall v s.
Unbox v =>
Int
-> Int -> ReadValue v -> State s v -> ST s (Maybe (Int, State s v))
readArrayElemsST (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
r ReadValue v
rv (Int
i, Vector Int
shape, STVector s v
arr, ByteString
t')
        -- Not OK to have zero values after a comma.
        case Maybe (Int, State s v)
next of
          Just (Int
0, State s v
_) -> Maybe (Int, State s v) -> ST s (Maybe (Int, State s v))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Int, State s v)
forall a. Maybe a
Nothing
          Maybe (Int, State s v)
_ -> Maybe (Int, State s v) -> ST s (Maybe (Int, State s v))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Int, State s v)
next
      | Bool
otherwise -> Maybe (Int, State s v) -> ST s (Maybe (Int, State s v))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Int, State s v) -> ST s (Maybe (Int, State s v)))
-> Maybe (Int, State s v) -> ST s (Maybe (Int, State s v))
forall a b. (a -> b) -> a -> b
$ (Int, State s v) -> Maybe (Int, State s v)
forall a. a -> Maybe a
Just (Int
j, (Int
i, Vector Int
shape, STVector s v
arr, ByteString
t))
    Maybe (State s v)
_ ->
      Maybe (Int, State s v) -> ST s (Maybe (Int, State s v))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Int, State s v) -> ST s (Maybe (Int, State s v)))
-> Maybe (Int, State s v) -> ST s (Maybe (Int, State s v))
forall a b. (a -> b) -> a -> b
$ (Int, State s v) -> Maybe (Int, State s v)
forall a. a -> Maybe a
Just (Int
0, State s v
s)

updateShape :: Int -> Int -> Vector Int -> Maybe (Vector Int)
updateShape :: Int -> Int -> Vector Int -> Maybe (Vector Int)
updateShape Int
d Int
n Vector Int
shape
  | Int
old_n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Vector Int -> Maybe (Vector Int)
forall a. a -> Maybe a
Just (Vector Int -> Maybe (Vector Int))
-> Vector Int -> Maybe (Vector Int)
forall a b. (a -> b) -> a -> b
$ Vector Int
shape Vector Int -> [(Int, Int)] -> Vector Int
forall a. Unbox a => Vector a -> [(Int, a)] -> Vector a
UVec.// [(Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
d, Int
n)]
  | Int
old_n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = Vector Int -> Maybe (Vector Int)
forall a. a -> Maybe a
Just Vector Int
shape
  | Bool
otherwise = Maybe (Vector Int)
forall a. Maybe a
Nothing
  where
    r :: Int
r = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
UVec.length Vector Int
shape
    old_n :: Int
old_n = Vector Int
shape Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
UVec.! (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
d)

growIfFilled :: UVec.Unbox v => Int -> STVector s v -> ST s (STVector s v)
growIfFilled :: Int -> STVector s v -> ST s (STVector s v)
growIfFilled Int
i STVector s v
arr =
  if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
capacity
    then MVector (PrimState (ST s)) v
-> Int -> ST s (MVector (PrimState (ST s)) v)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
UMVec.grow STVector s v
MVector (PrimState (ST s)) v
arr Int
capacity
    else STVector s v -> ST s (STVector s v)
forall (m :: * -> *) a. Monad m => a -> m a
return STVector s v
arr
  where
    capacity :: Int
capacity = STVector s v -> Int
forall a s. Unbox a => MVector s a -> Int
UMVec.length STVector s v
arr

readRankedArrayOfST ::
  UMVec.Unbox v =>
  Int ->
  ReadValue v ->
  State s v ->
  ST s (Maybe (State s v))
readRankedArrayOfST :: Int -> ReadValue v -> State s v -> ST s (Maybe (State s v))
readRankedArrayOfST Int
0 ReadValue v
rv (Int
i, Vector Int
shape, STVector s v
arr, ByteString
t)
  | Just (v
v, ByteString
t') <- ReadValue v
rv ByteString
t = do
    STVector s v
arr' <- Int -> STVector s v -> ST s (STVector s v)
forall v s. Unbox v => Int -> STVector s v -> ST s (STVector s v)
growIfFilled Int
i STVector s v
arr
    MVector (PrimState (ST s)) v -> Int -> v -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
UMVec.write STVector s v
MVector (PrimState (ST s)) v
arr' Int
i v
v
    Maybe (State s v) -> ST s (Maybe (State s v))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (State s v) -> ST s (Maybe (State s v)))
-> Maybe (State s v) -> ST s (Maybe (State s v))
forall a b. (a -> b) -> a -> b
$ State s v -> Maybe (State s v)
forall a. a -> Maybe a
Just (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Vector Int
shape, STVector s v
arr', ByteString
t')
readRankedArrayOfST Int
r ReadValue v
rv (Int
i, Vector Int
shape, STVector s v
arr, ByteString
t)
  | Just ByteString
t' <- Char -> ByteString -> Maybe ByteString
symbol Char
'[' ByteString
t = do
    Maybe (Int, State s v)
ms <- Int
-> Int -> ReadValue v -> State s v -> ST s (Maybe (Int, State s v))
forall v s.
Unbox v =>
Int
-> Int -> ReadValue v -> State s v -> ST s (Maybe (Int, State s v))
readArrayElemsST Int
1 (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ReadValue v
rv (Int
i, Vector Int
shape, STVector s v
arr, ByteString
t')
    Maybe (State s v) -> ST s (Maybe (State s v))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (State s v) -> ST s (Maybe (State s v)))
-> Maybe (State s v) -> ST s (Maybe (State s v))
forall a b. (a -> b) -> a -> b
$ do
      (Int
j, State s v
s) <- Maybe (Int, State s v)
ms
      Int -> Int -> State s v -> Maybe (State s v)
forall s v. Int -> Int -> State s v -> Maybe (State s v)
closeArray Int
r Int
j State s v
s
readRankedArrayOfST Int
_ ReadValue v
_ State s v
_ =
  Maybe (State s v) -> ST s (Maybe (State s v))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (State s v)
forall a. Maybe a
Nothing

closeArray :: Int -> Int -> State s v -> Maybe (State s v)
closeArray :: Int -> Int -> State s v -> Maybe (State s v)
closeArray Int
r Int
j (Int
i, Vector Int
shape, STVector s v
arr, ByteString
t) = do
  ByteString
t' <- Char -> ByteString -> Maybe ByteString
symbol Char
']' ByteString
t
  Vector Int
shape' <- Int -> Int -> Vector Int -> Maybe (Vector Int)
updateShape Int
r Int
j Vector Int
shape
  State s v -> Maybe (State s v)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
i, Vector Int
shape', STVector s v
arr, ByteString
t')

readRankedArrayOf ::
  UMVec.Unbox v =>
  Int ->
  ReadValue v ->
  BS.ByteString ->
  Maybe (Vector Int, Vector v, BS.ByteString)
readRankedArrayOf :: Int
-> ReadValue v
-> ByteString
-> Maybe (Vector Int, Vector v, ByteString)
readRankedArrayOf Int
r ReadValue v
rv ByteString
t = (forall s. ST s (Maybe (Vector Int, Vector v, ByteString)))
-> Maybe (Vector Int, Vector v, ByteString)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe (Vector Int, Vector v, ByteString)))
 -> Maybe (Vector Int, Vector v, ByteString))
-> (forall s. ST s (Maybe (Vector Int, Vector v, ByteString)))
-> Maybe (Vector Int, Vector v, ByteString)
forall a b. (a -> b) -> a -> b
$ do
  STVector s v
arr <- Int -> ST s (MVector (PrimState (ST s)) v)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
UMVec.new Int
1024
  Maybe (State s v)
ms <- Int -> ReadValue v -> State s v -> ST s (Maybe (State s v))
forall v s.
Unbox v =>
Int -> ReadValue v -> State s v -> ST s (Maybe (State s v))
readRankedArrayOfST Int
r ReadValue v
rv (Int
0, Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
UVec.replicate Int
r (-Int
1), STVector s v
arr, ByteString
t)
  case Maybe (State s v)
ms of
    Just (Int
i, Vector Int
shape, STVector s v
arr', ByteString
t') -> do
      Vector v
arr'' <- 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 (Int -> Int -> STVector s v -> STVector s v
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
UMVec.slice Int
0 Int
i STVector s v
arr')
      Maybe (Vector Int, Vector v, ByteString)
-> ST s (Maybe (Vector Int, Vector v, ByteString))
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe (Vector Int, Vector v, ByteString)
 -> ST s (Maybe (Vector Int, Vector v, ByteString)))
-> Maybe (Vector Int, Vector v, ByteString)
-> ST s (Maybe (Vector Int, Vector v, ByteString))
forall a b. (a -> b) -> a -> b
$ (Vector Int, Vector v, ByteString)
-> Maybe (Vector Int, Vector v, ByteString)
forall a. a -> Maybe a
Just (Vector Int
shape, Vector v
arr'', ByteString
t')
    Maybe (State s v)
Nothing ->
      Maybe (Vector Int, Vector v, ByteString)
-> ST s (Maybe (Vector Int, Vector v, ByteString))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (Vector Int, Vector v, ByteString)
forall a. Maybe a
Nothing

-- | A character that can be part of a value.  This doesn't work for
-- string and character literals.
constituent :: Char -> Bool
constituent :: Char -> Bool
constituent Char
',' = Bool
False
constituent Char
']' = Bool
False
constituent Char
')' = Bool
False
constituent Char
c = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Char -> Bool
isSpace Char
c

readIntegral :: Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral :: (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe int
f ByteString
t = do
  int
v <- case ([L Token], Pos) -> [L Token]
forall a b. (a, b) -> a
fst (([L Token], Pos) -> [L Token])
-> Either String ([L Token], Pos) -> Either String [L Token]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pos -> ByteString -> Either String ([L Token], Pos)
scanTokens (String -> Int -> Int -> Int -> Pos
Pos String
"" Int
1 Int
1 Int
0) ByteString
a of
    Right [L SrcLoc
_ Token
NEGATE, L SrcLoc
_ (INTLIT Integer
x)] -> int -> Maybe int
forall a. a -> Maybe a
Just (int -> Maybe int) -> int -> Maybe int
forall a b. (a -> b) -> a -> b
$ int -> int
forall a. Num a => a -> a
negate (int -> int) -> int -> int
forall a b. (a -> b) -> a -> b
$ Integer -> int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
x
    Right [L SrcLoc
_ (INTLIT Integer
x)] -> int -> Maybe int
forall a. a -> Maybe a
Just (int -> Maybe int) -> int -> Maybe int
forall a b. (a -> b) -> a -> b
$ Integer -> int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
x
    Right [L SrcLoc
_ Token
tok] -> Token -> Maybe int
f Token
tok
    Right [L SrcLoc
_ Token
NEGATE, L SrcLoc
_ Token
tok] -> int -> int
forall a. Num a => a -> a
negate (int -> int) -> Maybe int -> Maybe int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Token -> Maybe int
f Token
tok
    Either String [L Token]
_ -> Maybe int
forall a. Maybe a
Nothing
  (int, ByteString) -> Maybe (int, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (int
v, ByteString -> ByteString
dropSpaces ByteString
b)
  where
    (ByteString
a, ByteString
b) = (Char -> Bool) -> ByteString -> (ByteString, ByteString)
BS.span Char -> Bool
constituent ByteString
t

readInt8 :: ReadValue Int8
readInt8 :: ReadValue Int8
readInt8 = (Token -> Maybe Int8) -> ReadValue Int8
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe Int8
f
  where
    f :: Token -> Maybe Int8
f (I8LIT Int8
x) = Int8 -> Maybe Int8
forall a. a -> Maybe a
Just Int8
x
    f Token
_ = Maybe Int8
forall a. Maybe a
Nothing

readInt16 :: ReadValue Int16
readInt16 :: ReadValue Int16
readInt16 = (Token -> Maybe Int16) -> ReadValue Int16
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe Int16
f
  where
    f :: Token -> Maybe Int16
f (I16LIT Int16
x) = Int16 -> Maybe Int16
forall a. a -> Maybe a
Just Int16
x
    f Token
_ = Maybe Int16
forall a. Maybe a
Nothing

readInt32 :: ReadValue Int32
readInt32 :: ReadValue Int32
readInt32 = (Token -> Maybe Int32) -> ReadValue Int32
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe Int32
f
  where
    f :: Token -> Maybe Int32
f (I32LIT Int32
x) = Int32 -> Maybe Int32
forall a. a -> Maybe a
Just Int32
x
    f Token
_ = Maybe Int32
forall a. Maybe a
Nothing

readInt64 :: ReadValue Int64
readInt64 :: ReadValue Int64
readInt64 = (Token -> Maybe Int64) -> ReadValue Int64
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe Int64
f
  where
    f :: Token -> Maybe Int64
f (I64LIT Int64
x) = Int64 -> Maybe Int64
forall a. a -> Maybe a
Just Int64
x
    f Token
_ = Maybe Int64
forall a. Maybe a
Nothing

readWord8 :: ReadValue Word8
readWord8 :: ReadValue Word8
readWord8 = (Token -> Maybe Word8) -> ReadValue Word8
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe Word8
f
  where
    f :: Token -> Maybe Word8
f (U8LIT Word8
x) = Word8 -> Maybe Word8
forall a. a -> Maybe a
Just Word8
x
    f Token
_ = Maybe Word8
forall a. Maybe a
Nothing

readWord16 :: ReadValue Word16
readWord16 :: ReadValue Word16
readWord16 = (Token -> Maybe Word16) -> ReadValue Word16
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe Word16
f
  where
    f :: Token -> Maybe Word16
f (U16LIT Word16
x) = Word16 -> Maybe Word16
forall a. a -> Maybe a
Just Word16
x
    f Token
_ = Maybe Word16
forall a. Maybe a
Nothing

readWord32 :: ReadValue Word32
readWord32 :: ReadValue Word32
readWord32 = (Token -> Maybe Word32) -> ReadValue Word32
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe Word32
f
  where
    f :: Token -> Maybe Word32
f (U32LIT Word32
x) = Word32 -> Maybe Word32
forall a. a -> Maybe a
Just Word32
x
    f Token
_ = Maybe Word32
forall a. Maybe a
Nothing

readWord64 :: ReadValue Word64
readWord64 :: ReadValue Word64
readWord64 = (Token -> Maybe Word64) -> ReadValue Word64
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral Token -> Maybe Word64
f
  where
    f :: Token -> Maybe Word64
f (U64LIT Word64
x) = Word64 -> Maybe Word64
forall a. a -> Maybe a
Just Word64
x
    f Token
_ = Maybe Word64
forall a. Maybe a
Nothing

readFloat :: RealFloat float => ([Token] -> Maybe float) -> ReadValue float
readFloat :: ([Token] -> Maybe float) -> ReadValue float
readFloat [Token] -> Maybe float
f ByteString
t = do
  float
v <- case (L Token -> Token) -> [L Token] -> [Token]
forall a b. (a -> b) -> [a] -> [b]
map L Token -> Token
forall a. L a -> a
unLoc ([L Token] -> [Token])
-> (([L Token], Pos) -> [L Token]) -> ([L Token], Pos) -> [Token]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([L Token], Pos) -> [L Token]
forall a b. (a, b) -> a
fst (([L Token], Pos) -> [Token])
-> Either String ([L Token], Pos) -> Either String [Token]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pos -> ByteString -> Either String ([L Token], Pos)
scanTokens (String -> Int -> Int -> Int -> Pos
Pos String
"" Int
1 Int
1 Int
0) ByteString
a of
    Right [Token
NEGATE, FLOATLIT Double
x] -> float -> Maybe float
forall a. a -> Maybe a
Just (float -> Maybe float) -> float -> Maybe float
forall a b. (a -> b) -> a -> b
$ float -> float
forall a. Num a => a -> a
negate (float -> float) -> float -> float
forall a b. (a -> b) -> a -> b
$ Double -> float
fromDouble Double
x
    Right [FLOATLIT Double
x] -> float -> Maybe float
forall a. a -> Maybe a
Just (float -> Maybe float) -> float -> Maybe float
forall a b. (a -> b) -> a -> b
$ Double -> float
fromDouble Double
x
    Right (Token
NEGATE : [Token]
toks) -> float -> float
forall a. Num a => a -> a
negate (float -> float) -> Maybe float -> Maybe float
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Token] -> Maybe float
f [Token]
toks
    Right [Token]
toks -> [Token] -> Maybe float
f [Token]
toks
    Either String [Token]
_ -> Maybe float
forall a. Maybe a
Nothing
  (float, ByteString) -> Maybe (float, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (float
v, ByteString -> ByteString
dropSpaces ByteString
b)
  where
    (ByteString
a, ByteString
b) = (Char -> Bool) -> ByteString -> (ByteString, ByteString)
BS.span Char -> Bool
constituent ByteString
t
    fromDouble :: Double -> float
fromDouble = (Integer -> Int -> float) -> (Integer, Int) -> float
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Integer -> Int -> float
forall a. RealFloat a => Integer -> Int -> a
encodeFloat ((Integer, Int) -> float)
-> (Double -> (Integer, Int)) -> Double -> float
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> (Integer, Int)
forall a. RealFloat a => a -> (Integer, Int)
decodeFloat
    unLoc :: L a -> a
unLoc (L SrcLoc
_ a
x) = a
x

readFloat32 :: ReadValue Float
readFloat32 :: ReadValue Float
readFloat32 = ([Token] -> Maybe Float) -> ReadValue Float
forall float.
RealFloat float =>
([Token] -> Maybe float) -> ReadValue float
readFloat [Token] -> Maybe Float
lexFloat32
  where
    lexFloat32 :: [Token] -> Maybe Float
lexFloat32 [F32LIT Float
x] = Float -> Maybe Float
forall a. a -> Maybe a
Just Float
x
    lexFloat32 [ID Name
"f32", PROJ_FIELD Name
"inf"] = Float -> Maybe Float
forall a. a -> Maybe a
Just (Float -> Maybe Float) -> Float -> Maybe Float
forall a b. (a -> b) -> a -> b
$ Float
1 Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
0
    lexFloat32 [ID Name
"f32", PROJ_FIELD Name
"nan"] = Float -> Maybe Float
forall a. a -> Maybe a
Just (Float -> Maybe Float) -> Float -> Maybe Float
forall a b. (a -> b) -> a -> b
$ Float
0 Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
0
    lexFloat32 [Token]
_ = Maybe Float
forall a. Maybe a
Nothing

readFloat64 :: ReadValue Double
readFloat64 :: ReadValue Double
readFloat64 = ([Token] -> Maybe Double) -> ReadValue Double
forall float.
RealFloat float =>
([Token] -> Maybe float) -> ReadValue float
readFloat [Token] -> Maybe Double
lexFloat64
  where
    lexFloat64 :: [Token] -> Maybe Double
lexFloat64 [F64LIT Double
x] = Double -> Maybe Double
forall a. a -> Maybe a
Just Double
x
    lexFloat64 [ID Name
"f64", PROJ_FIELD Name
"inf"] = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$ Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0
    lexFloat64 [ID Name
"f64", PROJ_FIELD Name
"nan"] = Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$ Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0
    lexFloat64 [Token]
_ = Maybe Double
forall a. Maybe a
Nothing

readBool :: ReadValue Bool
readBool :: ReadValue Bool
readBool ByteString
t = do
  Bool
v <- case ([L Token], Pos) -> [L Token]
forall a b. (a, b) -> a
fst (([L Token], Pos) -> [L Token])
-> Either String ([L Token], Pos) -> Either String [L Token]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pos -> ByteString -> Either String ([L Token], Pos)
scanTokens (String -> Int -> Int -> Int -> Pos
Pos String
"" Int
1 Int
1 Int
0) ByteString
a of
    Right [L SrcLoc
_ Token
TRUE] -> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
True
    Right [L SrcLoc
_ Token
FALSE] -> Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
False
    Either String [L Token]
_ -> Maybe Bool
forall a. Maybe a
Nothing
  (Bool, ByteString) -> Maybe (Bool, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool
v, ByteString -> ByteString
dropSpaces ByteString
b)
  where
    (ByteString
a, ByteString
b) = (Char -> Bool) -> ByteString -> (ByteString, ByteString)
BS.span Char -> Bool
constituent ByteString
t

readPrimType :: ReadValue String
readPrimType :: ReadValue String
readPrimType ByteString
t = do
  String
pt <- case ([L Token], Pos) -> [L Token]
forall a b. (a, b) -> a
fst (([L Token], Pos) -> [L Token])
-> Either String ([L Token], Pos) -> Either String [L Token]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pos -> ByteString -> Either String ([L Token], Pos)
scanTokens (String -> Int -> Int -> Int -> Pos
Pos String
"" Int
1 Int
1 Int
0) ByteString
a of
    Right [L SrcLoc
_ (ID Name
s)] -> String -> Maybe String
forall a. a -> Maybe a
Just (String -> Maybe String) -> String -> Maybe String
forall a b. (a -> b) -> a -> b
$ Name -> String
F.nameToString Name
s
    Either String [L Token]
_ -> Maybe String
forall a. Maybe a
Nothing
  (String, ByteString) -> Maybe (String, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (String
pt, ByteString -> ByteString
dropSpaces ByteString
b)
  where
    (ByteString
a, ByteString
b) = (Char -> Bool) -> ByteString -> (ByteString, ByteString)
BS.span Char -> Bool
constituent ByteString
t

readEmptyArrayOfShape :: [Int] -> BS.ByteString -> Maybe (Value, BS.ByteString)
readEmptyArrayOfShape :: [Int] -> ByteString -> Maybe (Value, ByteString)
readEmptyArrayOfShape [Int]
shape ByteString
t
  | Just ByteString
t' <- Char -> ByteString -> Maybe ByteString
symbol Char
'[' ByteString
t,
    Just (Int
d, ByteString
t'') <- (Token -> Maybe Int) -> ReadValue Int
forall int. Integral int => (Token -> Maybe int) -> ReadValue int
readIntegral (Maybe Int -> Token -> Maybe Int
forall a b. a -> b -> a
const Maybe Int
forall a. Maybe a
Nothing) ByteString
t',
    Just ByteString
t''' <- Char -> ByteString -> Maybe ByteString
symbol Char
']' ByteString
t'' =
    [Int] -> ByteString -> Maybe (Value, ByteString)
readEmptyArrayOfShape ([Int]
shape [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
d]) ByteString
t'''
  | Bool
otherwise = do
    (String
pt, ByteString
t') <- ReadValue String
readPrimType ByteString
t
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem Int
0 [Int]
shape
    Value
v <- case String
pt of
      String
"i8" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int8 -> Value
Int8Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Int8
forall a. Unbox a => Vector a
UVec.empty
      String
"i16" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int16 -> Value
Int16Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Int16
forall a. Unbox a => Vector a
UVec.empty
      String
"i32" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int32 -> Value
Int32Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Int32
forall a. Unbox a => Vector a
UVec.empty
      String
"i64" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int64 -> Value
Int64Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Int64
forall a. Unbox a => Vector a
UVec.empty
      String
"u8" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word8 -> Value
Word8Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Word8
forall a. Unbox a => Vector a
UVec.empty
      String
"u16" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word16 -> Value
Word16Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Word16
forall a. Unbox a => Vector a
UVec.empty
      String
"u32" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word32 -> Value
Word32Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Word32
forall a. Unbox a => Vector a
UVec.empty
      String
"u64" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word64 -> Value
Word64Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Word64
forall a. Unbox a => Vector a
UVec.empty
      String
"f32" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Float -> Value
Float32Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Float
forall a. Unbox a => Vector a
UVec.empty
      String
"f64" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Double -> Value
Float64Value ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Double
forall a. Unbox a => Vector a
UVec.empty
      String
"bool" -> Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Bool -> Value
BoolValue ([Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UVec.fromList [Int]
shape) Vector Bool
forall a. Unbox a => Vector a
UVec.empty
      String
_ -> Maybe Value
forall a. Maybe a
Nothing
    (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Value
v, ByteString
t')

readEmptyArray :: BS.ByteString -> Maybe (Value, BS.ByteString)
readEmptyArray :: ByteString -> Maybe (Value, ByteString)
readEmptyArray ByteString
t = do
  ByteString
t' <- Char -> ByteString -> Maybe ByteString
symbol Char
'(' (ByteString -> Maybe ByteString)
-> Maybe ByteString -> Maybe ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ByteString -> ByteString -> Maybe ByteString
lexeme ByteString
"empty" ByteString
t
  (Value
v, ByteString
t'') <- [Int] -> ByteString -> Maybe (Value, ByteString)
readEmptyArrayOfShape [] ByteString
t'
  ByteString
t''' <- Char -> ByteString -> Maybe ByteString
symbol Char
')' ByteString
t''
  (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Value
v, ByteString
t''')

readValue :: BS.ByteString -> Maybe (Value, BS.ByteString)
readValue :: ByteString -> Maybe (Value, ByteString)
readValue ByteString
full_t
  | Right (ByteString
t', Int64
_, Value
v) <- ByteString
-> Either (ByteString, Int64, String) (ByteString, Int64, Value)
forall a.
Binary a =>
ByteString
-> Either (ByteString, Int64, String) (ByteString, Int64, a)
decodeOrFail ByteString
full_t =
    (Value, ByteString) -> Maybe (Value, ByteString)
forall a. a -> Maybe a
Just (Value
v, ByteString -> ByteString
dropSpaces ByteString
t')
  | Bool
otherwise = ByteString -> Maybe (Value, ByteString)
readEmptyArray ByteString
full_t Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` Int -> ByteString -> Maybe (Value, ByteString)
insideBrackets Int
0 ByteString
full_t
  where
    insideBrackets :: Int -> ByteString -> Maybe (Value, ByteString)
insideBrackets Int
r ByteString
t = Maybe (Value, ByteString)
-> (ByteString -> Maybe (Value, ByteString))
-> Maybe ByteString
-> Maybe (Value, ByteString)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int -> ByteString -> Maybe (Value, ByteString)
tryValueAndReadValue Int
r ByteString
t) (Int -> ByteString -> Maybe (Value, ByteString)
insideBrackets (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Maybe ByteString -> Maybe (Value, ByteString))
-> Maybe ByteString -> Maybe (Value, ByteString)
forall a b. (a -> b) -> a -> b
$ Char -> ByteString -> Maybe ByteString
symbol Char
'[' ByteString
t
    tryWith :: (ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ByteString -> Maybe (v, ByteString)
f Vector Int -> Vector v -> a
mk Int
r ByteString
t
      | Just (v, ByteString)
_ <- ByteString -> Maybe (v, ByteString)
f ByteString
t = do
        (Vector Int
shape, Vector v
arr, ByteString
rest_t) <- Int
-> (ByteString -> Maybe (v, ByteString))
-> ByteString
-> Maybe (Vector Int, Vector v, ByteString)
forall v.
Unbox v =>
Int
-> ReadValue v
-> ByteString
-> Maybe (Vector Int, Vector v, ByteString)
readRankedArrayOf Int
r ByteString -> Maybe (v, ByteString)
f ByteString
full_t
        (a, ByteString) -> Maybe (a, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Vector Int -> Vector v -> a
mk Vector Int
shape Vector v
arr, ByteString
rest_t)
      | Bool
otherwise = Maybe (a, ByteString)
forall a. Maybe a
Nothing
    tryValueAndReadValue :: Int -> ByteString -> Maybe (Value, ByteString)
tryValueAndReadValue Int
r ByteString
t =
      -- 32-bit signed integers come first such that we parse
      -- unsuffixed integer constants as of that type.
      ReadValue Int32
-> (Vector Int -> Vector Int32 -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Int32
readInt32 Vector Int -> Vector Int32 -> Value
Int32Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Int8
-> (Vector Int -> Vector Int8 -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Int8
readInt8 Vector Int -> Vector Int8 -> Value
Int8Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Int16
-> (Vector Int -> Vector Int16 -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Int16
readInt16 Vector Int -> Vector Int16 -> Value
Int16Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Int64
-> (Vector Int -> Vector Int64 -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Int64
readInt64 Vector Int -> Vector Int64 -> Value
Int64Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Word8
-> (Vector Int -> Vector Word8 -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Word8
readWord8 Vector Int -> Vector Word8 -> Value
Word8Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Word16
-> (Vector Int -> Vector Word16 -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Word16
readWord16 Vector Int -> Vector Word16 -> Value
Word16Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Word32
-> (Vector Int -> Vector Word32 -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Word32
readWord32 Vector Int -> Vector Word32 -> Value
Word32Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Word64
-> (Vector Int -> Vector Word64 -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Word64
readWord64 Vector Int -> Vector Word64 -> Value
Word64Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Double
-> (Vector Int -> Vector Double -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Double
readFloat64 Vector Int -> Vector Double -> Value
Float64Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Float
-> (Vector Int -> Vector Float -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Float
readFloat32 Vector Int -> Vector Float -> Value
Float32Value Int
r ByteString
t
        Maybe (Value, ByteString)
-> Maybe (Value, ByteString) -> Maybe (Value, ByteString)
forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ReadValue Bool
-> (Vector Int -> Vector Bool -> Value)
-> Int
-> ByteString
-> Maybe (Value, ByteString)
forall v a.
Unbox v =>
(ByteString -> Maybe (v, ByteString))
-> (Vector Int -> Vector v -> a)
-> Int
-> ByteString
-> Maybe (a, ByteString)
tryWith ReadValue Bool
readBool Vector Int -> Vector Bool -> Value
BoolValue Int
r ByteString
t

-- | Parse Futhark values from the given bytestring.
readValues :: BS.ByteString -> Maybe [Value]
readValues :: ByteString -> Maybe [Value]
readValues = ByteString -> Maybe [Value]
readValues' (ByteString -> Maybe [Value])
-> (ByteString -> ByteString) -> ByteString -> Maybe [Value]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
dropSpaces
  where
    readValues' :: ByteString -> Maybe [Value]
readValues' ByteString
t
      | ByteString -> Bool
BS.null ByteString
t = [Value] -> Maybe [Value]
forall a. a -> Maybe a
Just []
      | Bool
otherwise = do
        (Value
a, ByteString
t') <- ByteString -> Maybe (Value, ByteString)
readValue ByteString
t
        (Value
a Value -> [Value] -> [Value]
forall a. a -> [a] -> [a]
:) ([Value] -> [Value]) -> Maybe [Value] -> Maybe [Value]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteString -> Maybe [Value]
readValues' ByteString
t'

-- Comparisons

-- | Two values differ in some way.  The 'Show' instance produces a
-- human-readable explanation.
data Mismatch
  = -- | The position the value number and a flat index
    -- into the array.
    PrimValueMismatch (Int, Int) PrimValue PrimValue
  | ArrayShapeMismatch Int [Int] [Int]
  | TypeMismatch Int String String
  | ValueCountMismatch Int Int

instance Show Mismatch where
  show :: Mismatch -> String
show (PrimValueMismatch (Int
i, Int
j) PrimValue
got PrimValue
expected) =
    (Int, Int) -> String -> PrimValue -> PrimValue -> String
forall i a. (Show i, Pretty a) => i -> String -> a -> a -> String
explainMismatch (Int
i, Int
j) String
"" PrimValue
got PrimValue
expected
  show (ArrayShapeMismatch Int
i [Int]
got [Int]
expected) =
    Int -> String -> [Int] -> [Int] -> String
forall i a. (Show i, Pretty a) => i -> String -> a -> a -> String
explainMismatch Int
i String
"array of shape " [Int]
got [Int]
expected
  show (TypeMismatch Int
i String
got String
expected) =
    Int -> String -> String -> ShowS
forall i a. (Show i, Pretty a) => i -> String -> a -> a -> String
explainMismatch Int
i String
"value of type " String
got String
expected
  show (ValueCountMismatch Int
got Int
expected) =
    String
"Expected " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
expected String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" values, got " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
got

-- | A human-readable description of how two values are not the same.
explainMismatch :: (Show i, PP.Pretty a) => i -> String -> a -> a -> String
explainMismatch :: i -> String -> a -> a -> String
explainMismatch i
i String
what a
got a
expected =
  String
"Value " String -> ShowS
forall a. [a] -> [a] -> [a]
++ i -> String
forall a. Show a => a -> String
show i
i String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" expected " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
what String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Pretty a => a -> String
PP.pretty a
expected String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", got " String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Pretty a => a -> String
PP.pretty a
got

-- | Compare two sets of Futhark values for equality.  Shapes and
-- types must also match.
compareValues :: [Value] -> [Value] -> [Mismatch]
compareValues :: [Value] -> [Value] -> [Mismatch]
compareValues [Value]
got [Value]
expected
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m = [Int -> Int -> Mismatch
ValueCountMismatch Int
n Int
m]
  | Bool
otherwise = [[Mismatch]] -> [Mismatch]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Mismatch]] -> [Mismatch]) -> [[Mismatch]] -> [Mismatch]
forall a b. (a -> b) -> a -> b
$ (Int -> Value -> Value -> [Mismatch])
-> [Int] -> [Value] -> [Value] -> [[Mismatch]]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 Int -> Value -> Value -> [Mismatch]
compareValue [Int
0 ..] [Value]
got [Value]
expected
  where
    n :: Int
n = [Value] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Value]
got
    m :: Int
m = [Value] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Value]
expected

-- | As 'compareValues', but only reports one mismatch.
compareValues1 :: [Value] -> [Value] -> Maybe Mismatch
compareValues1 :: [Value] -> [Value] -> Maybe Mismatch
compareValues1 [Value]
got [Value]
expected = [Mismatch] -> Maybe Mismatch
forall a. [a] -> Maybe a
maybeHead ([Mismatch] -> Maybe Mismatch) -> [Mismatch] -> Maybe Mismatch
forall a b. (a -> b) -> a -> b
$ [Value] -> [Value] -> [Mismatch]
compareValues [Value]
got [Value]
expected

compareValue :: Int -> Value -> Value -> [Mismatch]
compareValue :: Int -> Value -> Value -> [Mismatch]
compareValue Int
i Value
got_v Value
expected_v
  | Value -> [Int]
valueShape Value
got_v [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== Value -> [Int]
valueShape Value
expected_v =
    case (Value
got_v, Value
expected_v) of
      (Int8Value Vector Int
_ Vector Int8
got_vs, Int8Value Vector Int
_ Vector Int8
expected_vs) ->
        Int8 -> Vector Int8 -> Vector Int8 -> [Mismatch]
forall a.
(Unbox a, Ord a, Num a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareNum Int8
1 Vector Int8
got_vs Vector Int8
expected_vs
      (Int16Value Vector Int
_ Vector Int16
got_vs, Int16Value Vector Int
_ Vector Int16
expected_vs) ->
        Int16 -> Vector Int16 -> Vector Int16 -> [Mismatch]
forall a.
(Unbox a, Ord a, Num a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareNum Int16
1 Vector Int16
got_vs Vector Int16
expected_vs
      (Int32Value Vector Int
_ Vector Int32
got_vs, Int32Value Vector Int
_ Vector Int32
expected_vs) ->
        Int32 -> Vector Int32 -> Vector Int32 -> [Mismatch]
forall a.
(Unbox a, Ord a, Num a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareNum Int32
1 Vector Int32
got_vs Vector Int32
expected_vs
      (Int64Value Vector Int
_ Vector Int64
got_vs, Int64Value Vector Int
_ Vector Int64
expected_vs) ->
        Int64 -> Vector Int64 -> Vector Int64 -> [Mismatch]
forall a.
(Unbox a, Ord a, Num a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareNum Int64
1 Vector Int64
got_vs Vector Int64
expected_vs
      (Word8Value Vector Int
_ Vector Word8
got_vs, Word8Value Vector Int
_ Vector Word8
expected_vs) ->
        Word8 -> Vector Word8 -> Vector Word8 -> [Mismatch]
forall a.
(Unbox a, Ord a, Num a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareNum Word8
1 Vector Word8
got_vs Vector Word8
expected_vs
      (Word16Value Vector Int
_ Vector Word16
got_vs, Word16Value Vector Int
_ Vector Word16
expected_vs) ->
        Word16 -> Vector Word16 -> Vector Word16 -> [Mismatch]
forall a.
(Unbox a, Ord a, Num a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareNum Word16
1 Vector Word16
got_vs Vector Word16
expected_vs
      (Word32Value Vector Int
_ Vector Word32
got_vs, Word32Value Vector Int
_ Vector Word32
expected_vs) ->
        Word32 -> Vector Word32 -> Vector Word32 -> [Mismatch]
forall a.
(Unbox a, Ord a, Num a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareNum Word32
1 Vector Word32
got_vs Vector Word32
expected_vs
      (Word64Value Vector Int
_ Vector Word64
got_vs, Word64Value Vector Int
_ Vector Word64
expected_vs) ->
        Word64 -> Vector Word64 -> Vector Word64 -> [Mismatch]
forall a.
(Unbox a, Ord a, Num a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareNum Word64
1 Vector Word64
got_vs Vector Word64
expected_vs
      (Float32Value Vector Int
_ Vector Float
got_vs, Float32Value Vector Int
_ Vector Float
expected_vs) ->
        Float -> Vector Float -> Vector Float -> [Mismatch]
forall a.
(Unbox a, RealFloat a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareFloat (Vector Float -> Float
forall a. (RealFloat a, Unbox a) => Vector a -> a
tolerance Vector Float
expected_vs) Vector Float
got_vs Vector Float
expected_vs
      (Float64Value Vector Int
_ Vector Double
got_vs, Float64Value Vector Int
_ Vector Double
expected_vs) ->
        Double -> Vector Double -> Vector Double -> [Mismatch]
forall a.
(Unbox a, RealFloat a, IsValue a) =>
a -> Vector a -> Vector a -> [Mismatch]
compareFloat (Vector Double -> Double
forall a. (RealFloat a, Unbox a) => Vector a -> a
tolerance Vector Double
expected_vs) Vector Double
got_vs Vector Double
expected_vs
      (BoolValue Vector Int
_ Vector Bool
got_vs, BoolValue Vector Int
_ Vector Bool
expected_vs) ->
        ((Int, Bool) -> Bool -> [Mismatch])
-> Vector Bool -> Vector Bool -> [Mismatch]
forall a b a.
(Unbox a, Unbox b) =>
((Int, a) -> b -> [a]) -> Vector a -> Vector b -> [a]
compareGen (Int, Bool) -> Bool -> [Mismatch]
forall a. (Eq a, IsValue a) => (Int, a) -> a -> [Mismatch]
compareBool Vector Bool
got_vs Vector Bool
expected_vs
      (Value, Value)
_ ->
        [Int -> String -> String -> Mismatch
TypeMismatch Int
i (PrimType -> String
forall a. Pretty a => a -> String
pretty (PrimType -> String) -> PrimType -> String
forall a b. (a -> b) -> a -> b
$ Value -> PrimType
valueElemType Value
got_v) (PrimType -> String
forall a. Pretty a => a -> String
pretty (PrimType -> String) -> PrimType -> String
forall a b. (a -> b) -> a -> b
$ Value -> PrimType
valueElemType Value
expected_v)]
  | Bool
otherwise =
    [Int -> [Int] -> [Int] -> Mismatch
ArrayShapeMismatch Int
i (Value -> [Int]
valueShape Value
got_v) (Value -> [Int]
valueShape Value
expected_v)]
  where
    compareNum :: a -> Vector a -> Vector a -> [Mismatch]
compareNum a
tol = ((Int, a) -> a -> [Mismatch]) -> Vector a -> Vector a -> [Mismatch]
forall a b a.
(Unbox a, Unbox b) =>
((Int, a) -> b -> [a]) -> Vector a -> Vector b -> [a]
compareGen (((Int, a) -> a -> [Mismatch])
 -> Vector a -> Vector a -> [Mismatch])
-> ((Int, a) -> a -> [Mismatch])
-> Vector a
-> Vector a
-> [Mismatch]
forall a b. (a -> b) -> a -> b
$ a -> (Int, a) -> a -> [Mismatch]
forall a.
(Ord a, Num a, IsValue a) =>
a -> (Int, a) -> a -> [Mismatch]
compareElement a
tol
    compareFloat :: a -> Vector a -> Vector a -> [Mismatch]
compareFloat a
tol = ((Int, a) -> a -> [Mismatch]) -> Vector a -> Vector a -> [Mismatch]
forall a b a.
(Unbox a, Unbox b) =>
((Int, a) -> b -> [a]) -> Vector a -> Vector b -> [a]
compareGen (((Int, a) -> a -> [Mismatch])
 -> Vector a -> Vector a -> [Mismatch])
-> ((Int, a) -> a -> [Mismatch])
-> Vector a
-> Vector a
-> [Mismatch]
forall a b. (a -> b) -> a -> b
$ a -> (Int, a) -> a -> [Mismatch]
forall a.
(RealFloat a, IsValue a) =>
a -> (Int, a) -> a -> [Mismatch]
compareFloatElement a
tol

    compareGen :: ((Int, a) -> b -> [a]) -> Vector a -> Vector b -> [a]
compareGen (Int, a) -> b -> [a]
cmp Vector a
got Vector b
expected =
      [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a]) -> [[a]] -> [a]
forall a b. (a -> b) -> a -> b
$
        ((Int, a) -> b -> [a]) -> [(Int, a)] -> [b] -> [[a]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Int, a) -> b -> [a]
cmp (Vector (Int, a) -> [(Int, a)]
forall a. Unbox a => Vector a -> [a]
UVec.toList (Vector (Int, a) -> [(Int, a)]) -> Vector (Int, a) -> [(Int, a)]
forall a b. (a -> b) -> a -> b
$ Vector a -> Vector (Int, a)
forall a. Unbox a => Vector a -> Vector (Int, a)
UVec.indexed Vector a
got) (Vector b -> [b]
forall a. Unbox a => Vector a -> [a]
UVec.toList Vector b
expected)

    compareElement :: a -> (Int, a) -> a -> [Mismatch]
compareElement a
tol (Int
j, a
got) a
expected
      | a -> a -> a -> Bool
forall num. (Ord num, Num num) => num -> num -> num -> Bool
comparePrimValue a
tol a
got a
expected = []
      | Bool
otherwise = [(Int, Int) -> PrimValue -> PrimValue -> Mismatch
PrimValueMismatch (Int
i, Int
j) (a -> PrimValue
forall a. IsValue a => a -> PrimValue
value a
got) (a -> PrimValue
forall a. IsValue a => a -> PrimValue
value a
expected)]

    compareFloatElement :: a -> (Int, a) -> a -> [Mismatch]
compareFloatElement a
tol (Int
j, a
got) a
expected
      | a -> Bool
forall a. RealFloat a => a -> Bool
isNaN a
got, a -> Bool
forall a. RealFloat a => a -> Bool
isNaN a
expected = []
      | a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite a
got,
        a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite a
expected,
        a -> a
forall a. Num a => a -> a
signum a
got a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a -> a
forall a. Num a => a -> a
signum a
expected =
        []
      | Bool
otherwise = a -> (Int, a) -> a -> [Mismatch]
forall a.
(Ord a, Num a, IsValue a) =>
a -> (Int, a) -> a -> [Mismatch]
compareElement a
tol (Int
j, a
got) a
expected

    compareBool :: (Int, a) -> a -> [Mismatch]
compareBool (Int
j, a
got) a
expected
      | a
got a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
expected = []
      | Bool
otherwise = [(Int, Int) -> PrimValue -> PrimValue -> Mismatch
PrimValueMismatch (Int
i, Int
j) (a -> PrimValue
forall a. IsValue a => a -> PrimValue
value a
got) (a -> PrimValue
forall a. IsValue a => a -> PrimValue
value a
expected)]

comparePrimValue ::
  (Ord num, Num num) =>
  num ->
  num ->
  num ->
  Bool
comparePrimValue :: num -> num -> num -> Bool
comparePrimValue num
tol num
x num
y =
  num
diff num -> num -> Bool
forall a. Ord a => a -> a -> Bool
< num
tol
  where
    diff :: num
diff = num -> num
forall a. Num a => a -> a
abs (num -> num) -> num -> num
forall a b. (a -> b) -> a -> b
$ num
x num -> num -> num
forall a. Num a => a -> a -> a
- num
y

minTolerance :: Fractional a => a
minTolerance :: a
minTolerance = a
0.002 -- 0.2%

tolerance :: (RealFloat a, UVec.Unbox a) => Vector a -> a
tolerance :: Vector a -> a
tolerance = (a -> a -> a) -> a -> Vector a -> a
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
UVec.foldl a -> a -> a
forall a. (Ord a, Fractional a) => a -> a -> a
tolerance' a
forall a. Fractional a => a
minTolerance (Vector a -> a) -> (Vector a -> Vector a) -> Vector a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Bool) -> Vector a -> Vector a
forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
UVec.filter (Bool -> Bool
not (Bool -> Bool) -> (a -> Bool) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
forall a. RealFloat a => a -> Bool
nanOrInf)
  where
    tolerance' :: a -> a -> a
tolerance' a
t a
v = a -> a -> a
forall a. Ord a => a -> a -> a
max a
t (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ a
forall a. Fractional a => a
minTolerance a -> a -> a
forall a. Num a => a -> a -> a
* a
v
    nanOrInf :: a -> Bool
nanOrInf a
x = a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite a
x Bool -> Bool -> Bool
|| a -> Bool
forall a. RealFloat a => a -> Bool
isNaN a
x