Merge branch 'trunk' into runarorama/avroMagic

aryairani · web-flow · commit 30d87f637ccb · 2025-07-16T05:50:45.000+03:00
diff --git a/unison-runtime/src/Unison/Runtime/ANF.hs b/unison-runtime/src/Unison/Runtime/ANF.hs
@@ -125,6 +125,7 @@ import Unison.Reference (Id, Reference, Reference' (Builtin, DerivedId), toShort
 import Unison.Referent (Referent, pattern Con, pattern Ref)
 import Unison.Runtime.Array qualified as PA
 import Unison.Runtime.Foreign.Function.Type (ForeignFunc (..))
+import Unison.Runtime.Referenced
 import Unison.Runtime.TypeTags (CTag (..), PackedTag (..), RTag (..), Tag (..), maskTags, packTags, unpackTags)
 import Unison.ShortHash (shortenTo)
 import Unison.Symbol (Symbol)
@@ -1609,28 +1610,6 @@ type ANFD v = Compose (ANFM v) (Directed ())
 data GroupRef = GR Reference Word64
   deriving (Show, Eq)
 
--- A value with optional optimization information for serialization.
--- The references are required for serialization V5, and are assumed
--- to be the only references used in the value _up to in-memory
--- uniqueness_.
---
--- This is parameterized so that it can be used with both Value and
--- Code.
---
--- Also note, the stored referenced might not be 'tight' in the sense
--- that they all actually occur in the value. Maintaining this
--- invariant together with actual canonicalization would be onerous
--- and isn't done at this time.
-data Referenced a
-  = -- types, terms
-    WithRefs [Reference] [Reference] a
-  | Plain a
-  deriving (Show, Eq)
-
-dereference :: Referenced a -> a
-dereference (WithRefs _ _ x) = x
-dereference (Plain x) = x
-
 -- | A list of either unboxed or boxed values.
 -- Each slot is one of unboxed or boxed but not both.
 type ValList = [Value]
diff --git a/unison-runtime/src/Unison/Runtime/ANF/Serialize.hs b/unison-runtime/src/Unison/Runtime/ANF/Serialize.hs
@@ -7,6 +7,7 @@ module Unison.Runtime.ANF.Serialize where
 
 import Control.Monad
 import Control.Monad.Reader
+import Control.Monad.State.Strict (StateT (..))
 import Data.Bifunctor (bimap, first)
 import Data.Binary.Get (runGetOrFail)
 import Data.Binary.Get qualified as BGet
@@ -37,8 +38,10 @@ import Unison.Runtime.ANF.Optimize as ANF
 import Unison.Runtime.ANF.Serialize.CodeV4 qualified as CodeV4
 import Unison.Runtime.ANF.Serialize.Tags
 import Unison.Runtime.ANF.Serialize.ValueV5 qualified as ValueV5
+import Unison.Runtime.Canonicalizer qualified as C
 import Unison.Runtime.Exception
 import Unison.Runtime.Foreign.Function.Type (ForeignFunc)
+import Unison.Runtime.Referenced
 import Unison.Runtime.Serialize
 import Unison.Util.Text qualified as Util.Text
 import Unison.Var (Type (ANFBlank), Var (..))
@@ -964,27 +967,30 @@ deserializeCode bs = runGetS go bs
 -- Boolean argument determines whether ForeignFunc occurrences are
 -- allowed to be serialized. For interchange, this should be False.
 serializeCode :: Bool -> Referenced Code -> ByteString
-serializeCode fops (WithRefs tys tms co) =
-  runPutS (putWord32be 4 *> CodeV4.putCodeWithHeader tys tms fops co)
-serializeCode fops (Plain co) = runPutS (putVersion *> putCode fops co)
+serializeCode fops (dereference -> co) =
+  runPutS (putVersion *> putCode fops co)
   where
     putVersion = putWord32be codeVersion
 
 serializeCodeWithVersion ::
-  Word64 -> Bool -> Referenced Code -> Either String L.ByteString
-serializeCodeWithVersion v fops = \case
-  WithRefs tys tms co
-    | v == 4 ->
-        Right . runPutL $
+  Word64 -> Bool -> Referenced Code -> IO (Either String L.ByteString)
+serializeCodeWithVersion v fops rco
+  | v == 4 =
+      enreference rco >>= \(tys, tms, co) ->
+        pure . Right . runPutL $
           putWord32be 4 *> CodeV4.putCodeWithHeader tys tms fops co
-  rco
-    | v == 3 ->
-        Right . runPutL $
-          putWord32be 3 *> putCode fops (dereference rco)
-    | v == 4 ->
-        Left "could not serialize plain code at v4"
-    | otherwise ->
-        Left $ "unsupported code serialization version: " ++ show v
+  | v == 3 =
+      pure . Right . runPutL $
+        putWord32be 3 *> putCode fops (dereference rco)
+  | otherwise =
+      pure . Left $ "unsupported code serialization version: " ++ show v
+  where
+    enreference (WithRefs tys tms co) = pure (tys, tms, co)
+    enreference (Plain co) =
+      runStateT
+        (canonicalizeRefs traverseCodeRefs co)
+        (C.empty, [], [])
+        >>= \(co, (_, tys, tms)) -> pure (tys, tms, co)
 
 -- | Serializes a `SuperGroup` for rehashing.
 --
@@ -1042,14 +1048,28 @@ serializeValue (dereference -> v) =
   where
     putVersion = putWord32be valueVersion
 
-serializeValueWithVersion :: Word64 -> Referenced Value -> L.ByteString
-serializeValueWithVersion v = \case
-  WithRefs tys tms x
-    | v == 5 ->
-        runPutL $ putWord32be 5 *> ValueV5.putValueWithHeader tys tms x
-  rval
-    | n <- fromIntegral v ->
-        runPutL $ putWord32be n *> putValue (Transfer n) (dereference rval)
+serializeValueWithVersion ::
+  Word64 -> Referenced Value -> IO L.ByteString
+serializeValueWithVersion v rval
+  | v == 5 = case rval of
+      WithRefs tys tms x -> v5ser tys tms x
+      Plain x -> do
+        (x, (_, tys, tms)) <-
+          runStateT
+            (canonicalizeRefs traverseValueRefs x)
+            (C.empty, [], [])
+        v5ser tys tms x
+  | v < 5,
+    n <- fromIntegral v =
+      pure . runPutL $
+        putWord32be n
+          *> putValue (Transfer n) (dereference rval)
+  | otherwise =
+      die $ "Value.serialize.versioned: unrecognized version: " ++ show v
+  where
+    v5ser tys tms x =
+      pure . runPutL $
+        putWord32be 5 *> ValueV5.putValueWithHeader tys tms x
 
 -- This serializer is used exclusively for hashing unison values.
 -- For this reason, it doesn't prefix the string with the current
diff --git a/unison-runtime/src/Unison/Runtime/Foreign/Function.hs b/unison-runtime/src/Unison/Runtime/Foreign/Function.hs
@@ -519,7 +519,7 @@ foreignCallHelper = \case
       pure . Bytes.fromArray $ ANF.serializeCode False co
   Code_serialize_versioned -> mkForeign $
     \(ver :: Word64, co :: ANF.Referenced ANF.Code) ->
-      case ANF.serializeCodeWithVersion ver False co of
+      ANF.serializeCodeWithVersion ver False co >>= \case
         Left err -> die err
         Right bs -> pure $ Bytes.fromLazyByteString bs
   Code_deserialize ->
@@ -536,7 +536,7 @@ foreignCallHelper = \case
       pure . Bytes.fromArray . ANF.serializeValue
   Value_serialize_versioned ->
     mkForeign $
-      pure . Bytes.fromLazyByteString . uncurry ANF.serializeValueWithVersion
+      fmap Bytes.fromLazyByteString . uncurry ANF.serializeValueWithVersion
   Value_deserialize ->
     mkForeign $
       pure . ANF.deserializeValue . Bytes.toLazyByteString
diff --git a/unison-runtime/src/Unison/Runtime/Machine.hs b/unison-runtime/src/Unison/Runtime/Machine.hs
@@ -87,6 +87,7 @@ import Unison.Runtime.Foreign.Function
 import Unison.Runtime.MCode
 import Unison.Runtime.Machine.Primops
 import Unison.Runtime.Machine.Types
+import Unison.Runtime.Referenced
 import Unison.Runtime.Stack
 import Unison.Runtime.TypeTags qualified as TT
 import Unison.Symbol (Symbol)
@@ -1415,85 +1416,24 @@ type Reflect = StateT ReflectionState IO
 emptyRS :: ReflectionState
 emptyRS = RS HM.empty HM.empty C.empty [] []
 
-type RTrav a =
-  forall f.
-  (Applicative f) =>
-  (Bool -> Reference -> f Reference) ->
-  (a -> f a)
+mediate ::
+  StateT (C.Canonicalizer Reference, [Reference], [Reference]) IO a ->
+  Reflect a
+mediate act = StateT \(RS sty stm canon tys tms) -> do
+  (v, (canon, tys, tms)) <- runStateT act (canon, tys, tms)
+  pure $ (v, RS sty stm canon tys tms)
+{-# INLINE mediate #-}
 
 canonicalizeReference :: Bool -> Reference -> Reflect Reference
-canonicalizeReference isTy r = StateT \st@(RS _ _ canon tys tms) ->
-  C.categorize canon r >>= \case
-    C.Canonical -> pure (r, st)
-    C.Equivalent s canon -> (s,) <$> evaluate (st {_canon = canon})
-    C.Novel canon ->
-      (r,)
-        <$> evaluate
-          st
-            { _canon = canon,
-              _tys = if isTy then r : tys else tys,
-              _tms = if isTy then tms else r : tms
-            }
+canonicalizeReference isTy = mediate . canonicalizeRefs (\f -> f isTy)
 
 canonicalizeReferent :: Referent -> Reflect Referent
 canonicalizeReferent (Ref r) = Ref <$> canonicalizeReference False r
 canonicalizeReferent (Con (ConstructorReference r i) j) =
   flip Con j . flip ConstructorReference i <$> canonicalizeReference True r
 
 canonicalizeReferenced :: RTrav a -> ANF.Referenced a -> Reflect a
-canonicalizeReferenced trav = \case
-  -- no stored refs, have to traverse
-  ANF.Plain v -> trav h v
-  ANF.WithRefs tys tms v -> do
-    typs <- mapMaybe id <$> traverse (g True) tys
-    tmps <- mapMaybe id <$> traverse (g False) tms
-
-    ctys <- lift $ C.fromList typs
-    ctms <- lift $ C.fromList tmps
-
-    let f False r = C.findWithDefault r r ctms
-        f True r = C.findWithDefault r r ctys
-
-    if null typs && null tmps
-      then -- all references are already canonical
-        pure v
-      else lift $ trav f v
-  where
-    -- traversal function for plain values
-    g isTy r = StateT \st@(RS _ _ canon tys tms) ->
-      C.categorize canon r >>= \case
-        C.Canonical -> pure (Nothing, st)
-        C.Novel canon ->
-          (Nothing,)
-            <$> evaluate
-              st
-                { _canon = canon,
-                  _tys = if isTy then r : tys else tys,
-                  _tms = if isTy then tms else r : tms
-                }
-        C.Equivalent s canon ->
-          (Just (r, s),) <$> evaluate (st {_canon = canon})
-
-    -- traversal function for remapping WithRefs values
-    h isTy r = StateT \st@(RS _ _ canon tys tms) ->
-      C.categorize canon r >>= \case
-        C.Canonical -> pure (r, st)
-        C.Novel canon ->
-          (r,)
-            <$> evaluate
-              if isTy
-                then
-                  st
-                    { _canon = canon,
-                      _tys = r : tys
-                    }
-                else
-                  st
-                    { _canon = canon,
-                      _tms = r : tms
-                    }
-        C.Equivalent r canon ->
-          (r,) <$> evaluate (st {_canon = canon})
+canonicalizeReferenced trav x = mediate $ recanonicalizeRefs trav x
 {-# INLINE canonicalizeReferenced #-}
 
 reflectValue :: CCache -> Val -> IO (ANF.Referenced ANF.Value)
diff --git a/unison-runtime/src/Unison/Runtime/Referenced.hs b/unison-runtime/src/Unison/Runtime/Referenced.hs
@@ -0,0 +1,104 @@
+module Unison.Runtime.Referenced
+  ( Referenced (..),
+    dereference,
+    RTrav,
+    Canonize,
+    canonicalizeRefs,
+    recanonicalizeRefs,
+  )
+where
+
+import Control.Monad.State.Strict
+import Data.Maybe (mapMaybe)
+import Unison.Reference
+import Unison.Runtime.Canonicalizer
+
+-- A value with optional optimization information for serialization.
+-- The references are required for serialization V5, and are assumed
+-- to be the only references used in the value _up to in-memory
+-- uniqueness_.
+--
+-- This is parameterized so that it can be used with both Value and
+-- Code.
+--
+-- Also note, the stored referenced might not be 'tight' in the sense
+-- that they all actually occur in the value. Maintaining this
+-- invariant together with actual canonicalization would be onerous
+-- and isn't done at this time.
+data Referenced a
+  = -- types, terms
+    WithRefs [Reference] [Reference] a
+  | Plain a
+  deriving (Show, Eq)
+
+dereference :: Referenced a -> a
+dereference (WithRefs _ _ x) = x
+dereference (Plain x) = x
+
+type RTrav a =
+  forall f.
+  (Applicative f) =>
+  (Bool -> Reference -> f Reference) ->
+  (a -> f a)
+
+type Canonize =
+  StateT (Canonicalizer Reference, [Reference], [Reference]) IO
+
+-- Given a reference traversal, canonicalizes the references in a
+-- value. The operation is presented as a state transformation, so
+-- that it can hook into a larger canonicalization procedure. The
+-- lists of canonical references of each sort are yielded as part of
+-- the state.
+canonicalizeRefs :: RTrav a -> a -> Canonize a
+canonicalizeRefs trav = trav h
+  where
+    h isTy r = StateT \st@(canon, tys, tms) ->
+      categorize canon r >>= \case
+        Canonical -> pure (r, st)
+        Novel canon
+          | isTy -> pure (r, (canon, r : tys, tms))
+          | otherwise -> pure (r, (canon, tys, r : tms))
+        Equivalent s canon -> pure (s, (canon, tys, tms))
+{-# INLINE canonicalizeRefs #-}
+
+-- Given a `Referenced` value, this recanonicalizes the references in
+-- the wrapped value. The intention is for this to be hooked into a
+-- larger canonicalization procedure, so that already canonicalized
+-- values can be more efficiently brought in line with other values
+-- that are already canonicalized.
+--
+-- If the `Referenced` value is `Plain`, then all we can do is
+-- traverse it, canonicalizing the references. However, if it is
+-- tagged with canonical refs, we can see if they all match existing
+-- canonical refs. If so, we don't need to traverse the value. Even if
+-- not, we can traverse with marginally more efficient lookups.
+recanonicalizeRefs :: RTrav a -> Referenced a -> Canonize a
+recanonicalizeRefs trav = \case
+  Plain v -> canonicalizeRefs trav v
+  WithRefs tys tms v -> do
+    typs <- mapMaybe id <$> traverse (g True) tys
+    tmps <- mapMaybe id <$> traverse (g False) tms
+
+    ctys <- lift $ fromList typs
+    ctms <- lift $ fromList tmps
+
+    let f False r = findWithDefault r r ctms
+        f True r = findWithDefault r r ctys
+
+    if null typs && null tmps
+      then pure v -- already canonical
+      else lift $ trav f v
+  where
+    g isTy r = StateT \st@(canon, tys, tms) ->
+      categorize canon r >>= \case
+        Canonical -> pure (Nothing, st)
+        Novel canon ->
+          pure
+            ( Nothing,
+              ( canon,
+                if isTy then r : tys else tys,
+                if isTy then tms else r : tms
+              )
+            )
+        Equivalent s canon -> pure (Just (r, s), (canon, tys, tms))
+{-# INLINE recanonicalizeRefs #-}
diff --git a/unison-runtime/unison-runtime.cabal b/unison-runtime/unison-runtime.cabal
@@ -66,6 +66,7 @@ library
       Unison.Runtime.MCode
       Unison.Runtime.MCode.Serialize
       Unison.Runtime.Pattern
+      Unison.Runtime.Referenced
       Unison.Runtime.Serialize
       Unison.Runtime.SparseVector
       Unison.Runtime.Stack
diff --git a/unison-src/transcripts-using-base/codeops.md b/unison-src/transcripts-using-base/codeops.md
diff --git a/unison-src/transcripts-using-base/codeops.output.md b/unison-src/transcripts-using-base/codeops.output.md
diff --git a/unison-src/transcripts/idempotent/affine-handlers.md b/unison-src/transcripts/idempotent/affine-handlers.md
