Make instructions for more Ref primitives

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

@@ -1417,8 +1417,12 @@ data POp
   | SDBL -- sandbox link list
   | SDBV -- sandbox check for Values
   -- Refs
-  | RREF -- Ref.read
-  | WREF -- Ref.write
+  | REFN -- Ref.new
+  | REFR -- Ref.read
+  | REFW -- Ref.write
+  | RCAS -- Ref.cas
+  | RRFC -- Ref.readForCas
+  | TIKR -- Ref.Ticket.read
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 type ANormal = ABTN.Term ANormalF

unison-runtime/src/Unison/Runtime/ANF/Serialize.hs

@@ -656,8 +656,12 @@ pOpCode op = case op of
   COMI -> 128
   DRPN -> 129
   TRNC -> 130
-  RREF -> 131
-  WREF -> 132
+  REFN -> 131
+  REFR -> 132
+  REFW -> 133
+  RCAS -> 134
+  RRFC -> 135
+  TIKR -> 136
 
 pOpAssoc :: [(POp, Word16)]
 pOpAssoc = map (\op -> (op, pOpCode op)) [minBound .. maxBound]

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

@@ -49,10 +49,6 @@ import Data.ByteString (hGet, hGetSome, hPut)
 import Data.ByteString.Lazy qualified as L
 import Data.Default (def)
 import Data.Digest.Murmur64 (asWord64, hash64)
-import Data.IORef as SYS
-  ( IORef,
-    newIORef,
-  )
 import Data.IP (IP)
 import Data.Map qualified as Map
 import Data.PEM (PEM, pemContent, pemParseLBS)
@@ -180,11 +176,7 @@ import Unison.Util.Bytes qualified as Bytes
 import Unison.Util.EnumContainers as EC
 import Unison.Util.RefPromise
   ( Promise,
-    Ticket,
-    casIORef,
     newPromise,
-    peekTicket,
-    readForCAS,
     readPromise,
     tryReadPromise,
     writePromise,
@@ -998,11 +990,29 @@ any'extract =
 
 ref'read :: SuperNormal Symbol
 ref'read =
-  unop0 0 $ \[ref] -> (TPrm RREF [ref])
+  unop0 0 $ \[ref] -> (TPrm REFR [ref])
 
 ref'write :: SuperNormal Symbol
 ref'write =
-  binop0 0 $ \[ref, val] -> (TPrm WREF [ref, val])
+  binop0 0 $ \[ref, val] -> (TPrm REFW [ref, val])
+
+ref'cas :: SuperNormal Symbol
+ref'cas =
+  Lambda [BX, BX, BX]
+    . TAbss [x, y, z]
+    . TLetD b UN (TPrm RCAS [x, y, z])
+    $ boolift b
+  where
+    (x, y, z, b) = fresh
+
+ref'ticket'read :: SuperNormal Symbol
+ref'ticket'read = unop0 0 $ TPrm TIKR
+
+ref'readForCas :: SuperNormal Symbol
+ref'readForCas = unop0 0 $ TPrm RRFC
+
+ref'new :: SuperNormal Symbol
+ref'new = unop0 0 $ TPrm REFN
 
 seek'handle :: ForeignOp
 seek'handle instr =
@@ -1882,7 +1892,12 @@ builtinLookup =
         ("sandboxLinks", (Tracked, sandbox'links)),
         ("IO.tryEval", (Tracked, try'eval)),
         ("Ref.read", (Tracked, ref'read)),
-        ("Ref.write", (Tracked, ref'write))
+        ("Ref.write", (Tracked, ref'write)),
+        ("Ref.cas", (Tracked, ref'cas)),
+        ("Ref.Ticket.read", (Tracked, ref'ticket'read)),
+        ("Ref.readForCas", (Tracked, ref'readForCas)),
+        ("Scope.ref", (Untracked, ref'new)),
+        ("IO.ref", (Tracked, ref'new))
       ]
       ++ foreignWrappers
 
@@ -2381,40 +2396,6 @@ declareForeigns = do
   declareForeign Tracked "STM.retry" unitDirect . mkForeign $
     \() -> unsafeSTMToIO STM.retry :: IO Val
 
-  -- Scope and Ref stuff
-  declareForeign Untracked "Scope.ref" (argNDirect 1)
-    . mkForeign
-    $ \(c :: Val) -> newIORef c
-
-  declareForeign Tracked "IO.ref" (argNDirect 1)
-    . mkForeign
-    $ \(c :: Val) -> evaluate c >>= newIORef
-
-  declareForeign Tracked "Ref.readForCas" (argNDirect 1) . mkForeign $
-    \(r :: IORef Val) -> readForCAS r
-
-  declareForeign Tracked "Ref.Ticket.read" (argNDirect 1) . mkForeign $
-    \(t :: Ticket Val) -> pure $ peekTicket t
-
-  -- In GHC, CAS returns both a Boolean and the current value of the
-  -- IORef, which can be used to retry a failed CAS.
-  -- This strategy is more efficient than returning a Boolean only
-  -- because it uses a single call to cmpxchg in assembly (see [1]) to
-  -- avoid an extra read per CAS iteration, however it's not supported
-  -- in Scheme.
-  -- Therefore, we adopt the more common signature that only returns a
-  -- Boolean, which doesn't even suffer from spurious failures because
-  -- GHC issues loads of mutable variables with memory_order_acquire
-  -- (see [2])
-  --
-  -- [1]: https://github.com/ghc/ghc/blob/master/rts/PrimOps.cmm#L697
-  -- [2]: https://github.com/ghc/ghc/blob/master/compiler/GHC/StgToCmm/Prim.hs#L285
-  declareForeign Tracked "Ref.cas" (argNToBool 3) . mkForeign $
-    \(r :: IORef Val, t :: Ticket Val, v :: Val) -> fmap fst $
-      do
-        t <- evaluate t
-        casIORef r t v
-
   declareForeign Tracked "Promise.new" unitDirect . mkForeign $
     \() -> newPromise @Val
 

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

@@ -22,6 +22,7 @@ where
 import Control.Concurrent (MVar, ThreadId)
 import Control.Concurrent.STM (TVar)
 import Crypto.Hash qualified as Hash
+import Data.Atomics qualified as Atomic
 import Data.IORef (IORef)
 import Data.Tagged (Tagged (..))
 import Data.X509 qualified as X509
@@ -261,6 +262,8 @@ instance BuiltinForeign Value where foreignRef = Tagged Ty.valueRef
 
 instance BuiltinForeign TimeSpec where foreignRef = Tagged Ty.timeSpecRef
 
+instance BuiltinForeign (Atomic.Ticket a) where foreignRef = Tagged Ty.ticketRef
+
 data HashAlgorithm where
   -- Reference is a reference to the hash algorithm
   HashAlgorithm :: (Hash.HashAlgorithm a) => Reference -> a -> HashAlgorithm

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

@@ -403,7 +403,10 @@ data BPrim1
   | DBTX -- debug text
   | SDBL -- sandbox link list
   | -- Refs
-    RREF -- Ref.read
+    REFN -- Ref.new
+  | REFR -- Ref.read
+  | RRFC
+  | TIKR
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 data BPrim2
@@ -440,7 +443,7 @@ data BPrim2
   | SDBX -- sandbox
   | SDBV -- sandbox Value
   -- Refs
-  | WREF -- Ref.write
+  | REFW -- Ref.write
   deriving (Show, Eq, Ord, Enum, Bounded)
 
 data MLit
@@ -478,6 +481,9 @@ data GInstr comb
       !BPrim2
       !Int
       !Int
+  | -- Use a check-and-set ticket to update a reference
+    -- (ref stack index, ticket stack index, new value stack index)
+    RefCAS !Int !Int !Int
   | -- Call out to a Haskell function. This is considerably slower
     -- for very simple operations, hence the primops.
     ForeignCall
@@ -1332,8 +1338,12 @@ 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
+emitPOp ANF.REFN = emitBP1 REFN
+emitPOp ANF.REFR = emitBP1 REFR
+emitPOp ANF.REFW = emitBP2 REFW
+emitPOp ANF.RCAS = refCAS
+emitPOp ANF.RRFC = emitBP1 RRFC
+emitPOp ANF.TIKR = emitBP1 TIKR
 -- non-prim translations
 emitPOp ANF.BLDS = Seq
 emitPOp ANF.FORK = \case
@@ -1391,6 +1401,13 @@ emitBP2 p a =
     "wrong number of args for binary boxed primop: "
       ++ show (p, a)
 
+refCAS :: Args -> Instr
+refCAS (VArgN (primArrayToList -> [i, j, k])) = RefCAS i j k
+refCAS a =
+  internalBug $
+    "wrong number of args for refCAS: "
+      ++ show a
+
 emitDataMatching ::
   (Var v) =>
   Reference ->

unison-runtime/src/Unison/Runtime/MCode/Serialize.hs

@@ -17,8 +17,8 @@ import Data.Bytes.VarInt
 import Data.Void (Void)
 import Data.Word (Word64)
 import GHC.Exts (IsList (..))
-import Unison.Runtime.Array (PrimArray)
 import Unison.Runtime.ANF (PackedTag (..))
+import Unison.Runtime.Array (PrimArray)
 import Unison.Runtime.MCode hiding (MatchT)
 import Unison.Runtime.Serialize
 import Unison.Util.Text qualified as Util.Text
@@ -160,6 +160,7 @@ data InstrT
   | AtomicallyT
   | SeqT
   | TryForceT
+  | RefCAST
 
 instance Tag InstrT where
   tag2word UPrim1T = 0
@@ -179,6 +180,7 @@ instance Tag InstrT where
   tag2word AtomicallyT = 14
   tag2word SeqT = 15
   tag2word TryForceT = 16
+  tag2word RefCAST = 17
 
   word2tag 0 = pure UPrim1T
   word2tag 1 = pure UPrim2T
@@ -197,6 +199,7 @@ instance Tag InstrT where
   word2tag 14 = pure AtomicallyT
   word2tag 15 = pure SeqT
   word2tag 16 = pure TryForceT
+  word2tag 17 = pure RefCAST
   word2tag n = unknownTag "InstrT" n
 
 putInstr :: (MonadPut m) => GInstr cix -> m ()
@@ -205,6 +208,7 @@ putInstr = \case
   (UPrim2 up i j) -> putTag UPrim2T *> putTag up *> pInt i *> pInt j
   (BPrim1 bp i) -> putTag BPrim1T *> putTag bp *> pInt i
   (BPrim2 bp i j) -> putTag BPrim2T *> putTag bp *> pInt i *> pInt j
+  (RefCAS i j k) -> putTag RefCAST *> pInt i *> pInt j *> pInt k
   (ForeignCall b w a) -> putTag ForeignCallT *> serialize b *> pWord w *> putArgs a
   (SetDyn w i) -> putTag SetDynT *> pWord w *> pInt i
   (Capture w) -> putTag CaptureT *> pWord w
@@ -226,6 +230,7 @@ getInstr =
     UPrim2T -> UPrim2 <$> getTag <*> gInt <*> gInt
     BPrim1T -> BPrim1 <$> getTag <*> gInt
     BPrim2T -> BPrim2 <$> getTag <*> gInt <*> gInt
+    RefCAST -> RefCAS <$> gInt <*> gInt <*> gInt
     ForeignCallT -> ForeignCall <$> deserialize <*> gWord <*> getArgs
     SetDynT -> SetDyn <$> gWord <*> gInt
     CaptureT -> Capture <$> gWord

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

@@ -8,6 +8,7 @@ import Control.Concurrent (ThreadId)
 import Control.Concurrent.STM as STM
 import Control.Exception
 import Control.Lens
+import Data.Atomics qualified as Atomic
 import Data.Bits
 import Data.Functor.Classes (Eq1 (..), Ord1 (..))
 import Data.IORef (IORef)
@@ -548,6 +549,15 @@ exec !env !denv !_activeThreads !stk !k _ (BPrim2 TRCE i j)
 exec !_ !denv !_trackThreads !stk !k _ (BPrim2 op i j) = do
   stk <- bprim2 stk op i j
   pure (denv, stk, k)
+exec !_ !denv !_activeThreads !stk !k _ (RefCAS refI ticketI valI) = do
+  (ref :: IORef Val) <- peekOffBi stk refI
+  (ticket :: Atomic.Ticket Val) <- peekOffBi stk ticketI
+  v <- peekOff stk valI
+  ticket <- evaluate ticket
+  (r, _) <- Atomic.casIORef ref ticket v
+  stk <- bump stk
+  pokeBool stk r
+  pure (denv, stk, k)
 exec !_ !denv !_activeThreads !stk !k _ (Pack r t args) = do
   clo <- buildData stk r t args
   stk <- bump stk
@@ -1623,12 +1633,30 @@ bprim1 !stk FLTB i = do
 -- [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
+bprim1 !stk REFR i = do
   (ref :: IORef Val) <- peekOffBi stk i
   v <- IORef.readIORef ref
   stk <- bump stk
   poke stk v
   pure stk
+bprim1 !stk REFN i = do
+  v <- peekOff stk i
+  ref <- IORef.newIORef v
+  stk <- bump stk
+  pokeBi stk ref
+  pure stk
+bprim1 !stk RRFC i = do
+  (ref :: IORef Val) <- peekOffBi stk i
+  ticket <- Atomic.readForCAS ref
+  stk <- bump stk
+  pokeBi stk ticket
+  pure stk
+bprim1 !stk TIKR i = do
+  (t :: Atomic.Ticket Val) <- peekOffBi stk i
+  stk <- bump stk
+  let v = Atomic.peekTicket t
+  poke stk v
+  pure stk
 
 -- impossible
 bprim1 !stk MISS _ = pure stk
@@ -1834,7 +1862,7 @@ bprim2 !stk CATB i j = do
   stk <- bump stk
   pokeBi stk (l <> r :: By.Bytes)
   pure stk
-bprim2 !stk WREF i j = do
+bprim2 !stk REFW i j = do
   (ref :: IORef Val) <- peekOffBi stk i
   v <- peekOff stk j
   IORef.writeIORef ref v

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

@@ -476,7 +476,10 @@ instance Tag BPrim1 where
   tag2word TLTT = 24
   tag2word DBTX = 25
   tag2word SDBL = 26
-  tag2word RREF = 27
+  tag2word REFN = 27
+  tag2word REFR = 28
+  tag2word RRFC = 29
+  tag2word TIKR = 30
 
   word2tag 0 = pure SIZT
   word2tag 1 = pure USNC
@@ -505,7 +508,10 @@ instance Tag BPrim1 where
   word2tag 24 = pure TLTT
   word2tag 25 = pure DBTX
   word2tag 26 = pure SDBL
-  word2tag 27 = pure RREF
+  word2tag 27 = pure REFN
+  word2tag 28 = pure REFR
+  word2tag 29 = pure RRFC
+  word2tag 30 = pure TIKR
   word2tag n = unknownTag "BPrim1" n
 
 instance Tag BPrim2 where
@@ -535,7 +541,7 @@ instance Tag BPrim2 where
   tag2word IXOT = 23
   tag2word IXOB = 24
   tag2word SDBV = 25
-  tag2word WREF = 26
+  tag2word REFW = 26
 
   word2tag 0 = pure EQLU
   word2tag 1 = pure CMPU
@@ -563,5 +569,5 @@ instance Tag BPrim2 where
   word2tag 23 = pure IXOT
   word2tag 24 = pure IXOB
   word2tag 25 = pure SDBV
-  word2tag 26 = pure WREF
+  word2tag 26 = pure REFW
   word2tag n = unknownTag "BPrim2" n