Fix triggering of Worker Wrapper optimization on Stack

.github/workflows/bundle-ucm.yaml

@@ -59,6 +59,7 @@ jobs:
                 --ghc-options='-O2' \
                 --local-bin-path ucm-bin \
                 --copy-bins \
+                --flag unison-runtime:optchecks \
                 && break;
           done
 

unison-runtime/package.yaml

@@ -2,7 +2,7 @@ name: unison-runtime
 github: unisonweb/unison
 copyright: Copyright (C) 2013-2024 Unison Computing, PBC and contributors
 
-ghc-options: -Wall -funbox-strict-fields -O2
+ghc-options: -fmax-worker-args=100 -Wall -funbox-strict-fields -O2
 
 flags:
   arraychecks:
@@ -11,6 +11,13 @@ flags:
   stackchecks:
     manual: true
     default: false
+
+  # Run optimization assertion tests, make sure this runs with O2
+  optchecks:
+    manual: true
+    default: false
+
+  # Dumps core for debugging to unison-runtime/.stack-work/dist/<arch>/ghc-x.y.z/build/
   dumpcore:
     manual: true
     default: false
@@ -20,8 +27,13 @@ when:
     cpp-options: -DARRAY_CHECK
   - condition: flag(stackchecks)
     cpp-options: -DSTACK_CHECK
+  - condition: flag(optchecks)
+    ghc-options: -O2
+    cpp-options: -DOPT_CHECK
+    dependencies:
+      - inspection-testing
   - condition: flag(dumpcore)
-    ghc-options: -ddump-simpl -ddump-stg-final -ddump-to-file -dsuppress-coercions -dsuppress-idinfo -dsuppress-module-prefixes # -dsuppress-type-applications -dsuppress-type-signatures
+    ghc-options: -ddump-simpl -ddump-stg-final -ddump-to-file -dsuppress-coercions -dsuppress-idinfo -dsuppress-module-prefixes -ddump-str-signatures -ddump-simpl-stats # -dsuppress-type-applications -dsuppress-type-signatures
 
 library:
   source-dirs: src
@@ -65,6 +77,8 @@ library:
     - tagged
     - temporary
     - text
+    - template-haskell
+    - inspection-testing
     - time
     - tls
     - unison-codebase-sqlite

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

@@ -1,4 +1,10 @@
-module Unison.Runtime.Exception where
+module Unison.Runtime.Exception
+  ( RuntimeExn (..),
+    die,
+    dieP,
+    exn,
+  )
+where
 
 import Control.Exception
 import Data.String (fromString)
@@ -17,9 +23,12 @@ instance Exception RuntimeExn
 
 die :: (HasCallStack) => String -> IO a
 die = throwIO . PE callStack . P.lit . fromString
+{-# INLINE die #-}
 
 dieP :: (HasCallStack) => P.Pretty P.ColorText -> IO a
 dieP = throwIO . PE callStack
+{-# INLINE dieP #-}
 
 exn :: (HasCallStack) => String -> a
 exn = throw . PE callStack . P.lit . fromString
+{-# INLINE exn #-}

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,6 +24,7 @@ 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)
@@ -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/Interface.hs

@@ -6,6 +6,7 @@
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE UnboxedTuples #-}
 
 module Unison.Runtime.Interface
   ( startRuntime,
@@ -858,8 +859,8 @@ prepareEvaluation ppe tm ctx = do
       Just r -> r
       Nothing -> error "prepareEvaluation: could not remap main ref"
 
-watchHook :: IORef Val -> Stack -> IO ()
-watchHook r stk = peek stk >>= writeIORef r
+watchHook :: IORef Val -> XStack -> IO ()
+watchHook r xstk = peek (packXStack xstk) >>= writeIORef r
 
 backReferenceTm ::
   EnumMap Word64 Reference ->