{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}

module Extism.PDK (module Extism.PDK, module Extism.Manifest) where

import Extism.PDK.Bindings
import Extism.JSON(JSValue, JSON)
import Extism.Manifest(toString)
import Data.Word
import Data.Int
import Data.ByteString as B
import Data.ByteString.Internal (c2w, w2c)
import Data.ByteString.Unsafe (unsafeUseAsCString)
import Text.JSON(JSON, decode, encode, resultToEither)
import qualified Extism.PDK.MsgPack(MsgPack, decode, encode)

newtype JSONValue a = JSONValue a
newtype MsgPackValue a = MsgPackValue a

-- | Represents a block of memory
data Memory = Memory MemoryOffset MemoryLength

-- | Helper function to convert a string to a bytestring
toByteString :: String -> ByteString
toByteString :: String -> ByteString
toByteString String
x = [Word8] -> ByteString
B.pack ((Char -> Word8) -> String -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
Prelude.map Char -> Word8
c2w String
x)

-- | Helper function to convert a bytestring to a string
fromByteString :: ByteString -> String
fromByteString :: ByteString -> String
fromByteString ByteString
bs = (Word8 -> Char) -> [Word8] -> String
forall a b. (a -> b) -> [a] -> [b]
Prelude.map Word8 -> Char
w2c ([Word8] -> String) -> [Word8] -> String
forall a b. (a -> b) -> a -> b
$ ByteString -> [Word8]
B.unpack ByteString
bs

class FromBytes a where
  fromBytes :: ByteString -> a

class ToBytes a where
  toBytes :: a -> ByteString

instance FromBytes ByteString where
  fromBytes :: ByteString -> ByteString
fromBytes ByteString
bs = ByteString
bs

instance ToBytes ByteString where
  toBytes :: ByteString -> ByteString
toBytes ByteString
bs = ByteString
bs

instance FromBytes String where
  fromBytes :: ByteString -> String
fromBytes = ByteString -> String
fromByteString

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

instance JSON a => FromBytes (JSONValue a) where
  fromBytes :: ByteString -> JSONValue a
fromBytes ByteString
x =
    case Result a -> Either String a
forall a. Result a -> Either String a
resultToEither (Result a -> Either String a) -> Result a -> Either String a
forall a b. (a -> b) -> a -> b
$ String -> Result a
forall a. JSON a => String -> Result a
decode (ByteString -> String
fromByteString ByteString
x) of
      Left String
e -> String -> JSONValue a
forall a. HasCallStack => String -> a
error String
e
      Right a
y -> a -> JSONValue a
forall a. a -> JSONValue a
JSONValue a
y

instance JSON a => ToBytes (JSONValue a) where
  toBytes :: JSONValue a -> ByteString
toBytes (JSONValue a
x) = String -> ByteString
toByteString (a -> String
forall a. JSON a => a -> String
encode a
x)

instance Extism.PDK.MsgPack.MsgPack a => FromBytes (MsgPackValue a) where
  fromBytes :: ByteString -> MsgPackValue a
fromBytes 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 -> MsgPackValue a
forall a. HasCallStack => String -> a
error String
e
      Right a
y -> a -> MsgPackValue a
forall a. a -> MsgPackValue a
MsgPackValue a
y

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

-- | Get plugin input as 'ByteString'
input :: FromBytes a => IO a
input :: forall a. FromBytes a => IO a
input = do
  MemoryOffset
len <- IO MemoryOffset
extismInputLength
  ByteString -> a
forall a. FromBytes a => ByteString -> a
fromBytes (ByteString -> a) -> IO ByteString -> IO a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MemoryOffset -> IO ByteString
readInputBytes MemoryOffset
len

-- | Get plugin input as 'Memory' block
inputMemory :: IO Memory
inputMemory :: IO Memory
inputMemory = do
  MemoryOffset
len <- IO MemoryOffset
extismInputLength
  MemoryOffset
offs <- MemoryOffset -> IO MemoryOffset
extismAlloc MemoryOffset
len
  (MemoryOffset -> IO (IO ())) -> [MemoryOffset] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
Prelude.mapM_ (\MemoryOffset
x ->
    MemoryOffset -> Word8 -> IO ()
extismStoreU8 (MemoryOffset
offs MemoryOffset -> MemoryOffset -> MemoryOffset
forall a. Num a => a -> a -> a
+ MemoryOffset
x) (Word8 -> IO ()) -> IO Word8 -> IO (IO ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MemoryOffset -> IO Word8
extismInputLoadU8 MemoryOffset
x) [MemoryOffset
0, MemoryOffset
1 .. 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

-- | Get input as 'JSON'
inputJSON :: JSON a => IO (Maybe a)
inputJSON :: forall a. JSON a => IO (Maybe a)
inputJSON = do
  String
s <- IO String
forall a. FromBytes a => IO a
input :: IO String
  case Result a -> Either String a
forall a. Result a -> Either String a
resultToEither (Result a -> Either String a) -> Result a -> Either String a
forall a b. (a -> b) -> a -> b
$ String -> Result a
forall a. JSON a => String -> Result a
decode String
s of
    Left String
_ -> Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
    Right a
x -> Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
x)

-- | Load data from 'Memory' block
load :: FromBytes a => Memory -> IO a
load :: forall a. FromBytes a => Memory -> IO a
load (Memory MemoryOffset
offs MemoryOffset
len) =
  ByteString -> a
forall a. FromBytes a => ByteString -> a
fromBytes (ByteString -> a) -> IO ByteString -> IO a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MemoryOffset -> MemoryOffset -> IO ByteString
readBytes MemoryOffset
offs MemoryOffset
len

-- | 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

-- | Set plugin output to the provided 'ByteString'
output :: ToBytes a => a -> IO ()
output :: forall a. ToBytes a => a -> IO ()
output a
x =
  let bs :: ByteString
bs = a -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes a
x in
  let len :: MemoryOffset
len = Int -> MemoryOffset
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> MemoryOffset) -> Int -> MemoryOffset
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
B.length ByteString
bs in
  do
    MemoryOffset
offs <- MemoryOffset -> IO MemoryOffset
extismAlloc MemoryOffset
len
    ()
b <- Memory -> ByteString -> IO ()
forall a. ToBytes a => Memory -> a -> IO ()
store (MemoryOffset -> MemoryOffset -> Memory
Memory MemoryOffset
offs MemoryOffset
len) ByteString
bs
    MemoryOffset -> MemoryOffset -> IO ()
extismSetOutput MemoryOffset
offs MemoryOffset
len

-- | Set plugin output to a JSON encoded version of the provided value
outputJSON :: JSON a => a -> IO ()
outputJSON :: forall a. JSON a => a -> IO ()
outputJSON a
x =
  String -> IO ()
forall a. ToBytes a => a -> IO ()
output (a -> String
forall a. JSON a => a -> String
toString a
x)

-- | Load string from 'Memory' block
loadString :: Memory -> IO String
loadString :: Memory -> IO String
loadString Memory
mem = do
  ByteString
bs <- Memory -> IO ByteString
forall a. FromBytes a => Memory -> IO a
load Memory
mem
  String -> IO String
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> IO String) -> String -> IO String
forall a b. (a -> b) -> a -> b
$ ByteString -> String
fromByteString ByteString
bs

-- | 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 ()
forall a. ToBytes a => Memory -> a -> IO ()
store Memory
mem ByteString
bs

-- | Allocate a new 'Memory' block
alloc :: Int -> IO Memory
alloc :: Int -> IO Memory
alloc 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 contents of the provided 'ByteString'
allocByteString :: ByteString -> IO Memory
allocByteString :: ByteString -> IO Memory
allocByteString ByteString
bs = do
  Memory
mem <- Int -> IO Memory
alloc (ByteString -> Int
B.length ByteString
bs)
  Memory -> ByteString -> IO ()
forall a. ToBytes a => Memory -> a -> IO ()
store Memory
mem ByteString
bs
  Memory -> IO Memory
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Memory
mem

-- | Allocate a new 'Memory' block and copy the contents of the provided 'String'
allocString :: String -> IO Memory
allocString :: String -> IO Memory
allocString String
s =
  let bs :: ByteString
bs = String -> ByteString
toByteString String
s in
  ByteString -> IO Memory
allocByteString ByteString
bs

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

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

-- | Find 'Memory' block by offset
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

-- | Get a variable from the Extism runtime
getVar :: String -> IO (Maybe ByteString)
getVar :: String -> IO (Maybe ByteString)
getVar String
key = do
  Memory
k <- String -> IO Memory
allocString String
key
  MemoryOffset
v <- MemoryOffset -> IO MemoryOffset
extismGetVar (Memory -> MemoryOffset
memoryOffset Memory
k)
  Memory -> IO ()
free Memory
k
  if MemoryOffset
v MemoryOffset -> MemoryOffset -> Bool
forall a. Eq a => a -> a -> Bool
== MemoryOffset
0 then
    Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
  else do
    Memory
mem <- MemoryOffset -> IO Memory
findMemory MemoryOffset
v
    ByteString
bs <- Memory -> IO ByteString
forall a. FromBytes a => Memory -> IO a
load Memory
mem
    Memory -> IO ()
free Memory
mem
    Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs)

-- | Set a variable
setVar :: ToBytes a => String -> Maybe a -> IO ()
setVar :: forall a. ToBytes a => String -> Maybe a -> IO ()
setVar String
key Maybe a
Nothing = do
  Memory
k <- String -> IO Memory
allocString String
key
  MemoryOffset -> MemoryOffset -> IO ()
extismSetVar (Memory -> MemoryOffset
memoryOffset Memory
k) MemoryOffset
0
  Memory -> IO ()
free Memory
k
setVar String
key (Just a
v) = do
  Memory
k <- String -> IO Memory
allocString String
key
  Memory
x <- ByteString -> IO Memory
allocByteString (a -> ByteString
forall a. ToBytes a => a -> ByteString
toBytes a
v)
  MemoryOffset -> MemoryOffset -> IO ()
extismSetVar (Memory -> MemoryOffset
memoryOffset Memory
k) (Memory -> MemoryOffset
memoryOffset Memory
x)
  Memory -> IO ()
free Memory
k
  Memory -> IO ()
free Memory
x

-- | Get a configuration value
getConfig :: String -> IO (Maybe String)
getConfig :: String -> IO (Maybe String)
getConfig String
key = do
  Memory
k <- String -> IO Memory
allocString String
key
  MemoryOffset
v <- MemoryOffset -> IO MemoryOffset
extismGetConfig (Memory -> MemoryOffset
memoryOffset Memory
k)
  Memory -> IO ()
free Memory
k
  if MemoryOffset
v MemoryOffset -> MemoryOffset -> Bool
forall a. Eq a => a -> a -> Bool
== MemoryOffset
0 then
    Maybe String -> IO (Maybe String)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe String
forall a. Maybe a
Nothing
  else do
    Memory
mem <- MemoryOffset -> IO Memory
findMemory MemoryOffset
v
    String
s <- Memory -> IO String
loadString Memory
mem
    Memory -> IO ()
free Memory
mem
    Maybe String -> IO (Maybe String)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe String -> IO (Maybe String))
-> Maybe String -> IO (Maybe String)
forall a b. (a -> b) -> a -> b
$ String -> Maybe String
forall a. a -> Maybe a
Just String
s

-- | Set the current error message
setError :: String -> IO ()
setError :: String -> IO ()
setError String
msg = do
  Memory
s <- String -> IO Memory
allocString String
msg
  MemoryOffset -> IO ()
extismSetError (MemoryOffset -> IO ()) -> MemoryOffset -> IO ()
forall a b. (a -> b) -> a -> b
$ Memory -> MemoryOffset
memoryOffset Memory
s
  Memory -> IO ()
free Memory
s

data LogLevel = Info | Debug | Warn | Error

log :: LogLevel -> String -> IO ()
log :: LogLevel -> String -> IO ()
log LogLevel
Info String
msg = do
  Memory
s <- String -> IO Memory
allocString String
msg
  MemoryOffset -> IO ()
extismLogInfo (Memory -> MemoryOffset
memoryOffset Memory
s)
  Memory -> IO ()
free Memory
s
log LogLevel
Debug String
msg = do
  Memory
s <- String -> IO Memory
allocString String
msg
  MemoryOffset -> IO ()
extismLogDebug (Memory -> MemoryOffset
memoryOffset Memory
s)
  Memory -> IO ()
free Memory
s
log LogLevel
Warn String
msg = do
  Memory
s <- String -> IO Memory
allocString String
msg
  MemoryOffset -> IO ()
extismLogWarn (Memory -> MemoryOffset
memoryOffset Memory
s)
  Memory -> IO ()
free Memory
s
log LogLevel
Error String
msg = do
  Memory
s <- String -> IO Memory
allocString String
msg
  MemoryOffset -> IO ()
extismLogError (Memory -> MemoryOffset
memoryOffset Memory
s)
  Memory -> IO ()
free Memory
s