Manually unbox functions in FF

unison-runtime/src/Unison/Runtime/Foreign/Function.hs

@@ -1,7 +1,9 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE ViewPatterns #-}
 
@@ -22,14 +24,15 @@ import Data.IORef (IORef)
 import Data.Sequence qualified as Sq
 import Data.Time.Clock.POSIX (POSIXTime)
 import Data.Word (Word16, Word32, Word64, Word8)
+import GHC.Base (IO (..))
 import GHC.IO.Exception (IOErrorType (..), IOException (..))
 import Network.Socket (Socket)
 import Network.UDP (UDPSocket)
 import System.IO (BufferMode (..), Handle, IOMode, SeekMode)
 import Unison.Builtin.Decls qualified as Ty
 import Unison.Reference (Reference)
-import Unison.Runtime.Array qualified as PA
 import Unison.Runtime.ANF (Code, PackedTag (..), Value, internalBug)
+import Unison.Runtime.Array qualified as PA
 import Unison.Runtime.Exception
 import Unison.Runtime.Foreign
 import Unison.Runtime.MCode
@@ -53,8 +56,8 @@ import Unison.Util.Text (Text, pack, unpack)
 -- Foreign functions operating on stacks
 data ForeignFunc where
   FF ::
-    (Stack -> Args -> IO a) ->
-    (Stack -> r -> IO Stack) ->
+    (XStack -> Args -> IO a) ->
+    (XStack -> r -> IOStack) ->
     (a -> IO r) ->
     ForeignFunc
 
@@ -74,12 +77,17 @@ class ForeignConvention a where
     Stack -> a -> IO Stack
 
 mkForeign ::
+  forall a r.
   (ForeignConvention a, ForeignConvention r) =>
   (a -> IO r) ->
   ForeignFunc
-mkForeign ev = FF readArgs writeForeign ev
+mkForeign ev = FF readArgs doWrite ev
   where
-    readArgs stk (argsToLists -> args) =
+    doWrite :: XStack -> r -> IOStack
+    doWrite stk a = case writeForeign (packXStack stk) a of
+      (IO f) -> \state -> case f state of
+        (# state', stk #) -> (# state', unpackXStack stk #)
+    readArgs (packXStack -> stk) (argsToLists -> args) =
       readForeign args stk >>= \case
         ([], a) -> pure a
         _ ->

unison-runtime/src/Unison/Runtime/Stack.hs

@@ -1,5 +1,8 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE UnboxedTuples #-}
 
 module Unison.Runtime.Stack
   ( K (..),
@@ -27,6 +30,16 @@ module Unison.Runtime.Stack
     Augment (..),
     Dump (..),
     Stack (..),
+    XStack,
+    pattern XStack,
+    packXStack,
+    unpackXStack,
+    IOStack,
+    apX,
+    fpX,
+    spX,
+    ustkX,
+    bstkX,
     Off,
     SZ,
     FP,
@@ -127,10 +140,10 @@ where
 
 import Control.Monad.Primitive
 import Data.Char qualified as Char
-import Data.Kind (Constraint)
 import Data.Primitive (sizeOf)
 import Data.Primitive.ByteArray qualified as BA
 import Data.Word
+import GHC.Base
 import GHC.Exts as L (IsList (..))
 import Unison.Prelude
 import Unison.Reference (Reference)
@@ -597,6 +610,26 @@ data Stack = Stack
     bstk :: {-# UNPACK #-} !(MutableArray (PrimState IO) Closure)
   }
 
+-- Unboxed representation of the Stack, used to force GHC optimizations in a few spots.
+type XStack = (# Int#, Int#, Int#, MutableByteArray# (PrimState IO), MutableArray# (PrimState IO) Closure #)
+
+type IOStack = State# RealWorld -> (# State# RealWorld, XStack #)
+
+pattern XStack :: Int# -> Int# -> Int# -> MutableByteArray# RealWorld -> MutableArray# RealWorld Closure -> Stack
+pattern XStack {apX, fpX, spX, ustkX, bstkX} = Stack (I# apX) (I# fpX) (I# spX) (MutableByteArray ustkX) (MutableArray bstkX)
+
+{-# COMPLETE XStack #-}
+
+{-# INLINE XStack #-}
+
+packXStack :: XStack -> Stack
+packXStack (# ap, fp, sp, ustk, bstk #) = Stack {ap = I# ap, fp = I# fp, sp = I# sp, ustk = MutableByteArray ustk, bstk = MutableArray bstk}
+{-# INLINE packXStack #-}
+
+unpackXStack :: Stack -> XStack
+unpackXStack (Stack (I# ap) (I# fp) (I# sp) (MutableByteArray ustk) (MutableArray bstk)) = (# ap, fp, sp, ustk, bstk #)
+{-# INLINE unpackXStack #-}
+
 instance Show Stack where
   show (Stack ap fp sp _ _) =
     "Stack " ++ show ap ++ " " ++ show fp ++ " " ++ show sp