From 64348e9d42d5c24a9c7416e5b3589a51df339235 Mon Sep 17 00:00:00 2001
From: Chris Penner <christopher.penner@gmail.com>
Date: Wed, 16 Oct 2024 16:45:02 -0700
Subject: [PATCH] Make Ref.read Ref.write instructions

---
 unison-runtime/src/Unison/Runtime/ANF.hs      |  3 ++
 .../src/Unison/Runtime/ANF/Serialize.hs       |  2 ++
 unison-runtime/src/Unison/Runtime/Builtin.hs  | 31 +++++++++----------
 unison-runtime/src/Unison/Runtime/Foreign.hs  |  4 +++
 unison-runtime/src/Unison/Runtime/MCode.hs    |  7 +++++
 unison-runtime/src/Unison/Runtime/Machine.hs  | 23 +++++++++++++-
 .../src/Unison/Runtime/Serialize.hs           |  4 +++
 7 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/unison-runtime/src/Unison/Runtime/ANF.hs b/unison-runtime/src/Unison/Runtime/ANF.hs
index 6f6a615da2..d22dbda337 100644
--- a/unison-runtime/src/Unison/Runtime/ANF.hs
+++ b/unison-runtime/src/Unison/Runtime/ANF.hs
@@ -1445,6 +1445,9 @@ data POp
   | TFRC -- try force
   | SDBL -- sandbox link list
   | SDBV -- sandbox check for Values
+  -- Refs
+  | RREF -- Ref.read
+  | WREF -- Ref.write
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 type ANormal = ABTN.Term ANormalF
diff --git a/unison-runtime/src/Unison/Runtime/ANF/Serialize.hs b/unison-runtime/src/Unison/Runtime/ANF/Serialize.hs
index 75c27ba79d..4752a44b48 100644
--- a/unison-runtime/src/Unison/Runtime/ANF/Serialize.hs
+++ b/unison-runtime/src/Unison/Runtime/ANF/Serialize.hs
@@ -649,6 +649,8 @@ pOpCode op = case op of
   IXOB -> 121
   SDBL -> 122
   SDBV -> 123
+  RREF -> 124
+  WREF -> 125
 
 pOpAssoc :: [(POp, Word16)]
 pOpAssoc = map (\op -> (op, pOpCode op)) [minBound .. maxBound]
diff --git a/unison-runtime/src/Unison/Runtime/Builtin.hs b/unison-runtime/src/Unison/Runtime/Builtin.hs
index 0a31bdce41..03c35acc7d 100644
--- a/unison-runtime/src/Unison/Runtime/Builtin.hs
+++ b/unison-runtime/src/Unison/Runtime/Builtin.hs
@@ -51,8 +51,6 @@ import Data.Digest.Murmur64 (asWord64, hash64)
 import Data.IORef as SYS
   ( IORef,
     newIORef,
-    readIORef,
-    writeIORef,
   )
 import Data.IP (IP)
 import Data.Map qualified as Map
@@ -1089,6 +1087,18 @@ any'extract =
       TMatch v $
         MatchData Ty.anyRef (mapSingleton 0 $ ([BX], TAbs v1 (TVar v1))) Nothing
 
+-- Refs
+
+ref'read :: SuperNormal Symbol
+ref'read =
+  unop0 0 $ \[ref] -> (TPrm RREF [ref])
+
+ref'write :: SuperNormal Symbol
+ref'write =
+  binop0 1 $ \[ref, val, unit] ->
+    TLetD unit BX (TCon Ty.unitRef 0 [])
+    $ (TPrm WREF [ref, val])
+
 seek'handle :: ForeignOp
 seek'handle instr =
   ([BX, BX, BX],)
@@ -2253,7 +2263,9 @@ builtinLookup =
         ("validateSandboxed", (Untracked, check'sandbox)),
         ("Value.validateSandboxed", (Tracked, value'sandbox)),
         ("sandboxLinks", (Tracked, sandbox'links)),
-        ("IO.tryEval", (Tracked, try'eval))
+        ("IO.tryEval", (Tracked, try'eval)),
+        ("Ref.read", (Tracked, ref'read)),
+        ("Ref.write", (Tracked, ref'write))
       ]
       ++ foreignWrappers
 
@@ -2761,19 +2773,6 @@ declareForeigns = do
     . mkForeign
     $ \(c :: Closure) -> evaluate c >>= newIORef
 
-  -- The docs for IORef state that IORef operations can be observed
-  -- out of order ([1]) but actually GHC does emit the appropriate
-  -- load and store barriers nowadays ([2], [3]).
-  --
-  -- [1] https://hackage.haskell.org/package/base-4.17.0.0/docs/Data-IORef.html#g:2
-  -- [2] https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L286
-  -- [3] https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L298
-  declareForeign Untracked "Ref.read" boxDirect . mkForeign $
-    \(r :: IORef Closure) -> readIORef r
-
-  declareForeign Untracked "Ref.write" boxBoxTo0 . mkForeign $
-    \(r :: IORef Closure, c :: Closure) -> evaluate c >>= writeIORef r
-
   declareForeign Tracked "Ref.readForCas" boxDirect . mkForeign $
     \(r :: IORef Closure) -> readForCAS r
 
diff --git a/unison-runtime/src/Unison/Runtime/Foreign.hs b/unison-runtime/src/Unison/Runtime/Foreign.hs
index 5559ce9b6c..d4a8853b48 100644
--- a/unison-runtime/src/Unison/Runtime/Foreign.hs
+++ b/unison-runtime/src/Unison/Runtime/Foreign.hs
@@ -256,10 +256,14 @@ instance BuiltinForeign FilePath where foreignRef = Tagged Ty.filePathRef
 instance BuiltinForeign TLS.Context where foreignRef = Tagged Ty.tlsRef
 
 instance BuiltinForeign Code where foreignRef = Tagged Ty.codeRef
+
 instance BuiltinForeign Value where foreignRef = Tagged Ty.valueRef
 
 instance BuiltinForeign TimeSpec where foreignRef = Tagged Ty.timeSpecRef
 
+-- TODO: investigate this
+instance BuiltinForeign (IORef closure) where foreignRef = Tagged Ty.refRef
+
 data HashAlgorithm where
   -- Reference is a reference to the hash algorithm
   HashAlgorithm :: (Hash.HashAlgorithm a) => Reference -> a -> HashAlgorithm
diff --git a/unison-runtime/src/Unison/Runtime/MCode.hs b/unison-runtime/src/Unison/Runtime/MCode.hs
index e013a47adf..e32ad5469a 100644
--- a/unison-runtime/src/Unison/Runtime/MCode.hs
+++ b/unison-runtime/src/Unison/Runtime/MCode.hs
@@ -387,6 +387,8 @@ data BPrim1
   -- debug
   | DBTX -- debug text
   | SDBL -- sandbox link list
+  | -- Refs
+    RREF -- Ref.read
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 data BPrim2
@@ -422,6 +424,8 @@ data BPrim2
   -- code
   | SDBX -- sandbox
   | SDBV -- sandbox Value
+  -- Refs
+  | WREF -- Ref.write
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 data MLit
@@ -1285,6 +1289,9 @@ emitPOp ANF.SDBV = emitBP2 SDBV
 emitPOp ANF.EROR = emitBP2 THRO
 emitPOp ANF.TRCE = emitBP2 TRCE
 emitPOp ANF.DBTX = emitBP1 DBTX
+-- Refs
+emitPOp ANF.RREF = emitBP1 RREF
+emitPOp ANF.WREF = emitBP2 WREF
 -- non-prim translations
 emitPOp ANF.BLDS = Seq
 emitPOp ANF.FORK = \case
diff --git a/unison-runtime/src/Unison/Runtime/Machine.hs b/unison-runtime/src/Unison/Runtime/Machine.hs
index 48cf202f27..8e69b8958c 100644
--- a/unison-runtime/src/Unison/Runtime/Machine.hs
+++ b/unison-runtime/src/Unison/Runtime/Machine.hs
@@ -61,8 +61,9 @@ import Unison.Util.Bytes qualified as By
 import Unison.Util.EnumContainers as EC
 import Unison.Util.Pretty (toPlainUnbroken)
 import Unison.Util.Text qualified as Util.Text
-import UnliftIO (IORef)
+import Data.IORef (IORef)
 import UnliftIO qualified
+import Data.IORef qualified as IORef
 import UnliftIO.Concurrent qualified as UnliftIO
 
 -- | A ref storing every currently active thread.
@@ -1519,6 +1520,21 @@ bprim1 !stk FLTB i = do
   stk <- bump stk
   pokeBi stk $ By.flatten b
   pure stk
+
+-- The docs for IORef state that IORef operations can be observed
+-- out of order ([1]) but actually GHC does emit the appropriate
+-- load and store barriers nowadays ([2], [3]).
+--
+-- [1] https://hackage.haskell.org/package/base-4.17.0.0/docs/Data-IORef.html#g:2
+-- [2] https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L286
+-- [3] https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L298
+bprim1 !stk RREF i = do
+  (ref :: IORef Closure) <- peekOffBi stk i
+  v <- IORef.readIORef ref
+  stk <- bump stk
+  bpoke stk v
+  pure stk
+
 -- impossible
 bprim1 !stk MISS _ = pure stk
 bprim1 !stk CACH _ = pure stk
@@ -1729,6 +1745,11 @@ bprim2 !stk CATB i j = do
   stk <- bump stk
   pokeBi stk (l <> r :: By.Bytes)
   pure stk
+bprim2 !stk WREF i j = do
+  (ref :: IORef Closure) <- peekOffBi stk i
+  v <- bpeekOff stk j
+  IORef.writeIORef ref v
+  pure stk
 bprim2 !stk THRO _ _ = pure stk -- impossible
 bprim2 !stk TRCE _ _ = pure stk -- impossible
 bprim2 !stk CMPU _ _ = pure stk -- impossible
diff --git a/unison-runtime/src/Unison/Runtime/Serialize.hs b/unison-runtime/src/Unison/Runtime/Serialize.hs
index 394b846a0b..0f7098059c 100644
--- a/unison-runtime/src/Unison/Runtime/Serialize.hs
+++ b/unison-runtime/src/Unison/Runtime/Serialize.hs
@@ -446,6 +446,7 @@ instance Tag BPrim1 where
   tag2word TLTT = 24
   tag2word DBTX = 25
   tag2word SDBL = 26
+  tag2word RREF = 27
 
   word2tag 0 = pure SIZT
   word2tag 1 = pure USNC
@@ -474,6 +475,7 @@ instance Tag BPrim1 where
   word2tag 24 = pure TLTT
   word2tag 25 = pure DBTX
   word2tag 26 = pure SDBL
+  word2tag 27 = pure RREF
   word2tag n = unknownTag "BPrim1" n
 
 instance Tag BPrim2 where
@@ -503,6 +505,7 @@ instance Tag BPrim2 where
   tag2word IXOT = 23
   tag2word IXOB = 24
   tag2word SDBV = 25
+  tag2word WREF = 26
 
   word2tag 0 = pure EQLU
   word2tag 1 = pure CMPU
@@ -530,4 +533,5 @@ instance Tag BPrim2 where
   word2tag 23 = pure IXOT
   word2tag 24 = pure IXOB
   word2tag 25 = pure SDBV
+  word2tag 26 = pure WREF
   word2tag n = unknownTag "BPrim2" n