{-# LANGUAGE OverloadedStrings #-}

module Binrep.Type.Assembly.Assemble where

import Binrep.Type.Assembly

import Heystone qualified as Keystone
import System.IO.Unsafe ( unsafeDupablePerformIO )
import Control.Monad.IO.Class
import Data.ByteString qualified as BS
import Data.Text.Short qualified as Text.Short
import Data.Text qualified as Text
import Data.List qualified as List

class Assemble arch where
    assemble :: [AsmInstr arch] -> Either String (MachineCode arch)

assemble1
    :: forall arch. Assemble arch
    => AsmInstr arch -> Either String (MachineCode arch)
assemble1 :: forall (arch :: Arch).
Assemble arch =>
AsmInstr arch -> Either String (MachineCode arch)
assemble1 AsmInstr arch
inst = [AsmInstr arch] -> Either String (MachineCode arch)
forall (arch :: Arch).
Assemble arch =>
[AsmInstr arch] -> Either String (MachineCode arch)
assemble [AsmInstr arch
inst]

instance Assemble 'ArmV8ThumbLE where
    assemble :: [AsmInstr 'ArmV8ThumbLE]
-> Either String (MachineCode 'ArmV8ThumbLE)
assemble =
          (ByteString -> MachineCode 'ArmV8ThumbLE)
-> Either String ByteString
-> Either String (MachineCode 'ArmV8ThumbLE)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> MachineCode 'ArmV8ThumbLE
forall (arch :: Arch). ByteString -> MachineCode arch
MachineCode
        (Either String ByteString
 -> Either String (MachineCode 'ArmV8ThumbLE))
-> ([AsmInstr 'ArmV8ThumbLE] -> Either String ByteString)
-> [AsmInstr 'ArmV8ThumbLE]
-> Either String (MachineCode 'ArmV8ThumbLE)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Either String ByteString) -> Either String ByteString
forall a. IO a -> a
unsafeDupablePerformIO
        (IO (Either String ByteString) -> Either String ByteString)
-> ([AsmInstr 'ArmV8ThumbLE] -> IO (Either String ByteString))
-> [AsmInstr 'ArmV8ThumbLE]
-> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Architecture -> [Mode] -> [String] -> IO (Either String ByteString)
forall (m :: * -> *).
MonadIO m =>
Architecture -> [Mode] -> [String] -> m (Either String ByteString)
assemble' Architecture
Keystone.ArchArm [Mode]
modeFlags
        ([String] -> IO (Either String ByteString))
-> ([AsmInstr 'ArmV8ThumbLE] -> [String])
-> [AsmInstr 'ArmV8ThumbLE]
-> IO (Either String ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [AsmInstr 'ArmV8ThumbLE] -> [String]
forall (arch :: Arch). [AsmInstr arch] -> [String]
prepInstrs
      where
        modeFlags :: [Mode]
modeFlags =
            [Mode
Keystone.ModeV8, Mode
Keystone.ModeThumb, Mode
Keystone.ModeLittleEndian]

-- | TODO This is stupid because the assembler takes a '[String]'. Great for
--   end-user, poor for performance. I want the option to give it a
--   'BS.ByteString' that I've already prepared (as is the interface).
prepInstrs :: forall arch. [AsmInstr arch] -> [String]
prepInstrs :: forall (arch :: Arch). [AsmInstr arch] -> [String]
prepInstrs =
      String -> [String]
forall a. a -> [a]
List.singleton
    (String -> [String])
-> ([AsmInstr arch] -> String) -> [AsmInstr arch] -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
Text.unpack
    (Text -> String)
-> ([AsmInstr arch] -> Text) -> [AsmInstr arch] -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text] -> Text
Text.intercalate (String -> Text
Text.pack String
";")
    ([Text] -> Text)
-> ([AsmInstr arch] -> [Text]) -> [AsmInstr arch] -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AsmInstr arch -> Text) -> [AsmInstr arch] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (ShortText -> Text
Text.Short.toText (ShortText -> Text)
-> (AsmInstr arch -> ShortText) -> AsmInstr arch -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AsmInstr arch -> ShortText
forall (arch :: Arch). AsmInstr arch -> ShortText
getAsmInstr)

-- | Ideally, the assembler takes a raw 'BS.ByteString'. A
--   'BS.Short.ShortByteString' isn't appropriate, because it could be quite
--   large. But this way, this function is basically "compose a bunch of short
--   bytestrings into one big one".
prepInstrs' :: forall arch. [AsmInstr arch] -> BS.ByteString
prepInstrs' :: forall (arch :: Arch). [AsmInstr arch] -> ByteString
prepInstrs' =
      ByteString -> [ByteString] -> ByteString
BS.intercalate ByteString
";"
    ([ByteString] -> ByteString)
-> ([AsmInstr arch] -> [ByteString])
-> [AsmInstr arch]
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AsmInstr arch -> ByteString) -> [AsmInstr arch] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (ShortText -> ByteString
Text.Short.toByteString (ShortText -> ByteString)
-> (AsmInstr arch -> ShortText) -> AsmInstr arch -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AsmInstr arch -> ShortText
forall (arch :: Arch). AsmInstr arch -> ShortText
getAsmInstr)

assemble'
    :: MonadIO m
    => Keystone.Architecture -> [Keystone.Mode]
    -> [String]
    -> m (Either String BS.ByteString)
assemble' :: forall (m :: * -> *).
MonadIO m =>
Architecture -> [Mode] -> [String] -> m (Either String ByteString)
assemble' Architecture
arch [Mode]
modes [String]
instrs = do
    let as' :: Assembler Engine
as' = Architecture -> [Mode] -> Assembler Engine
Keystone.open Architecture
arch [Mode]
modes
    IO (Either Error Engine) -> m (Either Error Engine)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Assembler Engine -> IO (Either Error Engine)
forall a. Assembler a -> IO (Either Error a)
Keystone.runAssembler Assembler Engine
as') m (Either Error Engine)
-> (Either Error Engine -> m (Either String ByteString))
-> m (Either String ByteString)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Left  Error
e  -> String -> m (Either String ByteString)
forall {a} {b}. a -> m (Either a b)
err (String -> m (Either String ByteString))
-> String -> m (Either String ByteString)
forall a b. (a -> b) -> a -> b
$ String
"failed to obtain assembler: "String -> String -> String
forall a. Semigroup a => a -> a -> a
<>Error -> String
forall a. Show a => a -> String
show Error
e
      Right Engine
as -> do
        let out' :: Assembler (ByteString, Int)
out' = Engine -> [String] -> Maybe Word64 -> Assembler (ByteString, Int)
Keystone.assemble Engine
as [String]
instrs Maybe Word64
forall a. Maybe a
Nothing
        -- TODO have to inspect engine to find error. probably say if x=1 OK, if
        -- x>1 weird error, if x<1 check errno->strerror
        IO (Either Error (ByteString, Int))
-> m (Either Error (ByteString, Int))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Assembler (ByteString, Int) -> IO (Either Error (ByteString, Int))
forall a. Assembler a -> IO (Either Error a)
Keystone.runAssembler Assembler (ByteString, Int)
out') m (Either Error (ByteString, Int))
-> (Either Error (ByteString, Int) -> m (Either String ByteString))
-> m (Either String ByteString)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Left Error
e -> String -> m (Either String ByteString)
forall {a} {b}. a -> m (Either a b)
err (String -> m (Either String ByteString))
-> String -> m (Either String ByteString)
forall a b. (a -> b) -> a -> b
$ String
"error while assembling: "String -> String -> String
forall a. Semigroup a => a -> a -> a
<>Error -> String
forall a. Show a => a -> String
show Error
e
          Right (ByteString
mc, Int
_count) -> Either String ByteString -> m (Either String ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String ByteString -> m (Either String ByteString))
-> Either String ByteString -> m (Either String ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString -> Either String ByteString
forall a b. b -> Either a b
Right ByteString
mc
  where err :: a -> m (Either a b)
err = Either a b -> m (Either a b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either a b -> m (Either a b))
-> (a -> Either a b) -> a -> m (Either a b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either a b
forall a b. a -> Either a b
Left