{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-partial-fields #-}

module Wrench.Translator.Types (
    Section (..),
    CodeToken (..),
    DataToken (..),
    DataValue (..),
    ByteSize (..),
    IsWord,
    markupOffsets,
    markupSectionOffsets,
    DerefMnemonic (..),
    resolveRef,
    Ref (..),
    derefSection,
    Preprocessor,
    preprocess,
) where

import Prelude qualified
import Relude
import Text.Megaparsec.Pos (SourcePos, sourcePosPretty)
import Wrench.Machine.Types

-- | Rewrites the source before it is parsed. It keeps every character at its
-- offset, so that errors are shown against the original source, and fails
-- with the offset of what it rejects.
type Preprocessor = String -> Either (Int, Text) String

preprocess :: [Preprocessor] -> String -> Either (Int, Text) String
preprocess :: [Preprocessor] -> Preprocessor
preprocess [Preprocessor]
preprocessors String
src = (String -> Preprocessor -> Either (Int, Text) String)
-> String -> [Preprocessor] -> Either (Int, Text) String
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM String -> Preprocessor -> Either (Int, Text) String
forall a b. a -> (a -> b) -> b
(&) String
src [Preprocessor]
preprocessors

class DerefMnemonic m w where
    derefMnemonic :: (Text -> Maybe w) -> w -> m (Ref w) -> m w

data Section isa w l
    = Code
        { forall isa w l. Section isa w l -> Maybe Int
org :: Maybe Int
        , forall isa w l. Section isa w l -> [CodeToken isa l]
codeTokens :: ![CodeToken isa l]
        }
    | Data
        { org :: Maybe Int
        , forall isa w l. Section isa w l -> [DataToken w l]
dataTokens :: ![DataToken w l]
        }
    deriving (Int -> Section isa w l -> ShowS
[Section isa w l] -> ShowS
Section isa w l -> String
(Int -> Section isa w l -> ShowS)
-> (Section isa w l -> String)
-> ([Section isa w l] -> ShowS)
-> Show (Section isa w l)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall isa w l.
(Show l, Show isa, Show w) =>
Int -> Section isa w l -> ShowS
forall isa w l.
(Show l, Show isa, Show w) =>
[Section isa w l] -> ShowS
forall isa w l.
(Show l, Show isa, Show w) =>
Section isa w l -> String
$cshowsPrec :: forall isa w l.
(Show l, Show isa, Show w) =>
Int -> Section isa w l -> ShowS
showsPrec :: Int -> Section isa w l -> ShowS
$cshow :: forall isa w l.
(Show l, Show isa, Show w) =>
Section isa w l -> String
show :: Section isa w l -> String
$cshowList :: forall isa w l.
(Show l, Show isa, Show w) =>
[Section isa w l] -> ShowS
showList :: [Section isa w l] -> ShowS
Show)

instance (ByteSize isa, ByteSizeT w) => ByteSize (Section isa w l) where
    byteSize :: Section isa w l -> Int
byteSize Code{[CodeToken isa l]
codeTokens :: forall isa w l. Section isa w l -> [CodeToken isa l]
codeTokens :: [CodeToken isa l]
codeTokens} = [Int] -> Int
forall a (f :: * -> *). (Foldable f, Num a) => f a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (CodeToken isa l -> Int) -> [CodeToken isa l] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map CodeToken isa l -> Int
forall t. ByteSize t => t -> Int
byteSize [CodeToken isa l]
codeTokens
    byteSize Data{[DataToken w l]
dataTokens :: forall isa w l. Section isa w l -> [DataToken w l]
dataTokens :: [DataToken w l]
dataTokens} = [Int] -> Int
forall a (f :: * -> *). (Foldable f, Num a) => f a -> a
sum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (DataToken w l -> Int) -> [DataToken w l] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map DataToken w l -> Int
forall t. ByteSize t => t -> Int
byteSize [DataToken w l]
dataTokens

derefSection ::
    forall isa w.
    (ByteSize (isa (Ref w)), DerefMnemonic isa w, IsWord w, Show (isa w)) =>
    (Text -> Maybe w)
    -> w
    -> Section (isa (Ref w)) w Text
    -> Section (isa w) w w
derefSection :: forall (isa :: * -> *) w.
(ByteSize (isa (Ref w)), DerefMnemonic isa w, IsWord w,
 Show (isa w)) =>
(Text -> Maybe w)
-> w -> Section (isa (Ref w)) w Text -> Section (isa w) w w
derefSection Text -> Maybe w
f w
offset code :: Section (isa (Ref w)) w Text
code@Code{[CodeToken (isa (Ref w)) Text]
codeTokens :: forall isa w l. Section isa w l -> [CodeToken isa l]
codeTokens :: [CodeToken (isa (Ref w)) Text]
codeTokens} =
    let mnemonics :: [isa (Ref w)]
mnemonics = [isa (Ref w)
m | Mnemonic isa (Ref w)
m <- [CodeToken (isa (Ref w)) Text]
codeTokens]
        marked :: [(w, isa (Ref w))]
        marked :: [(w, isa (Ref w))]
marked = w -> [isa (Ref w)] -> [(w, isa (Ref w))]
forall t w. (ByteSize t, IsWord w) => w -> [t] -> [(w, t)]
markupOffsets w
offset [isa (Ref w)]
mnemonics
     in Section (isa (Ref w)) w Text
code
            { codeTokens =
                map
                    ( \(w
offset', isa (Ref w)
m) ->
                        let m' :: isa w
m' = (Text -> Maybe w) -> w -> isa (Ref w) -> isa w
forall (m :: * -> *) w.
DerefMnemonic m w =>
(Text -> Maybe w) -> w -> m (Ref w) -> m w
derefMnemonic Text -> Maybe w
f w
offset' isa (Ref w)
m
                            -- Force every Ref-derived field of m' to WHNF so that
                            -- an unresolved label aborts translation here, not
                            -- lazily at execution when something happens to read
                            -- the value (see issue #143). Walking @show@ visits
                            -- every constructor field, which is enough since
                            -- @w@ is a machine word and forcing it to WHNF is
                            -- already full evaluation.
                            !Int
_ = String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (isa w -> String
forall b a. (Show a, IsString b) => a -> b
show isa w
m' :: String)
                         in isa w -> CodeToken (isa w) w
forall isa l. isa -> CodeToken isa l
Mnemonic isa w
m'
                    )
                    marked
            }
derefSection Text -> Maybe w
f w
_offset dt :: Section (isa (Ref w)) w Text
dt@Data{[DataToken w Text]
dataTokens :: forall isa w l. Section isa w l -> [DataToken w l]
dataTokens :: [DataToken w Text]
dataTokens} =
    Section (isa (Ref w)) w Text
dt
        { dataTokens =
            map
                ( \DataToken{Text
dtLabel :: Text
dtLabel :: forall w l. DataToken w l -> l
dtLabel, DataValue w
dtValue :: DataValue w
dtValue :: forall w l. DataToken w l -> DataValue w
dtValue} ->
                    DataToken
                        { dtLabel :: w
dtLabel = w -> Maybe w -> w
forall a. a -> Maybe a -> a
fromMaybe (Text -> w
forall a t. (HasCallStack, IsText t) => t -> a
error (Text -> w) -> Text -> w
forall a b. (a -> b) -> a -> b
$ Text
"unknown label: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
forall b a. (Show a, IsString b) => a -> b
show Text
dtLabel) (Maybe w -> w) -> Maybe w -> w
forall a b. (a -> b) -> a -> b
$ Text -> Maybe w
f Text
dtLabel
                        , dtValue :: DataValue w
dtValue = DataValue w
dtValue
                        }
                )
                dataTokens
        }

markupOffsets :: (ByteSize t, IsWord w) => w -> [t] -> [(w, t)]
markupOffsets :: forall t w. (ByteSize t, IsWord w) => w -> [t] -> [(w, t)]
markupOffsets w
_offset [] = []
markupOffsets w
offset (t
m : [t]
ms) = (w
offset, t
m) (w, t) -> [(w, t)] -> [(w, t)]
forall a. a -> [a] -> [a]
: w -> [t] -> [(w, t)]
forall t w. (ByteSize t, IsWord w) => w -> [t] -> [(w, t)]
markupOffsets (w
offset w -> w -> w
forall a. Num a => a -> a -> a
+ Int -> w
forall a. Enum a => Int -> a
toEnum (t -> Int
forall t. ByteSize t => t -> Int
byteSize t
m)) [t]
ms

markupSectionOffsets :: (ByteSize isa, IsWord w) => w -> [Section isa w l] -> [(w, Section isa w l)]
markupSectionOffsets :: forall isa w l.
(ByteSize isa, IsWord w) =>
w -> [Section isa w l] -> [(w, Section isa w l)]
markupSectionOffsets w
_offset [] = []
markupSectionOffsets w
offset (Section isa w l
s : [Section isa w l]
ss) =
    let offset' :: w
offset' = w -> (Int -> w) -> Maybe Int -> w
forall b a. b -> (a -> b) -> Maybe a -> b
Prelude.maybe w
offset Int -> w
forall a. Enum a => Int -> a
toEnum (Section isa w l -> Maybe Int
forall isa w l. Section isa w l -> Maybe Int
org Section isa w l
s)
     in (w
offset', Section isa w l
s) (w, Section isa w l)
-> [(w, Section isa w l)] -> [(w, Section isa w l)]
forall a. a -> [a] -> [a]
: w -> [Section isa w l] -> [(w, Section isa w l)]
forall isa w l.
(ByteSize isa, IsWord w) =>
w -> [Section isa w l] -> [(w, Section isa w l)]
markupSectionOffsets (w
offset' w -> w -> w
forall a. Num a => a -> a -> a
+ Int -> w
forall a. Enum a => Int -> a
toEnum (Section isa w l -> Int
forall t. ByteSize t => t -> Int
byteSize Section isa w l
s)) [Section isa w l]
ss

data CodeToken isa l
    = Label l
    | Mnemonic isa
    deriving (Int -> CodeToken isa l -> ShowS
[CodeToken isa l] -> ShowS
CodeToken isa l -> String
(Int -> CodeToken isa l -> ShowS)
-> (CodeToken isa l -> String)
-> ([CodeToken isa l] -> ShowS)
-> Show (CodeToken isa l)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall isa l. (Show l, Show isa) => Int -> CodeToken isa l -> ShowS
forall isa l. (Show l, Show isa) => [CodeToken isa l] -> ShowS
forall isa l. (Show l, Show isa) => CodeToken isa l -> String
$cshowsPrec :: forall isa l. (Show l, Show isa) => Int -> CodeToken isa l -> ShowS
showsPrec :: Int -> CodeToken isa l -> ShowS
$cshow :: forall isa l. (Show l, Show isa) => CodeToken isa l -> String
show :: CodeToken isa l -> String
$cshowList :: forall isa l. (Show l, Show isa) => [CodeToken isa l] -> ShowS
showList :: [CodeToken isa l] -> ShowS
Show)

instance (ByteSize isa) => ByteSize (CodeToken isa l) where
    byteSize :: CodeToken isa l -> Int
byteSize (Mnemonic isa
m) = isa -> Int
forall t. ByteSize t => t -> Int
byteSize isa
m
    byteSize CodeToken isa l
_ = Int
0

data Ref w
    = Ref
        { forall w. Ref w -> w -> w
refPrepare :: w -> w
        , forall w. Ref w -> SourcePos
refPos :: SourcePos
        , forall w. Ref w -> Text
refLabel :: Text
        }
    | ValueR
        { refPrepare :: w -> w
        , refPos :: SourcePos
        , forall w. Ref w -> w
refValue :: w
        }

instance (Eq w) => Eq (Ref w) where
    Ref{refLabel :: forall w. Ref w -> Text
refLabel = Text
l} == :: Ref w -> Ref w -> Bool
== Ref{refLabel :: forall w. Ref w -> Text
refLabel = Text
l'} = Text
l Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
l'
    ValueR{refValue :: forall w. Ref w -> w
refValue = w
x} == ValueR{refValue :: forall w. Ref w -> w
refValue = w
x'} = w
x w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
x'
    Ref w
_ == Ref w
_ = Bool
False

instance (Show w) => Show (Ref w) where
    show :: Ref w -> String
show Ref{Text
refLabel :: forall w. Ref w -> Text
refLabel :: Text
refLabel} = Text -> String
forall a. ToString a => a -> String
toString Text
refLabel
    show ValueR{w -> w
refPrepare :: forall w. Ref w -> w -> w
refPrepare :: w -> w
refPrepare, w
refValue :: forall w. Ref w -> w
refValue :: w
refValue} = w -> String
forall b a. (Show a, IsString b) => a -> b
show (w -> String) -> w -> String
forall a b. (a -> b) -> a -> b
$ w -> w
refPrepare w
refValue

-- | Resolve a 'Ref' against a label table and check the result against the
-- instruction field it is going to be encoded into. Strict, so that an
-- unresolved label aborts translation here rather than becoming a thunk that
-- only blows up later if something happens to read it. Either failure is
-- reported with the position the reference was written at.
resolveRef :: (w -> Either Text w) -> (Text -> Maybe w) -> Ref w -> w
resolveRef :: forall w. (w -> Either Text w) -> (Text -> Maybe w) -> Ref w -> w
resolveRef w -> Either Text w
valueGuard Text -> Maybe w
labelResolver Ref w
ref =
    let pos :: Text
pos = String -> Text
forall a. ToText a => a -> Text
toText (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ SourcePos -> String
sourcePosPretty (SourcePos -> String) -> SourcePos -> String
forall a b. (a -> b) -> a -> b
$ Ref w -> SourcePos
forall w. Ref w -> SourcePos
refPos Ref w
ref
        !value :: w
value = case Ref w
ref of
            Ref{w -> w
refPrepare :: forall w. Ref w -> w -> w
refPrepare :: w -> w
refPrepare, Text
refLabel :: forall w. Ref w -> Text
refLabel :: Text
refLabel}
                | Just w
x <- Text -> Maybe w
labelResolver Text
refLabel -> w -> w
refPrepare w
x
                | Bool
otherwise -> Text -> w
forall a t. (HasCallStack, IsText t) => t -> a
error (Text
pos Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": can't resolve label: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
forall b a. (Show a, IsString b) => a -> b
show Text
refLabel)
            ValueR{w -> w
refPrepare :: forall w. Ref w -> w -> w
refPrepare :: w -> w
refPrepare, w
refValue :: forall w. Ref w -> w
refValue :: w
refValue} -> w -> w
refPrepare w
refValue
     in case w -> Either Text w
valueGuard w
value of
            Right w
x -> w
x
            Left Text
err -> Text -> w
forall a t. (HasCallStack, IsText t) => t -> a
error (Text
pos Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
err)

data DataToken w l = DataToken
    { forall w l. DataToken w l -> l
dtLabel :: !l
    , forall w l. DataToken w l -> DataValue w
dtValue :: DataValue w
    }
    deriving (Int -> DataToken w l -> ShowS
[DataToken w l] -> ShowS
DataToken w l -> String
(Int -> DataToken w l -> ShowS)
-> (DataToken w l -> String)
-> ([DataToken w l] -> ShowS)
-> Show (DataToken w l)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall w l. (Show l, Show w) => Int -> DataToken w l -> ShowS
forall w l. (Show l, Show w) => [DataToken w l] -> ShowS
forall w l. (Show l, Show w) => DataToken w l -> String
$cshowsPrec :: forall w l. (Show l, Show w) => Int -> DataToken w l -> ShowS
showsPrec :: Int -> DataToken w l -> ShowS
$cshow :: forall w l. (Show l, Show w) => DataToken w l -> String
show :: DataToken w l -> String
$cshowList :: forall w l. (Show l, Show w) => [DataToken w l] -> ShowS
showList :: [DataToken w l] -> ShowS
Show)

instance (ByteSizeT w) => ByteSize (DataToken w l) where
    byteSize :: DataToken w l -> Int
byteSize DataToken{DataValue w
dtValue :: forall w l. DataToken w l -> DataValue w
dtValue :: DataValue w
dtValue} = DataValue w -> Int
forall t. ByteSize t => t -> Int
byteSize DataValue w
dtValue

data DataValue w
    = DByte [Word8]
    | DWord [w]
    deriving (Int -> DataValue w -> ShowS
[DataValue w] -> ShowS
DataValue w -> String
(Int -> DataValue w -> ShowS)
-> (DataValue w -> String)
-> ([DataValue w] -> ShowS)
-> Show (DataValue w)
forall w. Show w => Int -> DataValue w -> ShowS
forall w. Show w => [DataValue w] -> ShowS
forall w. Show w => DataValue w -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall w. Show w => Int -> DataValue w -> ShowS
showsPrec :: Int -> DataValue w -> ShowS
$cshow :: forall w. Show w => DataValue w -> String
show :: DataValue w -> String
$cshowList :: forall w. Show w => [DataValue w] -> ShowS
showList :: [DataValue w] -> ShowS
Show)

instance (ByteSizeT w) => ByteSize (DataValue w) where
    byteSize :: DataValue w -> Int
byteSize (DByte [Word8]
xs) = [Word8] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
xs
    byteSize (DWord [w]
xs) = forall t. ByteSizeT t => Int
byteSizeT @w Int -> Int -> Int
forall a. Num a => a -> a -> a
* [w] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [w]
xs