-- | IO operations exposed through the interpreter.
module Language.Futhark.Interpreter.IO
  ( IOOp (..),
    determineIO,
    doIOOp,
    ioRelativeTo,
  )
where

import Codec.BMP qualified as BMP
import Control.Exception (SomeException, displayException, try)
import Control.Monad
import Control.Monad.IO.Class
import Data.Bits
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as LBS
import Data.Maybe (fromMaybe)
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Data.Text.Read qualified as T
import Data.Vector.Storable qualified as SVec
import Data.Vector.Storable.ByteString qualified as SVec
import Futhark.Test.Values qualified as V
import Futhark.Util (runProgramWithExitCode)
import Language.Futhark
  ( FloatType (..),
    IntType (..),
    PrimType (..),
    RetTypeBase (..),
    ScalarTypeBase (..),
    Shape (..),
    TypeBase (..),
    ValueType,
    isTupleRecord,
    prettyString,
    toStruct,
  )
import Language.Futhark.Interpreter.FFI.ServerM qualified as FFI
import Language.Futhark.Interpreter.Values
import System.Exit
import System.FilePath
import System.IO.Temp (withSystemTempDirectory)

-- | The IO operation to perform. The idea of splitting it out into a type like
-- this is to enforce some kind of auditing.
data IOOp
  = LoadBytes FilePath
  | LoadImg FilePath
  | LoadAudio FilePath
  | -- | Load a value of this type.
    LoadValue ValueType FilePath

load ::
  (m ValueType -> FilePath -> m IOOp) ->
  Maybe (m ValueType -> Value m -> m IOOp)
load :: forall (m :: * -> *).
(m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
load m ValueType -> String -> m IOOp
c = (m ValueType -> Value m -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall a. a -> Maybe a
Just ((m ValueType -> Value m -> m IOOp)
 -> Maybe (m ValueType -> Value m -> m IOOp))
-> (m ValueType -> Value m -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall a b. (a -> b) -> a -> b
$ \m ValueType
t Value m
v -> case Value m -> Maybe ByteString
forall (m :: * -> *). Value m -> Maybe ByteString
asByteString Value m
v of
  Maybe ByteString
Nothing -> String -> m IOOp
forall a. HasCallStack => String -> a
error String
"loadbytes: not a string"
  Just ByteString
v' -> m ValueType -> String -> m IOOp
c m ValueType
t (String -> m IOOp) -> String -> m IOOp
forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack (Text -> String) -> Text -> String
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
T.decodeUtf8 ByteString
v'

primTypeToValueType :: PrimType -> V.PrimType
primTypeToValueType :: PrimType -> PrimType
primTypeToValueType (Signed IntType
Int8) = PrimType
V.I8
primTypeToValueType (Signed IntType
Int16) = PrimType
V.I16
primTypeToValueType (Signed IntType
Int32) = PrimType
V.I32
primTypeToValueType (Signed IntType
Int64) = PrimType
V.I64
primTypeToValueType (Unsigned IntType
Int8) = PrimType
V.U8
primTypeToValueType (Unsigned IntType
Int16) = PrimType
V.U16
primTypeToValueType (Unsigned IntType
Int32) = PrimType
V.U32
primTypeToValueType (Unsigned IntType
Int64) = PrimType
V.U64
primTypeToValueType (FloatType FloatType
Float16) = PrimType
V.F16
primTypeToValueType (FloatType FloatType
Float32) = PrimType
V.F32
primTypeToValueType (FloatType FloatType
Float64) = PrimType
V.F64
primTypeToValueType PrimType
Bool = PrimType
V.Bool

-- | The types of the values in a data file that can be loaded as a value of
-- this type. A tuple corresponds to one value per element.
typeToValueTypes :: ValueType -> Maybe [V.ValueType]
typeToValueTypes :: ValueType -> Maybe [ValueType]
typeToValueTypes ValueType
t = (ValueType -> Maybe ValueType) -> [ValueType] -> Maybe [ValueType]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ValueType -> Maybe ValueType
forall {a} {o}. Integral a => TypeBase a o -> Maybe ValueType
onValue ([ValueType] -> Maybe [ValueType])
-> [ValueType] -> Maybe [ValueType]
forall a b. (a -> b) -> a -> b
$ [ValueType] -> Maybe [ValueType] -> [ValueType]
forall a. a -> Maybe a -> a
fromMaybe [ValueType
t] (Maybe [ValueType] -> [ValueType])
-> Maybe [ValueType] -> [ValueType]
forall a b. (a -> b) -> a -> b
$ ValueType -> Maybe [ValueType]
forall dim o. TypeBase dim o -> Maybe [TypeBase dim o]
isTupleRecord ValueType
t
  where
    onValue :: TypeBase a o -> Maybe ValueType
onValue (Scalar (Prim PrimType
pt)) =
      ValueType -> Maybe ValueType
forall a. a -> Maybe a
Just (ValueType -> Maybe ValueType) -> ValueType -> Maybe ValueType
forall a b. (a -> b) -> a -> b
$ [Int] -> PrimType -> ValueType
V.ValueType [] (PrimType -> ValueType) -> PrimType -> ValueType
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
primTypeToValueType PrimType
pt
    onValue (Array o
_ (Shape [a]
ds) (Prim PrimType
pt)) =
      ValueType -> Maybe ValueType
forall a. a -> Maybe a
Just (ValueType -> Maybe ValueType) -> ValueType -> Maybe ValueType
forall a b. (a -> b) -> a -> b
$ [Int] -> PrimType -> ValueType
V.ValueType ((a -> Int) -> [a] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral [a]
ds) (PrimType -> ValueType) -> PrimType -> ValueType
forall a b. (a -> b) -> a -> b
$ PrimType -> PrimType
primTypeToValueType PrimType
pt
    onValue TypeBase a o
_ = Maybe ValueType
forall a. Maybe a
Nothing

loadResType :: ValueType -> ValueType
loadResType :: ValueType -> ValueType
loadResType (Scalar (Arrow NoMode
_ PName
_ Diet
_ ValueType
_ (RetType [VName]
_ TypeBase Int64 Freshness
rt))) = TypeBase Int64 Freshness -> ValueType
forall dim o. TypeBase dim o -> TypeBase dim NoMode
toStruct TypeBase Int64 Freshness
rt
loadResType ValueType
t =
  String -> ValueType
forall a. HasCallStack => String -> a
error (String -> ValueType) -> String -> ValueType
forall a b. (a -> b) -> a -> b
$ String
"loadResType: invalid type " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ValueType -> String
forall a. Pretty a => a -> String
prettyString ValueType
t

-- | Determine which IO operation this is.
--
-- If you want to add a new one, then remember to also add it as an intrinsic in
-- the type checker, and probably also to the prelude.
determineIO :: (Monad m) => T.Text -> Maybe (m ValueType -> Value m -> m IOOp)
determineIO :: forall (m :: * -> *).
Monad m =>
Text -> Maybe (m ValueType -> Value m -> m IOOp)
determineIO Text
"io_loadbytes" = (m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall (m :: * -> *).
(m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
load ((m ValueType -> String -> m IOOp)
 -> Maybe (m ValueType -> Value m -> m IOOp))
-> (m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall a b. (a -> b) -> a -> b
$ (String -> m IOOp) -> m ValueType -> String -> m IOOp
forall a b. a -> b -> a
const ((String -> m IOOp) -> m ValueType -> String -> m IOOp)
-> (String -> m IOOp) -> m ValueType -> String -> m IOOp
forall a b. (a -> b) -> a -> b
$ IOOp -> m IOOp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IOOp -> m IOOp) -> (String -> IOOp) -> String -> m IOOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IOOp
LoadBytes
determineIO Text
"io_loadimg" = (m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall (m :: * -> *).
(m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
load ((m ValueType -> String -> m IOOp)
 -> Maybe (m ValueType -> Value m -> m IOOp))
-> (m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall a b. (a -> b) -> a -> b
$ (String -> m IOOp) -> m ValueType -> String -> m IOOp
forall a b. a -> b -> a
const ((String -> m IOOp) -> m ValueType -> String -> m IOOp)
-> (String -> m IOOp) -> m ValueType -> String -> m IOOp
forall a b. (a -> b) -> a -> b
$ IOOp -> m IOOp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IOOp -> m IOOp) -> (String -> IOOp) -> String -> m IOOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IOOp
LoadImg
determineIO Text
"io_loadaudio" = (m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall (m :: * -> *).
(m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
load ((m ValueType -> String -> m IOOp)
 -> Maybe (m ValueType -> Value m -> m IOOp))
-> (m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall a b. (a -> b) -> a -> b
$ (String -> m IOOp) -> m ValueType -> String -> m IOOp
forall a b. a -> b -> a
const ((String -> m IOOp) -> m ValueType -> String -> m IOOp)
-> (String -> m IOOp) -> m ValueType -> String -> m IOOp
forall a b. (a -> b) -> a -> b
$ IOOp -> m IOOp
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IOOp -> m IOOp) -> (String -> IOOp) -> String -> m IOOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IOOp
LoadAudio
determineIO Text
"io_loadvalue" = (m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall (m :: * -> *).
(m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
load ((m ValueType -> String -> m IOOp)
 -> Maybe (m ValueType -> Value m -> m IOOp))
-> (m ValueType -> String -> m IOOp)
-> Maybe (m ValueType -> Value m -> m IOOp)
forall a b. (a -> b) -> a -> b
$ \m ValueType
t String
fname -> do
  t' <- m ValueType
t
  pure $ LoadValue (loadResType t') fname
determineIO Text
_ = Maybe (m ValueType -> Value m -> m IOOp)
forall a. Maybe a
Nothing
{-# NOINLINE determineIO #-}

withTempDir :: (FilePath -> IO a) -> IO a
withTempDir :: forall a. (String -> IO a) -> IO a
withTempDir = String -> (String -> IO a) -> IO a
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
String -> (String -> m a) -> m a
withSystemTempDirectory String
"futhark"

system ::
  FilePath ->
  [String] ->
  T.Text ->
  IO T.Text
system :: String -> [String] -> Text -> IO Text
system String
prog [String]
options Text
input = do
  res <- String
-> [String]
-> ByteString
-> IO (Either IOException (ExitCode, String, String))
runProgramWithExitCode String
prog [String]
options (ByteString -> IO (Either IOException (ExitCode, String, String)))
-> ByteString -> IO (Either IOException (ExitCode, String, String))
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
T.encodeUtf8 Text
input
  case res of
    Left IOException
err ->
      String -> IO Text
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO Text) -> String -> IO Text
forall a b. (a -> b) -> a -> b
$ String
prog' String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" failed: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> IOException -> String
forall a. Show a => a -> String
show IOException
err
    Right (ExitCode
ExitSuccess, String
stdout_t, String
_) ->
      Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> IO Text) -> Text -> IO Text
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
stdout_t
    Right (ExitFailure Int
code', String
_, String
stderr_t) ->
      String -> IO Text
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO Text) -> String -> IO Text
forall a b. (a -> b) -> a -> b
$
        String
prog'
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" failed with exit code "
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
code'
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" and stderr:\n"
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
stderr_t
  where
    prog' :: String
prog' = String
"\"" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
prog String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"\""

loadBMP :: FilePath -> IO V.Value
loadBMP :: String -> IO Value
loadBMP String
bmpfile = do
  res <- String -> IO (Either Error BMP)
BMP.readBMP String
bmpfile
  case res of
    Left Error
err ->
      String -> IO Value
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO Value) -> String -> IO Value
forall a b. (a -> b) -> a -> b
$ String
"Failed to read BMP:\n" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Error -> String
forall a. Show a => a -> String
show Error
err
    Right BMP
bmp -> do
      let bmp_bs :: ByteString
bmp_bs = BMP -> ByteString
BMP.unpackBMPToRGBA32 BMP
bmp
          (Int
w, Int
h) = BMP -> (Int, Int)
BMP.bmpDimensions BMP
bmp
          shape :: Vector Int
shape = [Int] -> Vector Int
forall a. Storable a => [a] -> Vector a
SVec.fromList [Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
h, Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
w]
          pix :: Int -> a
pix Int
l =
            let (Int
i, Int
j) = Int
l Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
w
                l' :: Int
l' = (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j
                r :: a
r = Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> a) -> Word8 -> a
forall a b. (a -> b) -> a -> b
$ ByteString
bmp_bs HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
`BS.index` (Int
l' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)
                g :: a
g = Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> a) -> Word8 -> a
forall a b. (a -> b) -> a -> b
$ ByteString
bmp_bs HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
`BS.index` (Int
l' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                b :: a
b = Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> a) -> Word8 -> a
forall a b. (a -> b) -> a -> b
$ ByteString
bmp_bs HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
`BS.index` (Int
l' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
                a :: a
a = Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> a) -> Word8 -> a
forall a b. (a -> b) -> a -> b
$ ByteString
bmp_bs HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
`BS.index` (Int
l' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)
             in (a
a a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
24) a -> a -> a
forall a. Bits a => a -> a -> a
.|. (a
r a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
16) a -> a -> a
forall a. Bits a => a -> a -> a
.|. (a
g a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
8) a -> a -> a
forall a. Bits a => a -> a -> a
.|. a
b
      Value -> IO Value
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value -> IO Value) -> Value -> IO Value
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Word32 -> Value
V.U32Value Vector Int
shape (Vector Word32 -> Value) -> Vector Word32 -> Value
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Word32) -> Vector Word32
forall a. Storable a => Int -> (Int -> a) -> Vector a
SVec.generate (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) Int -> Word32
forall {a}. (Bits a, Num a) => Int -> a
pix

loadImage :: FilePath -> IO V.Value
loadImage :: String -> IO Value
loadImage String
imgfile =
  (String -> IO Value) -> IO Value
forall a. (String -> IO a) -> IO a
withTempDir ((String -> IO Value) -> IO Value)
-> (String -> IO Value) -> IO Value
forall a b. (a -> b) -> a -> b
$ \String
dir -> do
    let bmpfile :: String
bmpfile = String
dir String -> String -> String
</> String -> String
takeBaseName String
imgfile String -> String -> String
`replaceExtension` String
"bmp"
    IO Text -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Text -> IO ()) -> IO Text -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> [String] -> Text -> IO Text
system String
"convert" [String
imgfile, String
"-type", String
"TrueColorAlpha", String
bmpfile] Text
forall a. Monoid a => a
mempty
    String -> IO Value
loadBMP String
bmpfile

loadPCM :: Int -> FilePath -> IO V.Value
loadPCM :: Int -> String -> IO Value
loadPCM Int
num_channels String
pcmfile = do
  contents <- String -> IO ByteString
LBS.readFile String
pcmfile
  let v = ByteString -> Vector Double
forall a. Storable a => ByteString -> Vector a
SVec.byteStringToVector (ByteString -> Vector Double) -> ByteString -> Vector Double
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
LBS.toStrict ByteString
contents
      channel_length = Vector Double -> Int
forall a. Storable a => Vector a -> Int
SVec.length Vector Double
v Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
num_channels
      shape =
        [Int] -> Vector Int
forall a. Storable a => [a] -> Vector a
SVec.fromList
          [ Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
num_channels,
            Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
channel_length
          ]
      -- ffmpeg outputs audio data in column-major format. `backPermuter` computes the
      -- tranposed indexes for a backpermutation.
      backPermuter Int
i = (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
channel_length) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
num_channels Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
channel_length
      perm = Int -> (Int -> Int) -> Vector Int
forall a. Storable a => Int -> (Int -> a) -> Vector a
SVec.generate (Vector Double -> Int
forall a. Storable a => Vector a -> Int
SVec.length Vector Double
v) Int -> Int
backPermuter
  pure $ V.F64Value shape $ SVec.backpermute v perm

loadAudio :: FilePath -> IO V.Value
loadAudio :: String -> IO Value
loadAudio String
audiofile = do
  s <- String -> [String] -> Text -> IO Text
system String
"ffprobe" [String
audiofile, String
"-show_entries", String
"stream=channels", String
"-select_streams", String
"a", String
"-of", String
"compact=p=0:nk=1", String
"-v", String
"0"] Text
forall a. Monoid a => a
mempty
  case T.decimal s of
    Right (Int
num_channels, Text
_) -> do
      (String -> IO Value) -> IO Value
forall a. (String -> IO a) -> IO a
withTempDir ((String -> IO Value) -> IO Value)
-> (String -> IO Value) -> IO Value
forall a b. (a -> b) -> a -> b
$ \String
dir -> do
        let pcmfile :: String
pcmfile = String
dir String -> String -> String
</> String -> String
takeBaseName String
audiofile String -> String -> String
`replaceExtension` String
"pcm"
        IO Text -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Text -> IO ()) -> IO Text -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> [String] -> Text -> IO Text
system String
"ffmpeg" [String
"-i", String
audiofile, String
"-c:a", String
"pcm_f64le", String
"-map", String
"0", String
"-f", String
"data", String
pcmfile] Text
forall a. Monoid a => a
mempty
        Int -> String -> IO Value
loadPCM Int
num_channels String
pcmfile
    Either String (Int, Text)
_ -> String -> IO Value
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"io.loadaudio failed to detect the number of channels in the audio input"

tryIO :: (MonadIO m) => IO a -> m (Either T.Text a)
tryIO :: forall (m :: * -> *) a. MonadIO m => IO a -> m (Either Text a)
tryIO =
  (SomeException -> m (Either Text a))
-> (a -> m (Either Text a))
-> Either SomeException a
-> m (Either Text a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either
    ( Either Text a -> m (Either Text a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        (Either Text a -> m (Either Text a))
-> (SomeException -> Either Text a)
-> SomeException
-> m (Either Text a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Either Text a
forall a b. a -> Either a b
Left
        (Text -> Either Text a)
-> (SomeException -> Text) -> SomeException -> Either Text a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack
        (String -> Text)
-> (SomeException -> String) -> SomeException -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SomeException -> String
forall e. Exception e => e -> String
displayException :: SomeException -> String)
    )
    (Either Text a -> m (Either Text a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Text a -> m (Either Text a))
-> (a -> Either Text a) -> a -> m (Either Text a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either Text a
forall a b. b -> Either a b
Right)
    (Either SomeException a -> m (Either Text a))
-> (IO a -> m (Either SomeException a))
-> IO a
-> m (Either Text a)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< IO (Either SomeException a) -> m (Either SomeException a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either SomeException a) -> m (Either SomeException a))
-> (IO a -> IO (Either SomeException a))
-> IO a
-> m (Either SomeException a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> IO (Either SomeException a)
forall e a. Exception e => IO a -> IO (Either e a)
try

loadValues :: FilePath -> IO [V.Value]
loadValues :: String -> IO [Value]
loadValues String
datafile = do
  contents <- IO ByteString -> IO ByteString
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ String -> IO ByteString
LBS.readFile String
datafile
  maybe (fail $ "Failed to read data file: " <> datafile) pure $
    V.readValues contents

-- | Resolve relative file paths in the operation relative to the given
-- directory, rather than the current working directory.
ioRelativeTo :: FilePath -> IOOp -> IOOp
ioRelativeTo :: String -> IOOp -> IOOp
ioRelativeTo String
dir (LoadBytes String
f) = String -> IOOp
LoadBytes (String -> IOOp) -> String -> IOOp
forall a b. (a -> b) -> a -> b
$ String
dir String -> String -> String
</> String
f
ioRelativeTo String
dir (LoadImg String
f) = String -> IOOp
LoadImg (String -> IOOp) -> String -> IOOp
forall a b. (a -> b) -> a -> b
$ String
dir String -> String -> String
</> String
f
ioRelativeTo String
dir (LoadAudio String
f) = String -> IOOp
LoadAudio (String -> IOOp) -> String -> IOOp
forall a b. (a -> b) -> a -> b
$ String
dir String -> String -> String
</> String
f
ioRelativeTo String
dir (LoadValue ValueType
t String
f) = ValueType -> String -> IOOp
LoadValue ValueType
t (String -> IOOp) -> String -> IOOp
forall a b. (a -> b) -> a -> b
$ String
dir String -> String -> String
</> String
f

-- | Turn a loaded value into an interpreter value. Arrays are put on the server,
-- if there is one, as they are very expensive to represent in the interpreter,
-- and are often just passed on to entry points. If the server cannot hold the
-- array, which happens when no entry point uses its type, it is kept in the
-- interpreter instead.
fromData :: Maybe FFI.Server -> V.Value -> IO (Value m)
fromData :: forall (m :: * -> *). Maybe Server -> Value -> IO (Value m)
fromData (Just Server
s) Value
v
  | dims :: [Int]
dims@(Int
_ : [Int]
_) <- Value -> [Int]
V.valueShape Value
v = do
      let shape :: Shape Int64
shape = (Int -> Shape Int64 -> Shape Int64)
-> Shape Int64 -> [Int] -> Shape Int64
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Int64 -> Shape Int64 -> Shape Int64
forall d. d -> Shape d -> Shape d
ShapeDim (Int64 -> Shape Int64 -> Shape Int64)
-> (Int -> Int64) -> Int -> Shape Int64 -> Shape Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Shape Int64
forall d. Shape d
ShapeLeaf [Int]
dims
      -- We run the FFI GC just as if we call an entry point, to avoid building
      -- up too much data.
      (String -> IO (Value m))
-> (ValueRef -> IO (Value m))
-> Either String ValueRef
-> IO (Value m)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (IO (Value m) -> String -> IO (Value m)
forall a b. a -> b -> a
const (IO (Value m) -> String -> IO (Value m))
-> IO (Value m) -> String -> IO (Value m)
forall a b. (a -> b) -> a -> b
$ Value m -> IO (Value m)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value m -> IO (Value m)) -> Value m -> IO (Value m)
forall a b. (a -> b) -> a -> b
$ Value -> Value m
forall (m :: * -> *). Value -> Value m
fromDataValue Value
v) (\ValueRef
ref -> Value m -> IO (Value m)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value m -> IO (Value m)) -> Value m -> IO (Value m)
forall a b. (a -> b) -> a -> b
$ Shape Int64 -> ValueRef -> [Int64] -> Value m
forall (m :: * -> *). Shape Int64 -> ValueRef -> [Int64] -> Value m
ValueLazyFFI Shape Int64
shape ValueRef
ref [])
        (Either String ValueRef -> IO (Value m))
-> IO (Either String ValueRef) -> IO (Value m)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Server -> ServerM ValueRef -> IO (Either String ValueRef)
forall a. Server -> ServerM a -> IO (Either String a)
FFI.runServerM Server
s (ServerM ()
FFI.gc ServerM () -> ServerM ValueRef -> ServerM ValueRef
forall a b. ServerM a -> ServerM b -> ServerM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Value -> ServerM ValueRef
FFI.putData Value
v)
fromData Maybe Server
_ Value
v = Value m -> IO (Value m)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Value m -> IO (Value m)) -> Value m -> IO (Value m)
forall a b. (a -> b) -> a -> b
$ Value -> Value m
forall (m :: * -> *). Value -> Value m
fromDataValue Value
v

-- | Run an IO operation. Arrays are put on the server, if one is given.
doIOOp :: Maybe FFI.Server -> IOOp -> IO (Either T.Text (Value m))
doIOOp :: forall (m :: * -> *).
Maybe Server -> IOOp -> IO (Either Text (Value m))
doIOOp Maybe Server
s (LoadBytes String
fname) =
  IO (Value m) -> IO (Either Text (Value m))
forall (m :: * -> *) a. MonadIO m => IO a -> m (Either Text a)
tryIO (IO (Value m) -> IO (Either Text (Value m)))
-> IO (Value m) -> IO (Either Text (Value m))
forall a b. (a -> b) -> a -> b
$ Maybe Server -> Value -> IO (Value m)
forall (m :: * -> *). Maybe Server -> Value -> IO (Value m)
fromData Maybe Server
s (Value -> IO (Value m))
-> (ByteString -> Value) -> ByteString -> IO (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Value
forall t. PutValue1 t => t -> Value
V.putValue1 (ByteString -> IO (Value m)) -> IO ByteString -> IO (Value m)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> IO ByteString
BS.readFile String
fname
doIOOp Maybe Server
s (LoadImg String
fname) =
  IO (Value m) -> IO (Either Text (Value m))
forall (m :: * -> *) a. MonadIO m => IO a -> m (Either Text a)
tryIO (IO (Value m) -> IO (Either Text (Value m)))
-> IO (Value m) -> IO (Either Text (Value m))
forall a b. (a -> b) -> a -> b
$ Maybe Server -> Value -> IO (Value m)
forall (m :: * -> *). Maybe Server -> Value -> IO (Value m)
fromData Maybe Server
s (Value -> IO (Value m)) -> IO Value -> IO (Value m)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> IO Value
loadImage String
fname
doIOOp Maybe Server
s (LoadAudio String
fname) =
  IO (Value m) -> IO (Either Text (Value m))
forall (m :: * -> *) a. MonadIO m => IO a -> m (Either Text a)
tryIO (IO (Value m) -> IO (Either Text (Value m)))
-> IO (Value m) -> IO (Either Text (Value m))
forall a b. (a -> b) -> a -> b
$ Maybe Server -> Value -> IO (Value m)
forall (m :: * -> *). Maybe Server -> Value -> IO (Value m)
fromData Maybe Server
s (Value -> IO (Value m)) -> IO Value -> IO (Value m)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> IO Value
loadAudio String
fname
doIOOp Maybe Server
s (LoadValue ValueType
t String
fname) =
  case ValueType -> Maybe [ValueType]
typeToValueTypes ValueType
t of
    Maybe [ValueType]
Nothing ->
      Either Text (Value m) -> IO (Either Text (Value m))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Text (Value m) -> IO (Either Text (Value m)))
-> (String -> Either Text (Value m))
-> String
-> IO (Either Text (Value m))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Either Text (Value m)
forall a b. a -> Either a b
Left (Text -> Either Text (Value m))
-> (String -> Text) -> String -> Either Text (Value m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack (String -> IO (Either Text (Value m)))
-> String -> IO (Either Text (Value m))
forall a b. (a -> b) -> a -> b
$
        String
"Cannot load a value of type "
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ValueType -> String
forall a. Pretty a => a -> String
prettyString ValueType
t
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": only primitives, arrays of primitives, and tuples of these can be loaded."
    Just [ValueType]
ts -> Maybe Server -> [ValueType] -> String -> IO (Either Text (Value m))
forall (m :: * -> *).
Maybe Server -> [ValueType] -> String -> IO (Either Text (Value m))
doLoadValue Maybe Server
s [ValueType]
ts String
fname
{-# NOINLINE doIOOp #-}

doLoadValue :: Maybe FFI.Server -> [V.ValueType] -> FilePath -> IO (Either T.Text (Value m))
doLoadValue :: forall (m :: * -> *).
Maybe Server -> [ValueType] -> String -> IO (Either Text (Value m))
doLoadValue Maybe Server
s [ValueType]
ts String
fname = IO (Value m) -> IO (Either Text (Value m))
forall (m :: * -> *) a. MonadIO m => IO a -> m (Either Text a)
tryIO (IO (Value m) -> IO (Either Text (Value m)))
-> IO (Value m) -> IO (Either Text (Value m))
forall a b. (a -> b) -> a -> b
$ do
  vs <- String -> IO [Value]
loadValues String
fname
  let vs_ts = (Value -> ValueType) -> [Value] -> [ValueType]
forall a b. (a -> b) -> [a] -> [b]
map Value -> ValueType
V.valueType [Value]
vs
  when (vs_ts /= ts) . fail $
    "Expected file \""
      <> fname
      <> "\" to contain data of types "
      <> unwords (map prettyString ts)
      <> " but found data of types "
      <> unwords (map prettyString vs_ts)
  asValue <$> mapM (fromData s) vs
  where
    asValue :: [Value m] -> Value m
asValue [Value m
v] = Value m
v
    asValue [Value m]
vs = [Value m] -> Value m
forall {m :: * -> *}. [Value m] -> Value m
toTuple [Value m]
vs