Add PrimCallConv support to GHCiwip/ghci-primcall

This adds support for calling Cmm code from bytecode using the native
calling convention, allowing modules that use `foreign import prim`
to be loaded and debugged in GHCi.

This patch introduces a new `PRIMCALL` bytecode instruction and
a helper stack frame `stg_primcall`. The code is based on the
existing functionality for dealing with unboxed tuples in bytecode,
which has been generalised to handle arbitrary calls.

Fixes #22051

3 files changed, 71 insertions, 52 deletions

diff --git a/compiler/GHC/ByteCode/Asm.hs b/compiler/GHC/ByteCode/Asm.hs
index 24e2645052..391949d448 100644
--- a/compiler/GHC/ByteCode/Asm.hs
+++ b/compiler/GHC/ByteCode/Asm.hs
@@ -12,7 +12,7 @@ module GHC.ByteCode.Asm (
         bcoFreeNames,
         SizedSeq, sizeSS, ssElts,
         iNTERP_STACK_CHECK_THRESH,
-        mkTupleInfoLit
+        mkNativeCallInfoLit
   ) where
 
 import GHC.Prelude
@@ -32,7 +32,6 @@ import GHC.Types.Unique.DSet
 
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
-import GHC.Utils.Panic.Plain
 
 import GHC.Core.TyCon
 import GHC.Data.FastString
@@ -40,7 +39,7 @@ import GHC.Data.SizedSeq
 
 import GHC.StgToCmm.Layout     ( ArgRep(..) )
 import GHC.Cmm.Expr
-import GHC.Cmm.CallConv        ( tupleRegsCover )
+import GHC.Cmm.CallConv        ( allArgRegsCover )
 import GHC.Platform
 import GHC.Platform.Profile
 
@@ -202,7 +201,8 @@ assembleBCO platform (ProtoBCO { protoBCOName       = nm
   (final_insns, final_lits, final_ptrs) <- flip execStateT initial_state $ runAsm platform long_jumps env asm
 
   -- precomputed size should be equal to final size
-  massert (n_insns == sizeSS final_insns)
+  massertPpr (n_insns == sizeSS final_insns)
+             (text "bytecode instruction count mismatch")
 
   let asm_insns = ssElts final_insns
       insns_arr = Array.listArray (0, fromIntegral n_insns - 1) asm_insns
@@ -351,7 +351,8 @@ largeArg platform w = case platformWordSize platform of
            fromIntegral (w `shiftR` 32),
            fromIntegral (w `shiftR` 16),
            fromIntegral w]
-   PW4 -> assert (w < fromIntegral (maxBound :: Word32)) $
+   PW4 -> assertPpr (w < fromIntegral (maxBound :: Word32))
+                    (text "largeArg too big:" <+> ppr w) $
           [fromIntegral (w `shiftR` 16),
            fromIntegral w]
 
@@ -388,14 +389,14 @@ assembleI platform i = case i of
                            -> do let ul_bco = assembleBCO platform proto
                                  p <- ioptr (liftM BCOPtrBCO ul_bco)
                                  emit (push_alts pk) [Op p]
-  PUSH_ALTS_TUPLE proto tuple_info tuple_proto
+  PUSH_ALTS_TUPLE proto call_info tuple_proto
                            -> do let ul_bco = assembleBCO platform proto
                                      ul_tuple_bco = assembleBCO platform
                                                                 tuple_proto
                                  p <- ioptr (liftM BCOPtrBCO ul_bco)
                                  p_tup <- ioptr (liftM BCOPtrBCO ul_tuple_bco)
                                  info <- int (fromIntegral $
-                                              mkTupleInfoSig platform tuple_info)
+                                              mkNativeCallInfoSig platform call_info)
                                  emit bci_PUSH_ALTS_T
                                       [Op p, Op info, Op p_tup]
   PUSH_PAD8                -> emit bci_PUSH_PAD8 []
@@ -491,6 +492,7 @@ assembleI platform i = case i of
   RETURN_TUPLE             -> emit bci_RETURN_T []
   CCALL off m_addr i       -> do np <- addr m_addr
                                  emit bci_CCALL [SmallOp off, Op np, SmallOp i]
+  PRIMCALL                 -> emit bci_PRIMCALL []
   BRK_FUN index uniq cc    -> do p1 <- ptr BCOPtrBreakArray
                                  q <- int (getKey uniq)
                                  np <- addr cc
@@ -580,41 +582,44 @@ return_unlifted V64 = error "return_unlifted: vector"
   maximum number of tuple elements may be larger. Elements can also
   take multiple words on the stack (for example Double# on a 32 bit
   platform).
-
  -}
-maxTupleNativeStackSize :: WordOff
-maxTupleNativeStackSize = 62
+maxTupleReturnNativeStackSize :: WordOff
+maxTupleReturnNativeStackSize = 62
 
 {-
-  Construct the tuple_info word that stg_ctoi_t and stg_ret_t use
-  to convert a tuple between the native calling convention and the
+  Construct the call_info word that stg_ctoi_t, stg_ret_t and stg_primcall
+  use to convert arguments between the native calling convention and the
   interpreter.
 
-  See Note [GHCi tuple layout] for more information.
+  See Note [GHCi and native call registers] for more information.
  -}
-mkTupleInfoSig :: Platform -> TupleInfo -> Word32
-mkTupleInfoSig platform TupleInfo{..}
-  | tupleNativeStackSize > maxTupleNativeStackSize
-  = pprPanic "mkTupleInfoSig: tuple too big for the bytecode compiler"
-             (ppr tupleNativeStackSize <+> text "stack words." <+>
+mkNativeCallInfoSig :: Platform -> NativeCallInfo -> Word32
+mkNativeCallInfoSig platform NativeCallInfo{..}
+  | nativeCallType == NativeTupleReturn && nativeCallStackSpillSize > maxTupleReturnNativeStackSize
+  = pprPanic "mkNativeCallInfoSig: tuple too big for the bytecode compiler"
+             (ppr nativeCallStackSpillSize <+> text "stack words." <+>
               text "Use -fobject-code to get around this limit"
              )
   | otherwise
-  = assert (length regs <= 24) {- 24 bits for bitmap -}
-    assert (tupleNativeStackSize < 255) {- 8 bits for stack size -}
-    assert (all (`elem` regs) (regSetToList tupleRegs)) {- all regs accounted for -}
-    foldl' reg_bit 0 (zip regs [0..]) .|.
-      (fromIntegral tupleNativeStackSize `shiftL` 24)
+  = assertPpr (length regs <= 24) (text "too many registers for bitmap:" <+> ppr (length regs)) {- 24 bits for register bitmap -}
+    assertPpr (cont_offset < 255) (text "continuation offset too large:" <+> ppr cont_offset) {- 8 bits for continuation offset (only for NativeTupleReturn) -}
+    assertPpr (all (`elem` regs) (regSetToList nativeCallRegs)) (text "not all registers accounted for") {- all regs accounted for -}
+    foldl' reg_bit 0 (zip regs [0..]) .|. (cont_offset `shiftL` 24)
   where
+    cont_offset :: Word32
+    cont_offset
+      | nativeCallType == NativeTupleReturn = fromIntegral nativeCallStackSpillSize
+      | otherwise                           = 0 -- there is no continuation for primcalls
+
     reg_bit :: Word32 -> (GlobalReg, Int) -> Word32
     reg_bit x (r, n)
-      | r `elemRegSet` tupleRegs = x .|. 1 `shiftL` n
-      | otherwise                = x
-    regs = tupleRegsCover platform
+      | r `elemRegSet` nativeCallRegs = x .|. 1 `shiftL` n
+      | otherwise                     = x
+    regs = allArgRegsCover platform
 
-mkTupleInfoLit :: Platform -> TupleInfo -> Literal
-mkTupleInfoLit platform tuple_info =
-  mkLitWord platform . fromIntegral $ mkTupleInfoSig platform tuple_info
+mkNativeCallInfoLit :: Platform -> NativeCallInfo -> Literal
+mkNativeCallInfoLit platform call_info =
+  mkLitWord platform . fromIntegral $ mkNativeCallInfoSig platform call_info
 
 -- Make lists of host-sized words for literals, so that when the
 -- words are placed in memory at increasing addresses, the
diff --git a/compiler/GHC/ByteCode/Instr.hs b/compiler/GHC/ByteCode/Instr.hs
index 498152c471..34baa57d40 100644
--- a/compiler/GHC/ByteCode/Instr.hs
+++ b/compiler/GHC/ByteCode/Instr.hs
@@ -90,7 +90,7 @@ data BCInstr
    | PUSH_ALTS          (ProtoBCO Name)
    | PUSH_ALTS_UNLIFTED (ProtoBCO Name) ArgRep
    | PUSH_ALTS_TUPLE    (ProtoBCO Name) -- continuation
-                        !TupleInfo
+                        !NativeCallInfo
                         (ProtoBCO Name) -- tuple return BCO
 
    -- Pushing 8, 16 and 32 bits of padding (for constructors).
@@ -184,6 +184,8 @@ data BCInstr
                                 -- (XXX: inefficient, but I don't know
                                 -- what the alignment constraints are.)
 
+   | PRIMCALL
+
    -- For doing magic ByteArray passing to foreign calls
    | SWIZZLE          Word16 -- to the ptr N words down the stack,
                       Word16 -- add M (interpreted as a signed 16-bit entity)
@@ -269,8 +271,8 @@ instance Outputable BCInstr where
 
    ppr (PUSH_ALTS bco)       = hang (text "PUSH_ALTS") 2 (ppr bco)
    ppr (PUSH_ALTS_UNLIFTED bco pk) = hang (text "PUSH_ALTS_UNLIFTED" <+> ppr pk) 2 (ppr bco)
-   ppr (PUSH_ALTS_TUPLE bco tuple_info tuple_bco) =
-                               hang (text "PUSH_ALTS_TUPLE" <+> ppr tuple_info)
+   ppr (PUSH_ALTS_TUPLE bco call_info tuple_bco) =
+                               hang (text "PUSH_ALTS_TUPLE" <+> ppr call_info)
                                     2
                                     (ppr tuple_bco $+$ ppr bco)
 
@@ -340,6 +342,7 @@ instance Outputable BCInstr where
                                                       0x1 -> text "(interruptible)"
                                                       0x2 -> text "(unsafe)"
                                                       _   -> empty)
+   ppr PRIMCALL              = text "PRIMCALL"
    ppr (SWIZZLE stkoff n)    = text "SWIZZLE " <+> text "stkoff" <+> ppr stkoff
                                                <+> text "by" <+> ppr n
    ppr ENTER                 = text "ENTER"
@@ -382,11 +385,11 @@ bciStackUse (PUSH_ALTS bco)       = 2 {- profiling only, restore CCCS -} +
 bciStackUse (PUSH_ALTS_UNLIFTED bco _) = 2 {- profiling only, restore CCCS -} +
                                          4 + protoBCOStackUse bco
 bciStackUse (PUSH_ALTS_TUPLE bco info _) =
-   -- (tuple_bco, tuple_info word, cont_bco, stg_ctoi_t)
+   -- (tuple_bco, call_info word, cont_bco, stg_ctoi_t)
    -- tuple
-   -- (tuple_info, tuple_bco, stg_ret_t)
+   -- (call_info, tuple_bco, stg_ret_t)
    1 {- profiling only -} +
-   7 + fromIntegral (tupleSize info) + protoBCOStackUse bco
+   7 + fromIntegral (nativeCallSize info) + protoBCOStackUse bco
 bciStackUse (PUSH_PAD8)           = 1  -- overapproximation
 bciStackUse (PUSH_PAD16)          = 1  -- overapproximation
 bciStackUse (PUSH_PAD32)          = 1  -- overapproximation on 64bit arch
@@ -443,6 +446,7 @@ bciStackUse RETURN{}              = 0
 bciStackUse RETURN_UNLIFTED{}     = 1 -- pushes stg_ret_X for some X
 bciStackUse RETURN_TUPLE{}        = 1 -- pushes stg_ret_t header
 bciStackUse CCALL{}               = 0
+bciStackUse PRIMCALL{}            = 1 -- pushes stg_primcall
 bciStackUse SWIZZLE{}             = 0
 bciStackUse BRK_FUN{}             = 0
 
diff --git a/compiler/GHC/ByteCode/Types.hs b/compiler/GHC/ByteCode/Types.hs
index 830b60a4ca..a4b025ce92 100644
--- a/compiler/GHC/ByteCode/Types.hs
+++ b/compiler/GHC/ByteCode/Types.hs
@@ -10,7 +10,7 @@ module GHC.ByteCode.Types
   ( CompiledByteCode(..), seqCompiledByteCode
   , FFIInfo(..)
   , RegBitmap(..)
-  , TupleInfo(..), voidTupleInfo
+  , NativeCallType(..), NativeCallInfo(..), voidTupleReturnInfo, voidPrimCallInfo
   , ByteOff(..), WordOff(..)
   , UnlinkedBCO(..), BCOPtr(..), BCONPtr(..)
   , ItblEnv, ItblPtr(..)
@@ -105,22 +105,32 @@ newtype RegBitmap = RegBitmap { unRegBitmap :: Word32 }
 
    See GHC.StgToByteCode.layoutTuple for more details.
 -}
-data TupleInfo = TupleInfo
-  { tupleSize            :: !WordOff   -- total size of tuple in words
-  , tupleRegs            :: !GlobalRegSet
-  , tupleNativeStackSize :: !WordOff {- words spilled on the stack by
-                                        GHCs native calling convention -}
-  } deriving (Show)
-
-instance Outputable TupleInfo where
-  ppr TupleInfo{..} = text "<size" <+> ppr tupleSize <+>
-                      text "stack" <+> ppr tupleNativeStackSize <+>
-                      text "regs"  <+>
-                      ppr (map (text @SDoc . show) $ regSetToList tupleRegs) <>
-                      char '>'
-
-voidTupleInfo :: TupleInfo
-voidTupleInfo = TupleInfo 0 emptyRegSet 0
+
+data NativeCallType = NativePrimCall
+                    | NativeTupleReturn
+  deriving (Eq)
+
+data NativeCallInfo = NativeCallInfo
+  { nativeCallType           :: !NativeCallType
+  , nativeCallSize           :: !WordOff   -- total size of arguments in words
+  , nativeCallRegs           :: !GlobalRegSet
+  , nativeCallStackSpillSize :: !WordOff {- words spilled on the stack by
+                                            GHCs native calling convention -}
+  }
+
+instance Outputable NativeCallInfo where
+  ppr NativeCallInfo{..} = text "<arg_size" <+> ppr nativeCallSize <+>
+                           text "stack" <+> ppr nativeCallStackSpillSize <+>
+                           text "regs"  <+>
+                           ppr (map (text @SDoc . show) $ regSetToList nativeCallRegs) <>
+                           char '>'
+
+
+voidTupleReturnInfo :: NativeCallInfo
+voidTupleReturnInfo = NativeCallInfo NativeTupleReturn 0 emptyRegSet 0
+
+voidPrimCallInfo :: NativeCallInfo
+voidPrimCallInfo = NativeCallInfo NativePrimCall 0 emptyRegSet 0
 
 type ItblEnv = NameEnv (Name, ItblPtr)
         -- We need the Name in the range so we know which