Asm.hs

{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE Strict #-}

module GHC.CmmToAsm.Wasm.Asm (asmTellEverything, execWasmAsmM) where

import Control.Monad
import Control.Monad.Trans.Reader
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.ByteString.Builder
import qualified Data.ByteString.Char8 as BS8
import Data.Coerce
import Data.Foldable
import Data.Maybe
import Data.Semigroup
import GHC.Cmm
import GHC.CmmToAsm.Ppr
import GHC.CmmToAsm.Wasm.FromCmm
import GHC.CmmToAsm.Wasm.Types
import GHC.CmmToAsm.Wasm.Utils
import GHC.Data.FastString
import GHC.Float
import GHC.Prelude
import GHC.Settings.Config (cProjectVersion)
import GHC.Types.Basic
import GHC.Types.Unique
import GHC.Types.Unique.Map
import GHC.Types.Unique.Set
import GHC.Utils.Monad.State.Strict
import GHC.Utils.Outputable hiding ((<>))
import GHC.Utils.Panic (panic)

-- | Reads current indentation, appends result to state
newtype WasmAsmM a = WasmAsmM (WasmAsmConfig -> Builder -> State Builder a)
  deriving
    ( Functor,
      Applicative,
      Monad
    )
    via (ReaderT WasmAsmConfig (ReaderT Builder (State Builder)))

instance Semigroup a => Semigroup (WasmAsmM a) where
  (<>) = liftA2 (<>)

instance Monoid a => Monoid (WasmAsmM a) where
  mempty = pure mempty

getConf :: WasmAsmM WasmAsmConfig
getConf = WasmAsmM $ \conf _ -> pure conf

-- | Default indent level is none
execWasmAsmM :: WasmAsmConfig -> WasmAsmM a -> Builder
execWasmAsmM conf (WasmAsmM m) =
  execState (m conf mempty) mempty

-- | Increase indent level by a tab
asmWithTab :: WasmAsmM a -> WasmAsmM a
asmWithTab (WasmAsmM m) =
  WasmAsmM $ \conf t -> m conf $! char7 '\t' <> t

-- | Writes a single line starting with the current indent
asmTellLine :: Builder -> WasmAsmM ()
asmTellLine b = WasmAsmM $ \_ t -> modify $ \acc -> acc <> t <> b <> char7 '\n'

-- | Writes a single line break
asmTellLF :: WasmAsmM ()
asmTellLF = WasmAsmM $ \_ _ -> modify $ \acc -> acc <> char7 '\n'

-- | Writes a line starting with a single tab, ignoring current indent
-- level
asmTellTabLine :: Builder -> WasmAsmM ()
asmTellTabLine b =
  WasmAsmM $ \_ _ -> modify $ \acc -> acc <> char7 '\t' <> b <> char7 '\n'

asmFromWasmType :: WasmTypeTag t -> Builder
asmFromWasmType ty = case ty of
  TagI32 -> "i32"
  TagI64 -> "i64"
  TagF32 -> "f32"
  TagF64 -> "f64"

asmFromSomeWasmType :: SomeWasmType -> Builder
asmFromSomeWasmType (SomeWasmType t) = asmFromWasmType t

asmFromSomeWasmTypes :: [SomeWasmType] -> Builder
asmFromSomeWasmTypes ts = "(" <> builderCommas asmFromSomeWasmType ts <> ")"

asmFromFuncType :: [SomeWasmType] -> [SomeWasmType] -> Builder
asmFromFuncType arg_tys ret_tys =
  asmFromSomeWasmTypes arg_tys <> " -> " <> asmFromSomeWasmTypes ret_tys

asmTellFuncType ::
  SymName -> ([SomeWasmType], [SomeWasmType]) -> WasmAsmM ()
asmTellFuncType sym (arg_tys, ret_tys) =
  asmTellTabLine $
    ".functype "
      <> asmFromSymName sym
      <> " "
      <> asmFromFuncType arg_tys ret_tys

asmTellLocals :: [SomeWasmType] -> WasmAsmM ()
asmTellLocals [] = mempty
asmTellLocals local_tys =
  asmTellTabLine $ ".local " <> builderCommas asmFromSomeWasmType local_tys

asmFromSymName :: SymName -> Builder
asmFromSymName = shortByteString . coerce fastStringToShortByteString

asmTellDefSym :: SymName -> WasmAsmM ()
asmTellDefSym sym = do
  WasmAsmConfig {..} <- getConf
  unless pic $ asmTellTabLine $ ".hidden " <> asm_sym
  asmTellTabLine $ ".globl " <> asm_sym
  where
    asm_sym = asmFromSymName sym

asmTellDataSectionContent :: WasmTypeTag w -> DataSectionContent -> WasmAsmM ()
asmTellDataSectionContent ty_word c = asmTellTabLine $ case c of
  DataI8 i -> ".int8 0x" <> word8Hex i
  DataI16 i -> ".int16 0x" <> word16Hex i
  DataI32 i -> ".int32 0x" <> word32Hex i
  DataI64 i -> ".int64 0x" <> word64Hex i
  DataF32 f -> ".int32 0x" <> word32Hex (castFloatToWord32 f)
  DataF64 d -> ".int64 0x" <> word64Hex (castDoubleToWord64 d)
  DataSym sym o ->
    ( case ty_word of
        TagI32 -> ".int32 "
        TagI64 -> ".int64 "
        _ -> panic "asmTellDataSectionContent: unreachable"
    )
      <> asmFromSymName sym
      <> ( case compare o 0 of
             EQ -> mempty
             GT -> "+" <> intDec o
             LT -> panic "asmTellDataSectionContent: negative offset"
         )
  DataSkip i -> ".skip " <> intDec i
  DataASCII s
    | not (BS.null s) && BS.last s == 0 ->
        ".asciz \""
          <> string7
            (showSDocOneLine defaultSDocContext $ pprASCII $ BS.init s)
          <> "\""
    | otherwise ->
        ".ascii \""
          <> string7
            (showSDocOneLine defaultSDocContext $ pprASCII s)
          <> "\""
  DataIncBin f _ ->
    ".incbin "
      <> string7
        (showSDocOneLine defaultSDocContext $ pprFilePathString f)

dataSectionContentSize :: WasmTypeTag w -> DataSectionContent -> Int
dataSectionContentSize ty_word c = case c of
  DataI8 {} -> 1
  DataI16 {} -> 2
  DataI32 {} -> 4
  DataI64 {} -> 8
  DataF32 {} -> 4
  DataF64 {} -> 8
  DataSym {} -> alignmentBytes $ alignmentFromWordType ty_word
  DataSkip i -> i
  DataASCII s -> BS.length s
  DataIncBin _ l -> l

dataSectionSize :: WasmTypeTag w -> [DataSectionContent] -> Int
dataSectionSize ty_word =
  coerce
    . foldMap'
      (Sum . dataSectionContentSize ty_word)

asmTellAlign :: Alignment -> WasmAsmM ()
asmTellAlign a = case alignmentBytes a of
  1 -> mempty
  i -> asmTellTabLine $ ".p2align " <> intDec (countTrailingZeros i)

asmTellSectionHeader :: Builder -> WasmAsmM ()
asmTellSectionHeader k = asmTellTabLine $ ".section " <> k <> ",\"\",@"

asmTellDataSection ::
  WasmTypeTag w -> UniqueSet -> SymName -> DataSection -> WasmAsmM ()
asmTellDataSection ty_word def_syms sym DataSection {..} = do
  when (getUnique sym `memberUniqueSet` def_syms) $ asmTellDefSym sym
  asmTellSectionHeader sec_name
  asmTellAlign dataSectionAlignment
  asmTellTabLine asm_size
  asmTellLine $ asm_sym <> ":"
  for_ dataSectionContents $ asmTellDataSectionContent ty_word
  asmTellLF
  where
    asm_sym = asmFromSymName sym

    sec_name =
      ( case dataSectionKind of
          SectionData -> ".data."
          SectionROData -> ".rodata."
      )
        <> asm_sym

    asm_size =
      ".size "
        <> asm_sym
        <> ", "
        <> intDec
          (dataSectionSize ty_word dataSectionContents)

asmFromWasmBlockType :: WasmTypeTag w -> WasmFunctionType pre post -> Builder
asmFromWasmBlockType
  _
  (WasmFunctionType {ft_pops = TypeListNil, ft_pushes = TypeListNil}) =
    mempty
asmFromWasmBlockType
  TagI32
  ( WasmFunctionType
      { ft_pops = TypeListNil,
        ft_pushes = TypeListCons TagI32 TypeListNil
      }
    ) =
    " i32"
asmFromWasmBlockType
  TagI64
  ( WasmFunctionType
      { ft_pops = TypeListNil,
        ft_pushes = TypeListCons TagI64 TypeListNil
      }
    ) =
    " i64"
asmFromWasmBlockType _ _ = panic "asmFromWasmBlockType: invalid block type"

asmFromAlignmentSpec :: AlignmentSpec -> Builder
asmFromAlignmentSpec NaturallyAligned = mempty
asmFromAlignmentSpec Unaligned = ":p2align=0"

asmTellWasmInstr :: WasmTypeTag w -> WasmInstr w pre post -> WasmAsmM ()
asmTellWasmInstr ty_word instr = case instr of
  WasmComment c -> asmTellLine $ stringUtf8 $ "# " <> c
  WasmNop -> mempty
  WasmDrop -> asmTellLine "drop"
  WasmUnreachable -> asmTellLine "unreachable"
  WasmConst TagI32 i -> asmTellLine $ "i32.const " <> integerDec i
  WasmConst TagI64 i -> asmTellLine $ "i64.const " <> integerDec i
  WasmConst {} -> panic "asmTellWasmInstr: unreachable"
  WasmSymConst sym -> do
    WasmAsmConfig {..} <- getConf
    let
      asm_sym = asmFromSymName sym
      (ty_const, ty_add) = case ty_word of
        TagI32 -> ("i32.const ", "i32.add")
        TagI64 -> ("i64.const ", "i64.add")
        _ -> panic "asmTellWasmInstr: invalid word type"
    traverse_ asmTellLine $ if
      | pic, getUnique sym `memberUniqueSet` mbrelSyms -> [
          "global.get __memory_base",
          ty_const <> asm_sym <> "@MBREL",
          ty_add
        ]
      | pic, getUnique sym `memberUniqueSet` tbrelSyms -> [
          "global.get __table_base",
          ty_const <> asm_sym <> "@TBREL",
          ty_add
        ]
      | pic -> [ "global.get " <> asm_sym <> "@GOT" ]
      | otherwise -> [ ty_const <> asm_sym ]
  WasmLoad ty (Just w) s o align ->
    asmTellLine $
      asmFromWasmType ty
        <> ".load"
        <> intDec w
        <> ( case s of
               Signed -> "_s"
               Unsigned -> "_u"
           )
        <> " "
        <> intDec o
        <> asmFromAlignmentSpec align
  WasmLoad ty Nothing _ o align ->
    asmTellLine $
      asmFromWasmType ty
        <> ".load"
        <> " "
        <> intDec o
        <> asmFromAlignmentSpec align
  WasmStore ty (Just w) o align ->
    asmTellLine $
      asmFromWasmType ty
        <> ".store"
        <> intDec w
        <> " "
        <> intDec o
        <> asmFromAlignmentSpec align
  WasmStore ty Nothing o align ->
    asmTellLine $
      asmFromWasmType ty
        <> ".store"
        <> " "
        <> intDec o
        <> asmFromAlignmentSpec align
  WasmGlobalGet _ sym -> asmTellLine $ "global.get " <> asmFromSymName sym
  WasmGlobalSet _ sym -> asmTellLine $ "global.set " <> asmFromSymName sym
  WasmLocalGet _ i -> asmTellLine $ "local.get " <> intDec i
  WasmLocalSet _ i -> asmTellLine $ "local.set " <> intDec i
  WasmLocalTee _ i -> asmTellLine $ "local.tee " <> intDec i
  WasmCCall sym -> asmTellLine $ "call " <> asmFromSymName sym
  WasmCCallIndirect arg_tys ret_tys ->
    asmTellLine $
      "call_indirect "
        <> asmFromFuncType
          (someWasmTypesFromTypeList arg_tys)
          (someWasmTypesFromTypeList ret_tys)
  WasmConcat instr0 instr1 -> do
    asmTellWasmInstr ty_word instr0
    asmTellWasmInstr ty_word instr1
  WasmReinterpret t0 t1 ->
    asmTellLine $
      asmFromWasmType t1 <> ".reinterpret_" <> asmFromWasmType t0
  WasmTruncSat Signed t0 t1 ->
    asmTellLine $
      asmFromWasmType t1 <> ".trunc_sat_" <> asmFromWasmType t0 <> "_s"
  WasmTruncSat Unsigned t0 t1 ->
    asmTellLine $
      asmFromWasmType t1 <> ".trunc_sat_" <> asmFromWasmType t0 <> "_u"
  WasmConvert Signed t0 t1 ->
    asmTellLine $
      asmFromWasmType t1 <> ".convert_" <> asmFromWasmType t0 <> "_s"
  WasmConvert Unsigned t0 t1 ->
    asmTellLine $
      asmFromWasmType t1 <> ".convert_" <> asmFromWasmType t0 <> "_u"
  WasmAdd ty -> asmTellLine $ asmFromWasmType ty <> ".add"
  WasmSub ty -> asmTellLine $ asmFromWasmType ty <> ".sub"
  WasmMul ty -> asmTellLine $ asmFromWasmType ty <> ".mul"
  WasmDiv _ TagF32 -> asmTellLine "f32.div"
  WasmDiv _ TagF64 -> asmTellLine "f64.div"
  WasmDiv Signed ty -> asmTellLine $ asmFromWasmType ty <> ".div_s"
  WasmDiv Unsigned ty -> asmTellLine $ asmFromWasmType ty <> ".div_u"
  WasmRem Signed ty -> asmTellLine $ asmFromWasmType ty <> ".rem_s"
  WasmRem Unsigned ty -> asmTellLine $ asmFromWasmType ty <> ".rem_u"
  WasmAnd ty -> asmTellLine $ asmFromWasmType ty <> ".and"
  WasmOr ty -> asmTellLine $ asmFromWasmType ty <> ".or"
  WasmXor ty -> asmTellLine $ asmFromWasmType ty <> ".xor"
  WasmEq ty -> asmTellLine $ asmFromWasmType ty <> ".eq"
  WasmNe ty -> asmTellLine $ asmFromWasmType ty <> ".ne"
  WasmLt _ TagF32 -> asmTellLine "f32.lt"
  WasmLt _ TagF64 -> asmTellLine "f64.lt"
  WasmLt Signed ty -> asmTellLine $ asmFromWasmType ty <> ".lt_s"
  WasmLt Unsigned ty -> asmTellLine $ asmFromWasmType ty <> ".lt_u"
  WasmGt _ TagF32 -> asmTellLine "f32.gt"
  WasmGt _ TagF64 -> asmTellLine "f64.gt"
  WasmGt Signed ty -> asmTellLine $ asmFromWasmType ty <> ".gt_s"
  WasmGt Unsigned ty -> asmTellLine $ asmFromWasmType ty <> ".gt_u"
  WasmLe _ TagF32 -> asmTellLine "f32.le"
  WasmLe _ TagF64 -> asmTellLine "f64.le"
  WasmLe Signed ty -> asmTellLine $ asmFromWasmType ty <> ".le_s"
  WasmLe Unsigned ty -> asmTellLine $ asmFromWasmType ty <> ".le_u"
  WasmGe _ TagF32 -> asmTellLine "f32.ge"
  WasmGe _ TagF64 -> asmTellLine "f64.ge"
  WasmGe Signed ty -> asmTellLine $ asmFromWasmType ty <> ".ge_s"
  WasmGe Unsigned ty -> asmTellLine $ asmFromWasmType ty <> ".ge_u"
  WasmShl ty -> asmTellLine $ asmFromWasmType ty <> ".shl"
  WasmShr Signed ty -> asmTellLine $ asmFromWasmType ty <> ".shr_s"
  WasmShr Unsigned ty -> asmTellLine $ asmFromWasmType ty <> ".shr_u"
  WasmI32Extend8S -> asmTellLine "i32.extend8_s"
  WasmI32Extend16S -> asmTellLine "i32.extend16_s"
  WasmI64Extend8S -> asmTellLine "i64.extend8_s"
  WasmI64Extend16S -> asmTellLine "i64.extend16_s"
  WasmI64Extend32S -> asmTellLine "i64.extend32_s"
  WasmI64ExtendI32 Signed -> asmTellLine "i64.extend_i32_s"
  WasmI64ExtendI32 Unsigned -> asmTellLine "i64.extend_i32_u"
  WasmI32WrapI64 -> asmTellLine "i32.wrap_i64"
  WasmF32DemoteF64 -> asmTellLine "f32.demote_f64"
  WasmF64PromoteF32 -> asmTellLine "f64.promote_f32"
  WasmAbs ty -> asmTellLine $ asmFromWasmType ty <> ".abs"
  WasmSqrt ty -> asmTellLine $ asmFromWasmType ty <> ".sqrt"
  WasmNeg ty -> asmTellLine $ asmFromWasmType ty <> ".neg"
  WasmMin ty -> asmTellLine $ asmFromWasmType ty <> ".min"
  WasmMax ty -> asmTellLine $ asmFromWasmType ty <> ".max"
  WasmCond t -> do
    asmTellLine "if"
    asmWithTab $ asmTellWasmInstr ty_word t
    asmTellLine "end_if"

asmTellWasmControl ::
  WasmTypeTag w ->
  WasmControl
    (WasmStatements w)
    (WasmExpr w a)
    pre
    post ->
  WasmAsmM ()
asmTellWasmControl ty_word c = case c of
  WasmPush _ (WasmExpr e) -> asmTellWasmInstr ty_word e
  WasmBlock bt c -> do
    asmTellLine $ "block" <> asmFromWasmBlockType ty_word bt
    asmWithTab $ asmTellWasmControl ty_word c
    asmTellLine "end_block"
  WasmLoop bt c -> do
    asmTellLine $ "loop" <> asmFromWasmBlockType ty_word bt
    asmWithTab $ asmTellWasmControl ty_word c
    asmTellLine "end_loop"
  WasmIfTop bt t f -> do
    asmTellLine $ "if" <> asmFromWasmBlockType ty_word bt
    asmWithTab $ asmTellWasmControl ty_word t
    asmTellLine "else"
    asmWithTab $ asmTellWasmControl ty_word f
    asmTellLine "end_if"
  WasmBr i -> asmTellLine $ "br " <> intDec i
  WasmFallthrough -> mempty
  WasmBrTable (WasmExpr e) _ ts t -> do
    asmTellWasmInstr ty_word e
    asmTellLine $ "br_table {" <> builderCommas intDec (ts <> [t]) <> "}"
  -- See Note [WasmTailCall]
  WasmTailCall (WasmExpr e) -> do
    WasmAsmConfig {..} <- getConf
    if
        | tailcall,
          WasmSymConst sym <- e ->
            asmTellLine $ "return_call " <> asmFromSymName sym
        | tailcall ->
            do
              asmTellWasmInstr ty_word e
              asmTellLine $
                "return_call_indirect "
                  <> asmFromFuncType
                    []
                    [SomeWasmType ty_word]
        | otherwise ->
            do
              asmTellWasmInstr ty_word e
              asmTellLine "return"
  WasmActions (WasmStatements a) -> asmTellWasmInstr ty_word a
  WasmSeq c0 c1 -> do
    asmTellWasmControl ty_word c0
    asmTellWasmControl ty_word c1

asmTellFunc ::
  WasmTypeTag w ->
  UniqueSet ->
  SymName ->
  (([SomeWasmType], [SomeWasmType]), FuncBody w) ->
  WasmAsmM ()
asmTellFunc ty_word def_syms sym (func_ty, FuncBody {..}) = do
  when (getUnique sym `memberUniqueSet` def_syms) $ asmTellDefSym sym
  asmTellSectionHeader $ ".text." <> asm_sym
  asmTellLine $ asm_sym <> ":"
  asmTellFuncType sym func_ty
  asmTellLocals funcLocals
  asmWithTab $ asmTellWasmControl ty_word funcBody
  asmTellTabLine "end_function"
  asmTellLF
  where
    asm_sym = asmFromSymName sym

asmTellGlobals :: WasmTypeTag w -> WasmAsmM ()
asmTellGlobals ty_word = do
  WasmAsmConfig {..} <- getConf
  when pic $ traverse_ asmTellTabLine [
      ".globaltype __memory_base, i32, immutable",
      ".globaltype __table_base, i32, immutable"
    ]
  for_ supportedCmmGlobalRegs $ \reg ->
    let
      (sym, ty) = fromJust $ globalInfoFromCmmGlobalReg ty_word reg
      asm_sym = asmFromSymName sym
     in do
      asmTellTabLine $
          ".globaltype "
            <> asm_sym
            <> ", "
            <> asmFromSomeWasmType ty
      when pic $ traverse_ asmTellTabLine [
          ".import_module " <> asm_sym <> ", regs",
          ".import_name " <> asm_sym <> ", " <> asm_sym
        ]
  asmTellLF

asmTellCtors :: WasmTypeTag w -> [SymName] -> WasmAsmM ()
asmTellCtors _ [] = mempty
asmTellCtors ty_word syms = do
  -- See Note [JSFFI initialization] for details
  asmTellSectionHeader ".init_array.101"
  asmTellAlign $ alignmentFromWordType ty_word
  for_ syms $ \sym ->
    asmTellTabLine $
      ( case ty_word of
          TagI32 -> ".int32 "
          TagI64 -> ".int64 "
          _ -> panic "asmTellCtors: unreachable"
      )
        <> asmFromSymName sym
  asmTellLF

asmTellBS :: ByteString -> WasmAsmM ()
asmTellBS s = do
  asmTellTabLine $ ".int8 " <> intDec (BS.length s)
  asmTellTabLine $
    ".ascii \""
      <> string7
        (showSDocOneLine defaultSDocContext $ pprASCII s)
      <> "\""

asmTellVec :: [WasmAsmM ()] -> WasmAsmM ()
asmTellVec xs = do
  asmTellTabLine $ ".int8 " <> intDec (length xs)
  sequence_ xs

asmTellProducers :: WasmAsmM ()
asmTellProducers = do
  asmTellSectionHeader ".custom_section.producers"
  asmTellVec
    [ do
        asmTellBS "processed-by"
        asmTellVec
          [ do
              asmTellBS "ghc"
              asmTellBS $ BS8.pack cProjectVersion
          ]
    ]

asmTellTargetFeatures :: WasmAsmM ()
asmTellTargetFeatures = do
  WasmAsmConfig {..} <- getConf
  asmTellSectionHeader ".custom_section.target_features"
  asmTellVec
    [ do
        asmTellTabLine ".int8 0x2b"
        asmTellBS feature
      | feature <-
          ["tail-call" | tailcall]
            <> [ "bulk-memory",
                 "mutable-globals",
                 "nontrapping-fptoint",
                 "reference-types",
                 "sign-ext"
               ]
    ]

asmTellEverything :: WasmTypeTag w -> WasmCodeGenState w -> WasmAsmM ()
asmTellEverything ty_word WasmCodeGenState {..} = do
  asmTellGlobals ty_word
  asm_functypes
  asm_funcs
  asm_data_secs
  asm_ctors
  asmTellProducers
  asmTellTargetFeatures
  where
    asm_functypes = do
      for_
        (detEltsUniqMap $ funcTypes `minusUniqMap` funcBodies)
        (uncurry asmTellFuncType)
      asmTellLF

    asm_funcs = do
      for_
        (detEltsUniqMap $ intersectUniqMap_C (,) funcTypes funcBodies)
        (uncurry $ asmTellFunc ty_word defaultSyms)
      asmTellLF

    asm_data_secs = do
      for_
        (detEltsUniqMap dataSections)
        (uncurry (asmTellDataSection ty_word defaultSyms))
      asmTellLF

    asm_ctors = asmTellCtors ty_word ctors