{-# 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 = 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 =
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (arch :: Arch). ByteString -> MachineCode arch
MachineCode
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IO a -> a
unsafeDupablePerformIO
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *).
MonadIO m =>
Architecture -> [Mode] -> [String] -> m (Either String ByteString)
assemble' Architecture
Keystone.ArchArm [Mode]
modeFlags
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 =
      forall a. a -> [a]
List.singleton
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
Text.unpack
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text] -> Text
Text.intercalate (String -> Text
Text.pack String
";")
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (ShortText -> Text
Text.Short.toText forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
";"
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (ShortText -> ByteString
Text.Short.toByteString forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Assembler a -> IO (Either Error a)
Keystone.runAssembler Assembler Engine
as') forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Left  Error
e  -> forall {a} {b}. a -> m (Either a b)
err forall a b. (a -> b) -> a -> b
$ String
"failed to obtain assembler: "forall a. Semigroup a => a -> a -> a
<>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 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
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Assembler a -> IO (Either Error a)
Keystone.runAssembler Assembler (ByteString, Int)
out') forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Left Error
e -> forall {a} {b}. a -> m (Either a b)
err forall a b. (a -> b) -> a -> b
$ String
"error while assembling: "forall a. Semigroup a => a -> a -> a
<>forall a. Show a => a -> String
show Error
e
          Right (ByteString
mc, Int
_count) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right ByteString
mc
  where err :: a -> m (Either a b)
err = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> Either a b
Left