{-# LANGUAGE FlexibleInstances #-}

-- |
-- Extism.PDK.Memory implements a low-level interface for interacting with Extism memory
module Extism.PDK.Memory
  ( Memory (..),
    MemoryOffset,
    MemoryLength,
    FromBytes (..),
    ToBytes (..),
    JSON (..),
    MsgPack (..),
    load,
    loadString,
    loadByteString,
    outputMemory,
    memAlloc,
    free,
    alloc,
    allocString,
    allocByteString,
    memoryOffset,
    memoryLength,
    findMemory,
  )
where

import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString as B
import Data.ByteString.Internal (c2w, w2c)
import Data.Int
import Data.Word
import Extism.PDK.Bindings
import qualified Extism.PDK.MsgPack (MsgPack, decode, encode)
import Extism.PDK.Util
import qualified Text.JSON (JSON, Result (..), decode, encode)
import qualified Text.JSON.Generic

-- | Represents a block of memory by offset and length
data Memory = Memory MemoryOffset MemoryLength

-- | Load data from 'Memory' block
load :: (FromBytes a) => Memory -> IO (Either String a)
load :: forall a. FromBytes a => Memory -> IO (Either String a)
load (Memory MemoryOffset
offs MemoryOffset
len) = do
  ByteString
x <- MemoryOffset -> MemoryOffset -> IO ByteString
readBytes MemoryOffset
offs MemoryOffset
len
  Either String a -> IO (Either String a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String a -> IO (Either String a))
-> Either String a -> IO (Either String a)
forall a b. (a -> b) -> a -> b
$ ByteString -> Either String a
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
x

-- | Store data into a 'Memory' block
store :: (ToBytes a) => Memory -> a -> IO ()
store :: forall a. ToBytes a => Memory -> a -> IO ()
store (Memory MemoryOffset
offs MemoryOffset
len) a
a =
  let bs :: ByteString
bs = a -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes a
a
   in MemoryOffset -> MemoryOffset -> ByteString -> IO ()
writeBytes MemoryOffset
offs MemoryOffset
len ByteString
bs

-- | Set plugin output to the provided 'Memory' block
outputMemory :: Memory -> IO ()
outputMemory :: Memory -> IO ()
outputMemory (Memory MemoryOffset
offs MemoryOffset
len) =
  MemoryOffset -> MemoryOffset -> IO ()
extismSetOutput MemoryOffset
offs MemoryOffset
len

-- | Load ByteString from 'Memory' block
loadByteString :: Memory -> IO B.ByteString
loadByteString :: Memory -> IO ByteString
loadByteString (Memory MemoryOffset
offs MemoryOffset
len) = do
  MemoryOffset -> MemoryOffset -> IO ByteString
readBytes MemoryOffset
offs MemoryOffset
len

-- | Load string from 'Memory' block
loadString :: Memory -> IO String
loadString :: Memory -> IO String
loadString (Memory MemoryOffset
offs MemoryOffset
len) =
  ByteString -> String
fromByteString (ByteString -> String) -> IO ByteString -> IO String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MemoryOffset -> MemoryOffset -> IO ByteString
readBytes MemoryOffset
offs MemoryOffset
len

-- | Store string in 'Memory' block
storeString :: Memory -> String -> IO ()
storeString :: Memory -> String -> IO ()
storeString Memory
mem String
s =
  let bs :: ByteString
bs = String -> ByteString
toByteString String
s
   in Memory -> ByteString -> IO ()
storeByteString Memory
mem ByteString
bs

-- | Store byte string in 'Memory' block
storeByteString :: Memory -> B.ByteString -> IO ()
storeByteString :: Memory -> ByteString -> IO ()
storeByteString (Memory MemoryOffset
offs MemoryOffset
len) =
  MemoryOffset -> MemoryOffset -> ByteString -> IO ()
writeBytes MemoryOffset
offs MemoryOffset
len

-- | Encode a value and copy it into Extism memory, returning the Memory block
alloc :: (ToBytes a) => a -> IO Memory
alloc :: forall a. ToBytes a => a -> IO Memory
alloc a
x =
  let bs :: ByteString
bs = a -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes a
x
   in do
        Memory MemoryOffset
offs MemoryOffset
len <- Int -> IO Memory
memAlloc (ByteString -> Int
B.length ByteString
bs)
        MemoryOffset -> MemoryOffset -> ByteString -> IO ()
writeBytes MemoryOffset
offs MemoryOffset
len ByteString
bs
        Memory -> IO Memory
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Memory -> IO Memory) -> Memory -> IO Memory
forall a b. (a -> b) -> a -> b
$ MemoryOffset -> MemoryOffset -> Memory
Memory MemoryOffset
offs MemoryOffset
len

-- | Allocate a new 'Memory' block
memAlloc :: Int -> IO Memory
memAlloc :: Int -> IO Memory
memAlloc Int
n =
  let len :: MemoryOffset
len = Int -> MemoryOffset
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
   in do
        MemoryOffset
offs <- MemoryOffset -> IO MemoryOffset
extismAlloc MemoryOffset
len
        Memory -> IO Memory
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Memory -> IO Memory) -> Memory -> IO Memory
forall a b. (a -> b) -> a -> b
$ MemoryOffset -> MemoryOffset -> Memory
Memory MemoryOffset
offs MemoryOffset
len

-- | Free a 'Memory' block
free :: Memory -> IO ()
free :: Memory -> IO ()
free (Memory MemoryOffset
0 MemoryOffset
_) = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
free (Memory MemoryOffset
_ MemoryOffset
0) = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
free (Memory MemoryOffset
offs MemoryOffset
_) =
  MemoryOffset -> IO ()
extismFree MemoryOffset
offs

-- | Allocate a new 'Memory' block and copy the encoded value
allocByteString :: B.ByteString -> IO Memory
allocByteString :: ByteString -> IO Memory
allocByteString ByteString
bs = do
  Memory MemoryOffset
offs MemoryOffset
len <- Int -> IO Memory
memAlloc (ByteString -> Int
B.length ByteString
bs)
  MemoryOffset -> MemoryOffset -> ByteString -> IO ()
writeBytes MemoryOffset
offs MemoryOffset
len ByteString
bs
  Memory -> IO Memory
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MemoryOffset -> MemoryOffset -> Memory
Memory MemoryOffset
offs MemoryOffset
len)

-- | Allocate a new 'Memory' block and copy the contents of the provided 'String'
allocString :: String -> IO Memory
allocString :: String -> IO Memory
allocString = ByteString -> IO Memory
allocByteString (ByteString -> IO Memory)
-> (String -> ByteString) -> String -> IO Memory
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ByteString
toByteString

-- | Get the offset of a 'Memory' block
memoryOffset :: Memory -> MemoryOffset
memoryOffset :: Memory -> MemoryOffset
memoryOffset (Memory MemoryOffset
offs MemoryOffset
_) = MemoryOffset
offs

-- | Get the length of a 'Memory' block
memoryLength :: Memory -> MemoryLength
memoryLength :: Memory -> MemoryOffset
memoryLength (Memory MemoryOffset
_ MemoryOffset
len) = MemoryOffset
len

-- | Find 'Memory' block by offset
findMemory :: MemoryOffset -> IO Memory
findMemory :: MemoryOffset -> IO Memory
findMemory MemoryOffset
offs = do
  MemoryOffset
len <- MemoryOffset -> IO MemoryOffset
extismLength MemoryOffset
offs
  Memory -> IO Memory
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Memory -> IO Memory) -> Memory -> IO Memory
forall a b. (a -> b) -> a -> b
$ MemoryOffset -> MemoryOffset -> Memory
Memory MemoryOffset
offs MemoryOffset
len

-- | A class used to convert values from bytes read from linear memory
class FromBytes a where
  fromBytes :: B.ByteString -> Either String a

-- | A class used to convert values to bytes to be written into linear memory
class ToBytes a where
  toBytes :: a -> B.ByteString

-- | A wrapper type for JSON encoded values
newtype JSON a = JSON a

-- | A wrapper type for MsgPack encoded values
newtype MsgPack a = MsgPack a

instance FromBytes B.ByteString where
  fromBytes :: ByteString -> Either String ByteString
fromBytes = ByteString -> Either String ByteString
forall a b. b -> Either a b
Right

instance ToBytes B.ByteString where
  toBytes :: ByteString -> ByteString
toBytes = ByteString -> ByteString
forall a. a -> a
id

instance FromBytes String where
  fromBytes :: ByteString -> Either String String
fromBytes ByteString
mem =
    let s :: Either String ByteString
s = ByteString -> Either String ByteString
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String ByteString
s of
          Left String
e -> String -> Either String String
forall a b. a -> Either a b
Left String
e
          Right ByteString
x -> String -> Either String String
forall a b. b -> Either a b
Right (String -> Either String String) -> String -> Either String String
forall a b. (a -> b) -> a -> b
$ ByteString -> String
fromByteString ByteString
x

instance ToBytes String where
  toBytes :: String -> ByteString
toBytes = String -> ByteString
toByteString

instance (Text.JSON.Generic.Data a) => FromBytes (JSON a) where
  fromBytes :: ByteString -> Either String (JSON a)
fromBytes ByteString
mem =
    let a :: Either String String
a = ByteString -> Either String String
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String String
a of
          Left String
e -> String -> Either String (JSON a)
forall a b. a -> Either a b
Left String
e
          Right String
x ->
            case String -> Result JSValue
forall a. JSON a => String -> Result a
Text.JSON.decode String
x of
              Text.JSON.Error String
e -> String -> Either String (JSON a)
forall a b. a -> Either a b
Left String
e
              Text.JSON.Ok JSValue
y ->
                case JSValue -> Result a
forall a. Data a => JSValue -> Result a
Text.JSON.Generic.fromJSON JSValue
y of
                  Text.JSON.Error String
e -> String -> Either String (JSON a)
forall a b. a -> Either a b
Left String
e
                  Text.JSON.Ok a
z -> JSON a -> Either String (JSON a)
forall a b. b -> Either a b
Right (a -> JSON a
forall a. a -> JSON a
JSON a
z)

instance (Text.JSON.Generic.Data a) => ToBytes (JSON a) where
  toBytes :: JSON a -> ByteString
toBytes (JSON a
x) = String -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes (a -> String
forall a. Data a => a -> String
Text.JSON.Generic.encodeJSON a
x)

instance (Extism.PDK.MsgPack.MsgPack a) => FromBytes (MsgPack a) where
  fromBytes :: ByteString -> Either String (MsgPack a)
fromBytes ByteString
mem =
    let a :: Either String ByteString
a = ByteString -> Either String ByteString
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String ByteString
a of
          Left String
e -> String -> Either String (MsgPack a)
forall a b. a -> Either a b
Left String
e
          Right ByteString
x ->
            case ByteString -> Either String a
forall a. MsgPack a => ByteString -> Either String a
Extism.PDK.MsgPack.decode ByteString
x of
              Left String
e -> String -> Either String (MsgPack a)
forall a b. a -> Either a b
Left String
e
              Right a
y -> MsgPack a -> Either String (MsgPack a)
forall a b. b -> Either a b
Right (a -> MsgPack a
forall a. a -> MsgPack a
MsgPack a
y)

instance (Extism.PDK.MsgPack.MsgPack a) => ToBytes (MsgPack a) where
  toBytes :: MsgPack a -> ByteString
toBytes (MsgPack a
x) = ByteString -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ a -> ByteString
forall a. MsgPack a => a -> ByteString
Extism.PDK.MsgPack.encode a
x

instance ToBytes Int32 where
  toBytes :: Int32 -> ByteString
toBytes Int32
i = ByteString -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
B.toStrict (Put -> ByteString
runPut (Int32 -> Put
putInt32le Int32
i))

instance FromBytes Int32 where
  fromBytes :: ByteString -> Either String Int32
fromBytes ByteString
mem =
    let bs :: Either String ByteString
bs = ByteString -> Either String ByteString
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String ByteString
bs of
          Left String
e -> String -> Either String Int32
forall a b. a -> Either a b
Left String
e
          Right ByteString
x ->
            case Get Int32
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Int32)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get Int32
getInt32le (ByteString -> ByteString
B.fromStrict ByteString
x) of
              Left (ByteString
_, ByteOffset
_, String
e) -> String -> Either String Int32
forall a b. a -> Either a b
Left String
e
              Right (ByteString
_, ByteOffset
_, Int32
x) -> Int32 -> Either String Int32
forall a b. b -> Either a b
Right Int32
x

instance ToBytes Int64 where
  toBytes :: ByteOffset -> ByteString
toBytes ByteOffset
i = ByteString -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
B.toStrict (Put -> ByteString
runPut (ByteOffset -> Put
putInt64le ByteOffset
i))

instance FromBytes Int64 where
  fromBytes :: ByteString -> Either String ByteOffset
fromBytes ByteString
mem =
    let bs :: Either String ByteString
bs = ByteString -> Either String ByteString
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String ByteString
bs of
          Left String
e -> String -> Either String ByteOffset
forall a b. a -> Either a b
Left String
e
          Right ByteString
x ->
            case Get ByteOffset
-> ByteString
-> Either
     (ByteString, ByteOffset, String)
     (ByteString, ByteOffset, ByteOffset)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get ByteOffset
getInt64le (ByteString -> ByteString
B.fromStrict ByteString
x) of
              Left (ByteString
_, ByteOffset
_, String
e) -> String -> Either String ByteOffset
forall a b. a -> Either a b
Left String
e
              Right (ByteString
_, ByteOffset
_, ByteOffset
x) -> ByteOffset -> Either String ByteOffset
forall a b. b -> Either a b
Right ByteOffset
x

instance ToBytes Word32 where
  toBytes :: Word32 -> ByteString
toBytes Word32
i = ByteString -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
B.toStrict (Put -> ByteString
runPut (Word32 -> Put
putWord32le Word32
i))

instance FromBytes Word32 where
  fromBytes :: ByteString -> Either String Word32
fromBytes ByteString
mem =
    let bs :: Either String ByteString
bs = ByteString -> Either String ByteString
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String ByteString
bs of
          Left String
e -> String -> Either String Word32
forall a b. a -> Either a b
Left String
e
          Right ByteString
x ->
            case Get Word32
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Word32)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get Word32
getWord32le (ByteString -> ByteString
B.fromStrict ByteString
x) of
              Left (ByteString
_, ByteOffset
_, String
e) -> String -> Either String Word32
forall a b. a -> Either a b
Left String
e
              Right (ByteString
_, ByteOffset
_, Word32
x) -> Word32 -> Either String Word32
forall a b. b -> Either a b
Right Word32
x

instance ToBytes Word64 where
  toBytes :: MemoryOffset -> ByteString
toBytes MemoryOffset
i = ByteString -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
B.toStrict (Put -> ByteString
runPut (MemoryOffset -> Put
putWord64le MemoryOffset
i))

instance FromBytes Word64 where
  fromBytes :: ByteString -> Either String MemoryOffset
fromBytes ByteString
mem =
    let bs :: Either String ByteString
bs = ByteString -> Either String ByteString
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String ByteString
bs of
          Left String
e -> String -> Either String MemoryOffset
forall a b. a -> Either a b
Left String
e
          Right ByteString
x ->
            case Get MemoryOffset
-> ByteString
-> Either
     (ByteString, ByteOffset, String)
     (ByteString, ByteOffset, MemoryOffset)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get MemoryOffset
getWord64le (ByteString -> ByteString
B.fromStrict ByteString
x) of
              Left (ByteString
_, ByteOffset
_, String
e) -> String -> Either String MemoryOffset
forall a b. a -> Either a b
Left String
e
              Right (ByteString
_, ByteOffset
_, MemoryOffset
x) -> MemoryOffset -> Either String MemoryOffset
forall a b. b -> Either a b
Right MemoryOffset
x

instance ToBytes Float where
  toBytes :: Float -> ByteString
toBytes Float
i = ByteString -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
B.toStrict (Put -> ByteString
runPut (Float -> Put
putFloatle Float
i))

instance FromBytes Float where
  fromBytes :: ByteString -> Either String Float
fromBytes ByteString
mem =
    let bs :: Either String ByteString
bs = ByteString -> Either String ByteString
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String ByteString
bs of
          Left String
e -> String -> Either String Float
forall a b. a -> Either a b
Left String
e
          Right ByteString
x ->
            case Get Float
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Float)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get Float
getFloatle (ByteString -> ByteString
B.fromStrict ByteString
x) of
              Left (ByteString
_, ByteOffset
_, String
e) -> String -> Either String Float
forall a b. a -> Either a b
Left String
e
              Right (ByteString
_, ByteOffset
_, Float
x) -> Float -> Either String Float
forall a b. b -> Either a b
Right Float
x

instance ToBytes Double where
  toBytes :: Double -> ByteString
toBytes Double
i = ByteString -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
B.toStrict (Put -> ByteString
runPut (Double -> Put
putDoublele Double
i))

instance FromBytes Double where
  fromBytes :: ByteString -> Either String Double
fromBytes ByteString
mem =
    let bs :: Either String ByteString
bs = ByteString -> Either String ByteString
forall a. FromBytes a => ByteString -> Either String a
fromBytes ByteString
mem
     in case Either String ByteString
bs of
          Left String
e -> String -> Either String Double
forall a b. a -> Either a b
Left String
e
          Right ByteString
x ->
            case Get Double
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, Double)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get Double
getDoublele (ByteString -> ByteString
B.fromStrict ByteString
x) of
              Left (ByteString
_, ByteOffset
_, String
e) -> String -> Either String Double
forall a b. a -> Either a b
Left String
e
              Right (ByteString
_, ByteOffset
_, Double
x) -> Double -> Either String Double
forall a b. b -> Either a b
Right Double
x