From d8d4f0800bf8de14fdf165b521ff1e02879fe9ce Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Fri, 25 Oct 2024 17:27:20 -0500
Subject: [PATCH 01/17] [wip] Towards #2107; add structure + homomorphic hashes
 to Ctx

---
 src/swarm-lang/Swarm/Language/Context.hs | 92 +++++++++++++++++++-----
 1 file changed, 73 insertions(+), 19 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index ac78de5e7..bef7c4c10 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE ViewPatterns #-}
@@ -16,8 +17,10 @@ import Control.Lens.Empty (AsEmpty (..), pattern Empty)
 import Control.Lens.Prism (prism)
 import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
 import Data.Data (Data)
+import Data.Hashable
 import Data.Map (Map)
 import Data.Map qualified as M
+import Data.Semigroup (Sum (..))
 import Data.Text (Text)
 import GHC.Generics (Generic)
 import Prettyprinter (brackets, emptyDoc, hsep, punctuate)
@@ -28,54 +31,96 @@ import Prelude hiding (lookup)
 -- | We use 'Text' values to represent variables.
 type Var = Text
 
+-- | A data type to record the structure of how a context was built,
+--   so that we can later destruct/serialize it effectively.
+data CtxStruct t
+  = CtxEmpty
+  | CtxSingle Var
+  | CtxDelete Var t (CtxStruct t)
+  | CtxUnion (CtxStruct t) (CtxStruct t)
+  deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic, ToJSON, FromJSON)
+
+-- | A context hash is a hash value used to identify contexts without
+--   having to compare them for equality.  Hash values are computed in
+--   homomorphically, so that two equal contexts will be guaranteed to
+--   have the same hash value, even if they were constructed
+--   by a different sequence of operations.
+--
+--   The downside of this approach is that, /in theory/, there could
+--   be a hash collision---two different contexts which nonetheless
+--   have the same hash value.  However, this is extremely unlikely.
+--   The benefit is that everything can be purely functional, without
+--   the need to thread around some kind of globally unique ID
+--   generation effect.
+newtype CtxHash = CtxHash {getCtxHash :: Int}
+  deriving (Eq, Show, Data, Generic, ToJSON, FromJSON)
+  deriving (Semigroup, Monoid) via Sum Int
+  deriving (Num) via Int
+
 -- | A context is a mapping from variable names to things.
-newtype Ctx t = Ctx {unCtx :: Map Var t}
-  deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic)
+data Ctx t = Ctx {unCtx :: Map Var t, ctxHash :: CtxHash, ctxStruct :: CtxStruct t}
+  deriving (Eq, Show, Functor, Traversable, Data, Generic)
+
+instance Hashable t => Hashable (Ctx t) where
+  hash = getCtxHash . ctxHash
+  hashWithSalt s = hashWithSalt s . getCtxHash . ctxHash
 
+instance Foldable Ctx where
+  foldMap f = foldMap f . unCtx
+
+-- XXX this instance will have to change!!
 instance ToJSON t => ToJSON (Ctx t) where
   toJSON = genericToJSON optionsUnwrapUnary
 
+-- XXX this instance will have to change!!
 instance FromJSON t => FromJSON (Ctx t) where
   parseJSON = genericParseJSON optionsUnwrapUnary
 
+-- XXX this instance will have to change!!
 instance (PrettyPrec t) => PrettyPrec (Ctx t) where
   prettyPrec _ Empty = emptyDoc
   prettyPrec _ (assocs -> bs) = brackets (hsep (punctuate "," (map prettyBinding bs)))
 
 -- | The semigroup operation for contexts is /right/-biased union.
-instance Semigroup (Ctx t) where
+instance Hashable t => Semigroup (Ctx t) where
   (<>) = union
 
-instance Monoid (Ctx t) where
+instance Hashable t => Monoid (Ctx t) where
   mempty = empty
   mappend = (<>)
 
 instance AsEmpty (Ctx t) where
   _Empty = prism (const empty) isEmpty
    where
-    isEmpty (Ctx c)
-      | M.null c = Right ()
-      | otherwise = Left (Ctx c)
+    isEmpty c
+      | M.null (unCtx c) = Right ()
+      | otherwise = Left c
 
 -- | The empty context.
 empty :: Ctx t
-empty = Ctx M.empty
+empty = Ctx M.empty mempty CtxEmpty
+
+-- | The hash for a single variable -> value binding.
+singletonHash :: Hashable t => Var -> t -> CtxHash
+singletonHash x t = CtxHash $ hashWithSalt (hash x) t
 
 -- | A singleton context.
-singleton :: Var -> t -> Ctx t
-singleton x t = Ctx (M.singleton x t)
+singleton :: Hashable t => Var -> t -> Ctx t
+singleton x t = Ctx (M.singleton x t) (singletonHash x t) (CtxSingle x)
 
 -- | Look up a variable in a context.
 lookup :: Var -> Ctx t -> Maybe t
-lookup x (Ctx c) = M.lookup x c
+lookup x (Ctx m _ _) = M.lookup x m
 
 -- | Look up a variable in a context in an ambient Reader effect.
 lookupR :: Has (Reader (Ctx t)) sig m => Var -> m (Maybe t)
 lookupR x = lookup x <$> ask
 
 -- | Delete a variable from a context.
-delete :: Var -> Ctx t -> Ctx t
-delete x (Ctx c) = Ctx (M.delete x c)
+delete :: Hashable t => Var -> Ctx t -> Ctx t
+delete x c@(Ctx m h s) = case M.lookup x m of
+  Nothing -> c
+  Just t -> Ctx (M.delete x m) (h - singletonHash x t) (CtxDelete x t s)
 
 -- | Get the list of key-value associations from a context.
 assocs :: Ctx t -> [(Var, t)]
@@ -87,18 +132,27 @@ vars = M.keys . unCtx
 
 -- | Add a key-value binding to a context (overwriting the old one if
 --   the key is already present).
-addBinding :: Var -> t -> Ctx t -> Ctx t
-addBinding x t (Ctx c) = Ctx (M.insert x t c)
+addBinding :: Hashable t => Var -> t -> Ctx t -> Ctx t
+addBinding x t (Ctx m h s) = Ctx (M.insert x t m) h' (CtxUnion s (CtxSingle x))
+ where
+  h' = case M.lookup x m of
+    Nothing -> h + singletonHash x t
+    Just t' -> h - singletonHash x t' + singletonHash x t
 
 -- | /Right/-biased union of contexts.
-union :: Ctx t -> Ctx t -> Ctx t
-union (Ctx c1) (Ctx c2) = Ctx (c2 `M.union` c1)
+union :: Hashable t => Ctx t -> Ctx t -> Ctx t
+union (Ctx m1 h1 s1) (Ctx m2 h2 s2) = Ctx (m2 `M.union` m1) h' (CtxUnion s1 s2)
+ where
+  -- `Data.Map.intersection l r` returns a map with common keys, but values from `l`
+  overwritten = M.intersection m1 m2
+  overwrittenHash = M.foldMapWithKey singletonHash overwritten
+  h' = h1 + h2 - overwrittenHash
 
 -- | Locally extend the context with an additional binding.
-withBinding :: Has (Reader (Ctx t)) sig m => Var -> t -> m a -> m a
+withBinding :: (Has (Reader (Ctx t)) sig m, Hashable t) => Var -> t -> m a -> m a
 withBinding x ty = local (addBinding x ty)
 
 -- | Locally extend the context with an additional context of
 --   bindings.
-withBindings :: Has (Reader (Ctx t)) sig m => Ctx t -> m a -> m a
+withBindings :: (Has (Reader (Ctx t)) sig m, Hashable t) => Ctx t -> m a -> m a
 withBindings ctx = local (`union` ctx)

From 7736a4769304d2b6608385baa9be6e30e23945f0 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Sat, 26 Oct 2024 17:06:42 -0500
Subject: [PATCH 02/17] [wip] start adding required Hashable instances

---
 src/swarm-lang/Swarm/Language/.#Types.hs      |  1 +
 src/swarm-lang/Swarm/Language/Context.hs      | 11 +++++++++--
 .../Swarm/Language/Requirements/Type.hs       |  2 +-
 src/swarm-lang/Swarm/Language/Typecheck.hs    |  2 +-
 src/swarm-lang/Swarm/Language/Types.hs        | 19 +++++++++++++------
 5 files changed, 25 insertions(+), 10 deletions(-)
 create mode 120000 src/swarm-lang/Swarm/Language/.#Types.hs

diff --git a/src/swarm-lang/Swarm/Language/.#Types.hs b/src/swarm-lang/Swarm/Language/.#Types.hs
new file mode 120000
index 000000000..6129f7487
--- /dev/null
+++ b/src/swarm-lang/Swarm/Language/.#Types.hs
@@ -0,0 +1 @@
+brent@diophantus.1236385:1729109353
\ No newline at end of file
diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index bef7c4c10..7e67658c1 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -104,10 +104,18 @@ empty = Ctx M.empty mempty CtxEmpty
 singletonHash :: Hashable t => Var -> t -> CtxHash
 singletonHash x t = CtxHash $ hashWithSalt (hash x) t
 
+-- | The hash for an entire Map's with of bindings.
+mapHash :: Hashable t => Map Var t -> CtxHash
+mapHash = M.foldMapWithKey singletonHash
+
 -- | A singleton context.
 singleton :: Hashable t => Var -> t -> Ctx t
 singleton x t = Ctx (M.singleton x t) (singletonHash x t) (CtxSingle x)
 
+-- | Create a Ctx from a Map.
+fromMap :: Hashable t => Map Var t -> Ctx t
+fromMap m = Ctx m (mapHash m) (M.foldrWithKey (\x _ -> CtxUnion (CtxSingle x)) CtxEmpty m)
+
 -- | Look up a variable in a context.
 lookup :: Var -> Ctx t -> Maybe t
 lookup x (Ctx m _ _) = M.lookup x m
@@ -145,8 +153,7 @@ union (Ctx m1 h1 s1) (Ctx m2 h2 s2) = Ctx (m2 `M.union` m1) h' (CtxUnion s1 s2)
  where
   -- `Data.Map.intersection l r` returns a map with common keys, but values from `l`
   overwritten = M.intersection m1 m2
-  overwrittenHash = M.foldMapWithKey singletonHash overwritten
-  h' = h1 + h2 - overwrittenHash
+  h' = h1 + h2 - mapHash overwritten
 
 -- | Locally extend the context with an additional binding.
 withBinding :: (Has (Reader (Ctx t)) sig m, Hashable t) => Var -> t -> m a -> m a
diff --git a/src/swarm-lang/Swarm/Language/Requirements/Type.hs b/src/swarm-lang/Swarm/Language/Requirements/Type.hs
index 5d74bc3b5..628a4b7a2 100644
--- a/src/swarm-lang/Swarm/Language/Requirements/Type.hs
+++ b/src/swarm-lang/Swarm/Language/Requirements/Type.hs
@@ -75,7 +75,7 @@ data Requirements = Requirements
   , devReqs :: Set Text
   , invReqs :: Map Text Int
   }
-  deriving (Eq, Ord, Show, Data, Generic, FromJSON, ToJSON)
+  deriving (Eq, Ord, Show, Data, Generic, FromJSON, ToJSON, Hashable)
 
 instance Semigroup Requirements where
   Requirements c1 d1 i1 <> Requirements c2 d2 i2 =
diff --git a/src/swarm-lang/Swarm/Language/Typecheck.hs b/src/swarm-lang/Swarm/Language/Typecheck.hs
index cb89e441b..6d8c702d8 100644
--- a/src/swarm-lang/Swarm/Language/Typecheck.hs
+++ b/src/swarm-lang/Swarm/Language/Typecheck.hs
@@ -405,7 +405,7 @@ skolemize (unPoly -> (xs, uty)) = do
   let xs' = map UTyVar skolemNames
       newSubst = M.fromList $ zip xs xs'
       s = M.mapKeys Left (newSubst `M.union` unCtx boundSubst)
-  pure (Ctx newSubst, substU s uty)
+  pure (Ctx.fromMap newSubst, substU s uty)
 
 -- | 'generalize' is the opposite of 'instantiate': add a 'Forall'
 --   which closes over all free type and unification variables.
diff --git a/src/swarm-lang/Swarm/Language/Types.hs b/src/swarm-lang/Swarm/Language/Types.hs
index d1198da9c..4c65437a8 100644
--- a/src/swarm-lang/Swarm/Language/Types.hs
+++ b/src/swarm-lang/Swarm/Language/Types.hs
@@ -131,6 +131,9 @@ import Data.Data (Data)
 import Data.Eq.Deriving (deriveEq1)
 import Data.Fix
 import Data.Foldable (fold)
+import Data.Functor.Classes (Eq1)
+import Data.Hashable (Hashable)
+import Data.Hashable.Lifted (Hashable1)
 import Data.Kind qualified
 import Data.List.NonEmpty ((<|))
 import Data.List.NonEmpty qualified as NE
@@ -178,7 +181,7 @@ data BaseTy
     BActor
   | -- | Keys, i.e. things that can be pressed on the keyboard
     BKey
-  deriving (Eq, Ord, Show, Bounded, Enum, Data, Generic, FromJSON, ToJSON)
+  deriving (Eq, Ord, Show, Bounded, Enum, Data, Generic, Hashable, FromJSON, ToJSON)
 
 baseTyName :: BaseTy -> Text
 baseTyName = into @Text . drop 1 . show
@@ -206,7 +209,7 @@ data TyCon
     TCFun
   | -- | User-defined type constructor.
     TCUser Var
-  deriving (Eq, Ord, Show, Data, Generic)
+  deriving (Eq, Ord, Show, Data, Generic, Hashable)
 
 instance ToJSON TyCon where
   toJSON = genericToJSON optionsMinimize
@@ -227,7 +230,7 @@ instance PrettyPrec TyCon where
 -- | The arity of a type, /i.e./ the number of type parameters it
 --   expects.
 newtype Arity = Arity {getArity :: Int}
-  deriving (Eq, Ord, Show, Generic, Data)
+  deriving (Eq, Ord, Show, Generic, Data, Hashable)
 
 instance ToJSON Arity where
   toJSON = genericToJSON optionsUnwrapUnary
@@ -246,7 +249,7 @@ instance PrettyPrec Arity where
 data Nat where
   NZ :: Nat
   NS :: Nat -> Nat
-  deriving (Eq, Ord, Show, Data, Generic, FromJSON, ToJSON)
+  deriving (Eq, Ord, Show, Data, Generic, Hashable, FromJSON, ToJSON)
 
 natToInt :: Nat -> Int
 natToInt NZ = 0
@@ -281,6 +284,8 @@ deriveEq1 ''TypeF
 deriveOrd1 ''TypeF
 deriveShow1 ''TypeF
 
+instance Hashable1 TypeF -- needs the Eq1 instance
+
 instance ToJSON1 TypeF where
   liftToJSON = genericLiftToJSON optionsMinimize
 
@@ -308,6 +313,8 @@ instance PrettyPrec IntVar where
 --   working with 'UType' as if it were defined directly.
 type UType = Free TypeF IntVar
 
+instance (Eq1 f, Hashable x, Hashable (f (Free f x))) => Hashable (Free f x)
+
 -- | A generic /fold/ for things defined via 'Free' (including, in
 --   particular, 'UType').
 ucata :: Functor t => (v -> a) -> (t a -> a) -> Free t v -> a
@@ -498,7 +505,7 @@ data ImplicitQuantification = Unquantified | Quantified
 --   only way to create a @Poly Quantified@ is through the 'quantify'
 --   function.
 data Poly (q :: ImplicitQuantification) t = Forall {_ptVars :: [Var], ptBody :: t}
-  deriving (Show, Eq, Functor, Foldable, Traversable, Data, Generic, FromJSON, ToJSON)
+  deriving (Show, Eq, Functor, Foldable, Traversable, Data, Generic, FromJSON, ToJSON, Hashable)
 
 -- | Create a raw, unquantified @Poly@ value.
 mkPoly :: [Var] -> t -> Poly 'Unquantified t
@@ -777,7 +784,7 @@ data TydefInfo = TydefInfo
   { _tydefType :: Polytype
   , _tydefArity :: Arity
   }
-  deriving (Eq, Show, Generic, Data, FromJSON, ToJSON)
+  deriving (Eq, Show, Generic, Data, FromJSON, ToJSON, Hashable)
 
 makeLenses ''TydefInfo
 

From d9c44e6b90cb9713ae8c3e373d1efcd2974f5e56 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Sat, 26 Oct 2024 22:28:13 -0500
Subject: [PATCH 03/17] finish adding Hashable instances

---
 src/swarm-lang/Swarm/Language/.#Types.hs         | 1 -
 src/swarm-lang/Swarm/Language/Key.hs             | 6 +++++-
 src/swarm-lang/Swarm/Language/Syntax/AST.hs      | 6 ++++--
 src/swarm-lang/Swarm/Language/Syntax/Comments.hs | 9 +++++----
 src/swarm-lang/Swarm/Language/Syntax/Loc.hs      | 5 +++--
 src/swarm-lang/Swarm/Language/Types.hs           | 5 +++--
 src/swarm-lang/Swarm/Language/Value.hs           | 5 +++--
 7 files changed, 23 insertions(+), 14 deletions(-)
 delete mode 120000 src/swarm-lang/Swarm/Language/.#Types.hs

diff --git a/src/swarm-lang/Swarm/Language/.#Types.hs b/src/swarm-lang/Swarm/Language/.#Types.hs
deleted file mode 120000
index 6129f7487..000000000
--- a/src/swarm-lang/Swarm/Language/.#Types.hs
+++ /dev/null
@@ -1 +0,0 @@
-brent@diophantus.1236385:1729109353
\ No newline at end of file
diff --git a/src/swarm-lang/Swarm/Language/Key.hs b/src/swarm-lang/Swarm/Language/Key.hs
index 788c6ba27..b4b2b0475 100644
--- a/src/swarm-lang/Swarm/Language/Key.hs
+++ b/src/swarm-lang/Swarm/Language/Key.hs
@@ -20,6 +20,7 @@ where
 
 import Data.Aeson (FromJSON, ToJSON)
 import Data.Foldable (asum)
+import Data.Hashable (Hashable)
 import Data.Kind qualified
 import Data.List (sort, (\\))
 import Data.Set (Set)
@@ -37,10 +38,13 @@ import Witch (from)
 ------------------------------------------------------------
 -- Parsing
 
+deriving instance Hashable V.Modifier
+deriving instance Hashable V.Key
+
 -- | A keyboard input, represented as a key + modifiers.  Invariant:
 --   the modifier list is always sorted.
 data KeyCombo = KeyCombo V.Key [V.Modifier]
-  deriving (Eq, Ord, Show, Generic, FromJSON, ToJSON)
+  deriving (Eq, Ord, Show, Generic, Hashable, FromJSON, ToJSON)
 
 deriving instance FromJSON V.Key
 deriving instance FromJSON V.Modifier
diff --git a/src/swarm-lang/Swarm/Language/Syntax/AST.hs b/src/swarm-lang/Swarm/Language/Syntax/AST.hs
index f281bdb3b..753fcaafa 100644
--- a/src/swarm-lang/Swarm/Language/Syntax/AST.hs
+++ b/src/swarm-lang/Swarm/Language/Syntax/AST.hs
@@ -16,6 +16,7 @@ import Control.Lens (Plated (..))
 import Data.Aeson.Types hiding (Key)
 import Data.Data (Data)
 import Data.Data.Lens (uniplate)
+import Data.Hashable (Hashable)
 import Data.Map.Strict (Map)
 import Data.Text (Text)
 import GHC.Generics (Generic)
@@ -37,7 +38,7 @@ data Syntax' ty = Syntax'
   , _sComments :: Comments
   , _sType :: ty
   }
-  deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic)
+  deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic, Hashable)
 
 instance Data ty => Plated (Syntax' ty) where
   plate = uniplate
@@ -46,7 +47,7 @@ instance Data ty => Plated (Syntax' ty) where
 --   as @def x = e1 end; e2@. This enumeration simply records which it
 --   was so that we can pretty-print appropriatly.
 data LetSyntax = LSLet | LSDef
-  deriving (Eq, Ord, Show, Bounded, Enum, Generic, Data, ToJSON, FromJSON)
+  deriving (Eq, Ord, Show, Bounded, Enum, Generic, Data, Hashable, ToJSON, FromJSON)
 
 ------------------------------------------------------------
 -- Term: basic syntax tree
@@ -148,6 +149,7 @@ data Term' ty
     , Foldable
     , Data
     , Generic
+    , Hashable
     , -- | The Traversable instance for Term (and for Syntax') is used during
       -- typechecking: during intermediate type inference, many of the type
       -- annotations placed on AST nodes will have unification variables in
diff --git a/src/swarm-lang/Swarm/Language/Syntax/Comments.hs b/src/swarm-lang/Swarm/Language/Syntax/Comments.hs
index ecdcd8e24..1e868dcfd 100644
--- a/src/swarm-lang/Swarm/Language/Syntax/Comments.hs
+++ b/src/swarm-lang/Swarm/Language/Syntax/Comments.hs
@@ -25,6 +25,7 @@ import Control.Lens (AsEmpty, makeLenses, pattern Empty)
 import Data.Aeson qualified as A
 import Data.Aeson.Types hiding (Key)
 import Data.Data (Data)
+import Data.Hashable (Hashable)
 import Data.Sequence (Seq)
 import Data.Text hiding (filter, length, map)
 import GHC.Generics (Generic)
@@ -34,12 +35,12 @@ import Swarm.Pretty (PrettyPrec (..))
 
 -- | Line vs block comments.
 data CommentType = LineComment | BlockComment
-  deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Data, ToJSON, FromJSON)
+  deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Data, Hashable, ToJSON, FromJSON)
 
 -- | Was a comment all by itself on a line, or did it occur after some
 --   other tokens on a line?
 data CommentSituation = StandaloneComment | SuffixComment
-  deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Data, ToJSON, FromJSON)
+  deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Data, Hashable, ToJSON, FromJSON)
 
 -- | Test whether a comment is a standalone comment or not.
 isStandalone :: Comment -> Bool
@@ -55,7 +56,7 @@ data Comment = Comment
   , commentSituation :: CommentSituation
   , commentText :: Text
   }
-  deriving (Eq, Show, Generic, Data, ToJSON, FromJSON)
+  deriving (Eq, Show, Generic, Data, ToJSON, FromJSON, Hashable)
 
 instance PrettyPrec Comment where
   prettyPrec _ (Comment _ LineComment _ txt) = "//" <> pretty txt
@@ -67,7 +68,7 @@ data Comments = Comments
   { _beforeComments :: Seq Comment
   , _afterComments :: Seq Comment
   }
-  deriving (Eq, Show, Generic, Data)
+  deriving (Eq, Show, Generic, Data, Hashable)
 
 makeLenses ''Comments
 
diff --git a/src/swarm-lang/Swarm/Language/Syntax/Loc.hs b/src/swarm-lang/Swarm/Language/Syntax/Loc.hs
index e3bdbf9a6..f09586523 100644
--- a/src/swarm-lang/Swarm/Language/Syntax/Loc.hs
+++ b/src/swarm-lang/Swarm/Language/Syntax/Loc.hs
@@ -14,6 +14,7 @@ module Swarm.Language.Syntax.Loc (
 
 import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
 import Data.Data (Data)
+import Data.Hashable (Hashable)
 import GHC.Generics (Generic)
 import Swarm.Language.Context (Var)
 import Swarm.Util.JSON (optionsUntagged)
@@ -28,7 +29,7 @@ data SrcLoc
   = NoLoc
   | -- | Half-open interval from start (inclusive) to end (exclusive)
     SrcLoc Int Int
-  deriving (Eq, Ord, Show, Data, Generic)
+  deriving (Eq, Ord, Show, Data, Generic, Hashable)
 
 instance ToJSON SrcLoc where
   toJSON = genericToJSON optionsUntagged
@@ -58,4 +59,4 @@ srcLocBefore _ _ = False
 --   binding sites. (Variable occurrences are a bare TVar which gets
 --   wrapped in a Syntax node, so we don't need LocVar for those.)
 data LocVar = LV {lvSrcLoc :: SrcLoc, lvVar :: Var}
-  deriving (Eq, Ord, Show, Data, Generic, FromJSON, ToJSON)
+  deriving (Eq, Ord, Show, Data, Generic, Hashable, FromJSON, ToJSON)
diff --git a/src/swarm-lang/Swarm/Language/Types.hs b/src/swarm-lang/Swarm/Language/Types.hs
index 4c65437a8..289b5e2c6 100644
--- a/src/swarm-lang/Swarm/Language/Types.hs
+++ b/src/swarm-lang/Swarm/Language/Types.hs
@@ -278,7 +278,7 @@ data TypeF t
     --   when pretty-printing; the actual bound variables are represented
     --   via de Bruijn indices.
     TyRecF Var t
-  deriving (Show, Eq, Ord, Functor, Foldable, Traversable, Generic, Generic1, Data)
+  deriving (Show, Eq, Ord, Functor, Foldable, Traversable, Generic, Generic1, Data, Hashable)
 
 deriveEq1 ''TypeF
 deriveOrd1 ''TypeF
@@ -299,7 +299,7 @@ instance FromJSON1 TypeF where
 type Type = Fix TypeF
 
 newtype IntVar = IntVar Int
-  deriving (Show, Data, Eq, Ord)
+  deriving (Show, Data, Eq, Ord, Generic, Hashable)
 
 instance PrettyPrec IntVar where
   prettyPrec _ = pretty . mkVarName "u"
@@ -313,6 +313,7 @@ instance PrettyPrec IntVar where
 --   working with 'UType' as if it were defined directly.
 type UType = Free TypeF IntVar
 
+-- XXX orphan instance
 instance (Eq1 f, Hashable x, Hashable (f (Free f x))) => Hashable (Free f x)
 
 -- | A generic /fold/ for things defined via 'Free' (including, in
diff --git a/src/swarm-lang/Swarm/Language/Value.hs b/src/swarm-lang/Swarm/Language/Value.hs
index c1b25e2ca..802146cc5 100644
--- a/src/swarm-lang/Swarm/Language/Value.hs
+++ b/src/swarm-lang/Swarm/Language/Value.hs
@@ -29,6 +29,7 @@ module Swarm.Language.Value (
 
 import Control.Lens hiding (Const)
 import Data.Bool (bool)
+import Data.Hashable (Hashable)
 import Data.List (foldl')
 import Data.Map (Map)
 import Data.Map qualified as M
@@ -118,7 +119,7 @@ data Value where
   --   <http://www.lel.ed.ac.uk/~gpullum/loopsnoop.html cannot detect
   --   /all/ infinite loops this way>.)
   VBlackhole :: Value
-  deriving (Eq, Show, Generic)
+  deriving (Eq, Show, Generic, Hashable)
 
 -- | A value context is a mapping from variable names to their runtime
 --   values.
@@ -141,7 +142,7 @@ data Env = Env
   , _envTydefs :: TDCtx
   -- ^ Type synonym definitions.
   }
-  deriving (Eq, Show, Generic)
+  deriving (Eq, Show, Generic, Hashable)
 
 makeLenses ''Env
 

From 2ba1a386dcb10c6a5e782b9a1ec361700c916877 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Sun, 27 Oct 2024 21:29:27 -0500
Subject: [PATCH 04/17] context folding

---
 src/swarm-lang/Swarm/Language/Context.hs | 107 +++++++++++++++++------
 1 file changed, 82 insertions(+), 25 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index 7e67658c1..e1dfd429b 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -1,8 +1,6 @@
 {-# LANGUAGE DeriveDataTypeable #-}
 {-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE PatternSynonyms #-}
-{-# LANGUAGE ViewPatterns #-}
 
 -- |
 -- SPDX-License-Identifier: BSD-3-Clause
@@ -13,31 +11,39 @@ module Swarm.Language.Context where
 
 import Control.Algebra (Has)
 import Control.Effect.Reader (Reader, ask, local)
-import Control.Lens.Empty (AsEmpty (..), pattern Empty)
+import Control.Effect.State (State, get, modify)
+import Control.Lens.Empty (AsEmpty (..))
 import Control.Lens.Prism (prism)
 import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
 import Data.Data (Data)
+import Data.Functor ((<&>))
 import Data.Hashable
 import Data.Map (Map)
 import Data.Map qualified as M
 import Data.Semigroup (Sum (..))
+import Data.Set (Set)
+import Data.Set qualified as S
 import Data.Text (Text)
 import GHC.Generics (Generic)
-import Prettyprinter (brackets, emptyDoc, hsep, punctuate)
-import Swarm.Pretty (PrettyPrec (..), prettyBinding)
+import Swarm.Pretty (PrettyPrec (..))
 import Swarm.Util.JSON (optionsUnwrapUnary)
 import Prelude hiding (lookup)
 
 -- | We use 'Text' values to represent variables.
 type Var = Text
 
+-- | A "shadow context" records the hash values and structure of a
+--   context but does not record the actual values associated to the
+--   variables.
+type SCtx t = (CtxHash, CtxStruct t)
+
 -- | A data type to record the structure of how a context was built,
 --   so that we can later destruct/serialize it effectively.
 data CtxStruct t
   = CtxEmpty
   | CtxSingle Var
-  | CtxDelete Var t (CtxStruct t)
-  | CtxUnion (CtxStruct t) (CtxStruct t)
+  | CtxDelete Var t (SCtx t)
+  | CtxUnion (SCtx t) (SCtx t)
   deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic, ToJSON, FromJSON)
 
 -- | A context hash is a hash value used to identify contexts without
@@ -53,13 +59,22 @@ data CtxStruct t
 --   the need to thread around some kind of globally unique ID
 --   generation effect.
 newtype CtxHash = CtxHash {getCtxHash :: Int}
-  deriving (Eq, Show, Data, Generic, ToJSON, FromJSON)
+  deriving (Eq, Ord, Show, Data, Generic, ToJSON, FromJSON)
   deriving (Semigroup, Monoid) via Sum Int
   deriving (Num) via Int
 
 -- | A context is a mapping from variable names to things.
-data Ctx t = Ctx {unCtx :: Map Var t, ctxHash :: CtxHash, ctxStruct :: CtxStruct t}
-  deriving (Eq, Show, Functor, Traversable, Data, Generic)
+data Ctx t = Ctx {unCtx :: Map Var t, sctx :: SCtx t}
+  deriving (Eq, Functor, Traversable, Data, Generic)
+
+ctxHash :: Ctx t -> CtxHash
+ctxHash (Ctx _ (h, _)) = h
+
+ctxStruct :: Ctx t -> CtxStruct t
+ctxStruct (Ctx _ (_, s)) = s
+
+instance Show (Ctx t) where
+  show _ = "<Ctx>"
 
 instance Hashable t => Hashable (Ctx t) where
   hash = getCtxHash . ctxHash
@@ -68,6 +83,45 @@ instance Hashable t => Hashable (Ctx t) where
 instance Foldable Ctx where
   foldMap f = foldMap f . unCtx
 
+-- Fold a context with sharing.  XXX
+foldCtx ::
+  (Has (State (Set CtxHash)) sig m, Has (Reader (Map Var t)) sig m) =>
+  r ->
+  (Var -> t -> m r) ->
+  (Var -> r -> r) ->
+  (r -> r -> r) ->
+  (CtxHash -> r) ->
+  Ctx t ->
+  m r
+foldCtx e sg del un sn (Ctx m s) = foldSCtx e sg del un sn m s
+
+-- XXX
+foldSCtx ::
+  forall sig m t r.
+  (Has (State (Set CtxHash)) sig m, Has (Reader (Map Var t)) sig m) =>
+  r ->
+  (Var -> t -> m r) ->
+  (Var -> r -> r) ->
+  (r -> r -> r) ->
+  (CtxHash -> r) ->
+  Map Var t ->
+  SCtx t ->
+  m r
+foldSCtx e sg del un sn m = go
+ where
+  go :: SCtx t -> m r
+  go (h, s) = do
+    seen <- get
+    case h `S.member` seen of
+      True -> pure $ sn h
+      False -> do
+        modify (S.insert h)
+        case s of
+          CtxEmpty -> pure e
+          CtxSingle x -> sg x (m M.! x)
+          CtxDelete x t sc -> local (M.insert x t) (go sc) <&> del x
+          CtxUnion sc1 sc2 -> un <$> go sc1 <*> go sc2
+
 -- XXX this instance will have to change!!
 instance ToJSON t => ToJSON (Ctx t) where
   toJSON = genericToJSON optionsUnwrapUnary
@@ -76,10 +130,8 @@ instance ToJSON t => ToJSON (Ctx t) where
 instance FromJSON t => FromJSON (Ctx t) where
   parseJSON = genericParseJSON optionsUnwrapUnary
 
--- XXX this instance will have to change!!
 instance (PrettyPrec t) => PrettyPrec (Ctx t) where
-  prettyPrec _ Empty = emptyDoc
-  prettyPrec _ (assocs -> bs) = brackets (hsep (punctuate "," (map prettyBinding bs)))
+  prettyPrec _ _ = "<Ctx>"
 
 -- | The semigroup operation for contexts is /right/-biased union.
 instance Hashable t => Semigroup (Ctx t) where
@@ -98,7 +150,7 @@ instance AsEmpty (Ctx t) where
 
 -- | The empty context.
 empty :: Ctx t
-empty = Ctx M.empty mempty CtxEmpty
+empty = Ctx M.empty (mempty, CtxEmpty)
 
 -- | The hash for a single variable -> value binding.
 singletonHash :: Hashable t => Var -> t -> CtxHash
@@ -110,15 +162,15 @@ mapHash = M.foldMapWithKey singletonHash
 
 -- | A singleton context.
 singleton :: Hashable t => Var -> t -> Ctx t
-singleton x t = Ctx (M.singleton x t) (singletonHash x t) (CtxSingle x)
+singleton x t = Ctx (M.singleton x t) (singletonHash x t, CtxSingle x)
 
 -- | Create a Ctx from a Map.
 fromMap :: Hashable t => Map Var t -> Ctx t
-fromMap m = Ctx m (mapHash m) (M.foldrWithKey (\x _ -> CtxUnion (CtxSingle x)) CtxEmpty m)
+fromMap m = Ctx m (M.foldrWithKey (insertSCtx Nothing) (mempty, CtxEmpty) m)
 
 -- | Look up a variable in a context.
 lookup :: Var -> Ctx t -> Maybe t
-lookup x (Ctx m _ _) = M.lookup x m
+lookup x (Ctx m _) = M.lookup x m
 
 -- | Look up a variable in a context in an ambient Reader effect.
 lookupR :: Has (Reader (Ctx t)) sig m => Var -> m (Maybe t)
@@ -126,9 +178,9 @@ lookupR x = lookup x <$> ask
 
 -- | Delete a variable from a context.
 delete :: Hashable t => Var -> Ctx t -> Ctx t
-delete x c@(Ctx m h s) = case M.lookup x m of
+delete x c@(Ctx m s@(h, _)) = case M.lookup x m of
   Nothing -> c
-  Just t -> Ctx (M.delete x m) (h - singletonHash x t) (CtxDelete x t s)
+  Just t -> Ctx (M.delete x m) (h - singletonHash x t, CtxDelete x t s)
 
 -- | Get the list of key-value associations from a context.
 assocs :: Ctx t -> [(Var, t)]
@@ -138,18 +190,23 @@ assocs = M.assocs . unCtx
 vars :: Ctx t -> [Var]
 vars = M.keys . unCtx
 
+-- | XXX
+insertSCtx :: Hashable t => Maybe t -> Var -> t -> SCtx t -> SCtx t
+insertSCtx old x new s@(h, _) = (h', CtxUnion s (tHash, CtxSingle x))
+ where
+  tHash = singletonHash x new
+  h' = case old of
+    Nothing -> h + tHash
+    Just t' -> h - singletonHash x t' + tHash
+
 -- | Add a key-value binding to a context (overwriting the old one if
 --   the key is already present).
 addBinding :: Hashable t => Var -> t -> Ctx t -> Ctx t
-addBinding x t (Ctx m h s) = Ctx (M.insert x t m) h' (CtxUnion s (CtxSingle x))
- where
-  h' = case M.lookup x m of
-    Nothing -> h + singletonHash x t
-    Just t' -> h - singletonHash x t' + singletonHash x t
+addBinding x t (Ctx m s) = Ctx (M.insert x t m) (insertSCtx (M.lookup x m) x t s)
 
 -- | /Right/-biased union of contexts.
 union :: Hashable t => Ctx t -> Ctx t -> Ctx t
-union (Ctx m1 h1 s1) (Ctx m2 h2 s2) = Ctx (m2 `M.union` m1) h' (CtxUnion s1 s2)
+union (Ctx m1 s1@(h1, _)) (Ctx m2 s2@(h2, _)) = Ctx (m2 `M.union` m1) (h', CtxUnion s1 s2)
  where
   -- `Data.Map.intersection l r` returns a map with common keys, but values from `l`
   overwritten = M.intersection m1 m2

From ae87867dd810e708e8b7e90f31eac79e5fefb93f Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Mon, 28 Oct 2024 14:00:31 -0500
Subject: [PATCH 05/17] better type for foldCtx

---
 src/swarm-lang/Swarm/Language/Context.hs | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index e1dfd429b..ed1b3cbad 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -10,6 +10,7 @@
 module Swarm.Language.Context where
 
 import Control.Algebra (Has)
+import Control.Carrier.Reader (ReaderC, runReader)
 import Control.Effect.Reader (Reader, ask, local)
 import Control.Effect.State (State, get, modify)
 import Control.Lens.Empty (AsEmpty (..))
@@ -85,9 +86,9 @@ instance Foldable Ctx where
 
 -- Fold a context with sharing.  XXX
 foldCtx ::
-  (Has (State (Set CtxHash)) sig m, Has (Reader (Map Var t)) sig m) =>
+  Has (State (Set CtxHash)) sig m =>
   r ->
-  (Var -> t -> m r) ->
+  (forall sig' m'. Has (State (Set CtxHash)) sig' m' => Var -> t -> m' r) ->
   (Var -> r -> r) ->
   (r -> r -> r) ->
   (CtxHash -> r) ->
@@ -98,18 +99,18 @@ foldCtx e sg del un sn (Ctx m s) = foldSCtx e sg del un sn m s
 -- XXX
 foldSCtx ::
   forall sig m t r.
-  (Has (State (Set CtxHash)) sig m, Has (Reader (Map Var t)) sig m) =>
+  Has (State (Set CtxHash)) sig m =>
   r ->
-  (Var -> t -> m r) ->
+  (forall sig' m'. Has (State (Set CtxHash)) sig' m' => Var -> t -> m' r) ->
   (Var -> r -> r) ->
   (r -> r -> r) ->
   (CtxHash -> r) ->
   Map Var t ->
   SCtx t ->
   m r
-foldSCtx e sg del un sn m = go
+foldSCtx e sg del un sn m = runReader m . go
  where
-  go :: SCtx t -> m r
+  go :: SCtx t -> ReaderC (Map Var t) m r
   go (h, s) = do
     seen <- get
     case h `S.member` seen of

From bd230ca9b580379cd5331b49b0d5b12a919f1c65 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Mon, 28 Oct 2024 14:03:02 -0500
Subject: [PATCH 06/17] even better type for foldCtx; work on getting rid of
 Empty base case for fromMap

---
 src/swarm-lang/Swarm/Language/Context.hs | 41 ++++++++++++++++++++----
 1 file changed, 34 insertions(+), 7 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index ed1b3cbad..e0d1d08f8 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -9,8 +9,9 @@
 -- types, values, or capability sets) used throughout the codebase.
 module Swarm.Language.Context where
 
-import Control.Algebra (Has)
+import Control.Algebra (Has, run)
 import Control.Carrier.Reader (ReaderC, runReader)
+import Control.Carrier.State.Strict (evalState)
 import Control.Effect.Reader (Reader, ask, local)
 import Control.Effect.State (State, get, modify)
 import Control.Lens.Empty (AsEmpty (..))
@@ -19,6 +20,7 @@ import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
 import Data.Data (Data)
 import Data.Functor ((<&>))
 import Data.Hashable
+import Data.List.NonEmpty qualified as NE
 import Data.Map (Map)
 import Data.Map qualified as M
 import Data.Semigroup (Sum (..))
@@ -28,6 +30,7 @@ import Data.Text (Text)
 import GHC.Generics (Generic)
 import Swarm.Pretty (PrettyPrec (..))
 import Swarm.Util.JSON (optionsUnwrapUnary)
+import Text.Printf (printf)
 import Prelude hiding (lookup)
 
 -- | We use 'Text' values to represent variables.
@@ -60,10 +63,13 @@ data CtxStruct t
 --   the need to thread around some kind of globally unique ID
 --   generation effect.
 newtype CtxHash = CtxHash {getCtxHash :: Int}
-  deriving (Eq, Ord, Show, Data, Generic, ToJSON, FromJSON)
+  deriving (Eq, Ord, Data, Generic, ToJSON, FromJSON)
   deriving (Semigroup, Monoid) via Sum Int
   deriving (Num) via Int
 
+instance Show CtxHash where
+  show (CtxHash h) = printf "%016x" h
+
 -- | A context is a mapping from variable names to things.
 data Ctx t = Ctx {unCtx :: Map Var t, sctx :: SCtx t}
   deriving (Eq, Functor, Traversable, Data, Generic)
@@ -84,8 +90,18 @@ instance Hashable t => Hashable (Ctx t) where
 instance Foldable Ctx where
   foldMap f = foldMap f . unCtx
 
--- Fold a context with sharing.  XXX
 foldCtx ::
+  r ->
+  (forall sig' m'. Has (State (Set CtxHash)) sig' m' => Var -> t -> m' r) ->
+  (Var -> r -> r) ->
+  (r -> r -> r) ->
+  (CtxHash -> r) ->
+  Ctx t ->
+  r
+foldCtx e sg del un sn = run . evalState @(Set CtxHash) S.empty . foldCtxWith e sg del un sn
+
+-- Fold a context with sharing.  XXX
+foldCtxWith ::
   Has (State (Set CtxHash)) sig m =>
   r ->
   (forall sig' m'. Has (State (Set CtxHash)) sig' m' => Var -> t -> m' r) ->
@@ -94,7 +110,7 @@ foldCtx ::
   (CtxHash -> r) ->
   Ctx t ->
   m r
-foldCtx e sg del un sn (Ctx m s) = foldSCtx e sg del un sn m s
+foldCtxWith e sg del un sn (Ctx m s) = foldSCtx e sg del un sn m s
 
 -- XXX
 foldSCtx ::
@@ -157,17 +173,26 @@ empty = Ctx M.empty (mempty, CtxEmpty)
 singletonHash :: Hashable t => Var -> t -> CtxHash
 singletonHash x t = CtxHash $ hashWithSalt (hash x) t
 
+singletonSCtx :: Hashable t => Var -> t -> SCtx t
+singletonSCtx x t = (singletonHash x t, CtxSingle x)
+
 -- | The hash for an entire Map's with of bindings.
 mapHash :: Hashable t => Map Var t -> CtxHash
 mapHash = M.foldMapWithKey singletonHash
 
 -- | A singleton context.
 singleton :: Hashable t => Var -> t -> Ctx t
-singleton x t = Ctx (M.singleton x t) (singletonHash x t, CtxSingle x)
+singleton x t = Ctx (M.singleton x t) (singletonSCtx x t)
 
 -- | Create a Ctx from a Map.
 fromMap :: Hashable t => Map Var t -> Ctx t
-fromMap m = Ctx m (M.foldrWithKey (insertSCtx Nothing) (mempty, CtxEmpty) m)
+fromMap m = Ctx m mapStruct
+ where
+  mapStruct = case NE.nonEmpty (map (uncurry singletonSCtx) (M.assocs m)) of
+    Nothing -> _
+    Just ne -> foldr1 CtxUnion ne
+
+-- (M.foldrWithKey (insertSCtx Nothing) (mempty, CtxEmpty) m)
 
 -- | Look up a variable in a context.
 lookup :: Var -> Ctx t -> Maybe t
@@ -193,7 +218,7 @@ vars = M.keys . unCtx
 
 -- | XXX
 insertSCtx :: Hashable t => Maybe t -> Var -> t -> SCtx t -> SCtx t
-insertSCtx old x new s@(h, _) = (h', CtxUnion s (tHash, CtxSingle x))
+insertSCtx old x new s@(h, _) = (h', CtxUnion (tHash, CtxSingle x) s)
  where
   tHash = singletonHash x new
   h' = case old of
@@ -205,6 +230,8 @@ insertSCtx old x new s@(h, _) = (h', CtxUnion s (tHash, CtxSingle x))
 addBinding :: Hashable t => Var -> t -> Ctx t -> Ctx t
 addBinding x t (Ctx m s) = Ctx (M.insert x t m) (insertSCtx (M.lookup x m) x t s)
 
+-- XXX how to encode structure of unioned context when there are overlapping variables?
+
 -- | /Right/-biased union of contexts.
 union :: Hashable t => Ctx t -> Ctx t -> Ctx t
 union (Ctx m1 s1@(h1, _)) (Ctx m2 s2@(h2, _)) = Ctx (m2 `M.union` m1) (h', CtxUnion s1 s2)

From 48cdd5855987f633665514b0b713c56f220433d9 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Sat, 2 Nov 2024 17:13:52 -0500
Subject: [PATCH 07/17] context dessicate/rehydrate works!!

---
 src/swarm-lang/Swarm/Language/Context.hs | 231 +++++++++++------------
 1 file changed, 115 insertions(+), 116 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index e0d1d08f8..be5b11bd3 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -10,22 +10,19 @@
 module Swarm.Language.Context where
 
 import Control.Algebra (Has, run)
-import Control.Carrier.Reader (ReaderC, runReader)
-import Control.Carrier.State.Strict (evalState)
+import Control.Carrier.State.Strict (execState)
 import Control.Effect.Reader (Reader, ask, local)
 import Control.Effect.State (State, get, modify)
 import Control.Lens.Empty (AsEmpty (..))
 import Control.Lens.Prism (prism)
 import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
 import Data.Data (Data)
-import Data.Functor ((<&>))
+import Data.Functor.Const
 import Data.Hashable
 import Data.List.NonEmpty qualified as NE
 import Data.Map (Map)
 import Data.Map qualified as M
 import Data.Semigroup (Sum (..))
-import Data.Set (Set)
-import Data.Set qualified as S
 import Data.Text (Text)
 import GHC.Generics (Generic)
 import Swarm.Pretty (PrettyPrec (..))
@@ -36,32 +33,21 @@ import Prelude hiding (lookup)
 -- | We use 'Text' values to represent variables.
 type Var = Text
 
--- | A "shadow context" records the hash values and structure of a
---   context but does not record the actual values associated to the
---   variables.
-type SCtx t = (CtxHash, CtxStruct t)
-
--- | A data type to record the structure of how a context was built,
---   so that we can later destruct/serialize it effectively.
-data CtxStruct t
-  = CtxEmpty
-  | CtxSingle Var
-  | CtxDelete Var t (SCtx t)
-  | CtxUnion (SCtx t) (SCtx t)
-  deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic, ToJSON, FromJSON)
+------------------------------------------------------------
+-- Context hash
 
 -- | A context hash is a hash value used to identify contexts without
---   having to compare them for equality.  Hash values are computed in
+--   having to compare them for equality.  Hash values are computed
 --   homomorphically, so that two equal contexts will be guaranteed to
---   have the same hash value, even if they were constructed
---   by a different sequence of operations.
+--   have the same hash value, even if they were constructed with a
+--   different sequence of operations.
 --
 --   The downside of this approach is that, /in theory/, there could
---   be a hash collision---two different contexts which nonetheless
---   have the same hash value.  However, this is extremely unlikely.
---   The benefit is that everything can be purely functional, without
---   the need to thread around some kind of globally unique ID
---   generation effect.
+--   be hash collisions, that is, two different contexts which
+--   nonetheless have the same hash value.  However, this is extremely
+--   unlikely.  The benefit is that everything can be purely
+--   functional, without the need to thread around some kind of
+--   globally unique ID generation effect.
 newtype CtxHash = CtxHash {getCtxHash :: Int}
   deriving (Eq, Ord, Data, Generic, ToJSON, FromJSON)
   deriving (Semigroup, Monoid) via Sum Int
@@ -70,15 +56,61 @@ newtype CtxHash = CtxHash {getCtxHash :: Int}
 instance Show CtxHash where
   show (CtxHash h) = printf "%016x" h
 
--- | A context is a mapping from variable names to things.
-data Ctx t = Ctx {unCtx :: Map Var t, sctx :: SCtx t}
+-- | The hash for a single variable -> value binding.
+singletonHash :: Hashable t => Var -> t -> CtxHash
+singletonHash x t = CtxHash $ hashWithSalt (hash x) t
+
+-- | The hash for an entire Map's worth of bindings.
+mapHash :: Hashable t => Map Var t -> CtxHash
+mapHash = M.foldMapWithKey singletonHash
+
+------------------------------------------------------------
+-- Context structure
+
+-- | 'CtxF' represents one level of structure of a context: a context
+--   is either empty, a singleton, or built via deletion or union.
+data CtxF f t
+  = CtxEmpty
+  | CtxSingle Var t
+  | CtxDelete Var t (f t)
+  | CtxUnion (f t) (f t)
+  deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic, ToJSON, FromJSON)
+
+-- | Map over the recursive structure stored in a 'CtxF'.
+restructure :: (f t -> g t) -> CtxF f t -> CtxF g t
+restructure _ CtxEmpty = CtxEmpty
+restructure _ (CtxSingle x t) = CtxSingle x t
+restructure h (CtxDelete x t f1) = CtxDelete x t (h f1)
+restructure h (CtxUnion f1 f2) = CtxUnion (h f1) (h f2)
+
+-- | A "context structure" is one possible representation of a
+--   context, consisting of a structured record of the process by
+--   which a context was constructed.  This representation would be
+--   terrible for doing efficient variable lookups, but it can be used
+--   to efficiently destruct/serialize the context while recovering
+--   sharing.
+--
+--   It stores a top-level hash of the context, along with a recursive
+--   tree built via 'CtxF'.
+data CtxStruct t = CtxStruct CtxHash (CtxF CtxStruct t)
+  deriving (Eq, Functor, Foldable, Traversable, Data, Generic, ToJSON, FromJSON, Show)
+
+-- | A 'CtxNode' is just a single level of structure for a context,
+--   with any recursive contexts replaced by their hash.
+type CtxNode t = CtxF (Const CtxHash) t
+
+------------------------------------------------------------
+-- Contexts
+
+-- | A context is a mapping from variable names to things.  We store
+--   both a 'Map' (for efficient lookup) as well as a 'CtxStruct' for
+--   sharing-aware serializing/deserializing of contexts.
+data Ctx t = Ctx {unCtx :: Map Var t, ctxStruct :: CtxStruct t}
   deriving (Eq, Functor, Traversable, Data, Generic)
 
+-- | Get the top-level hash of a context.
 ctxHash :: Ctx t -> CtxHash
-ctxHash (Ctx _ (h, _)) = h
-
-ctxStruct :: Ctx t -> CtxStruct t
-ctxStruct (Ctx _ (_, s)) = s
+ctxHash (Ctx _ (CtxStruct h _)) = h
 
 instance Show (Ctx t) where
   show _ = "<Ctx>"
@@ -90,54 +122,8 @@ instance Hashable t => Hashable (Ctx t) where
 instance Foldable Ctx where
   foldMap f = foldMap f . unCtx
 
-foldCtx ::
-  r ->
-  (forall sig' m'. Has (State (Set CtxHash)) sig' m' => Var -> t -> m' r) ->
-  (Var -> r -> r) ->
-  (r -> r -> r) ->
-  (CtxHash -> r) ->
-  Ctx t ->
-  r
-foldCtx e sg del un sn = run . evalState @(Set CtxHash) S.empty . foldCtxWith e sg del un sn
-
--- Fold a context with sharing.  XXX
-foldCtxWith ::
-  Has (State (Set CtxHash)) sig m =>
-  r ->
-  (forall sig' m'. Has (State (Set CtxHash)) sig' m' => Var -> t -> m' r) ->
-  (Var -> r -> r) ->
-  (r -> r -> r) ->
-  (CtxHash -> r) ->
-  Ctx t ->
-  m r
-foldCtxWith e sg del un sn (Ctx m s) = foldSCtx e sg del un sn m s
-
--- XXX
-foldSCtx ::
-  forall sig m t r.
-  Has (State (Set CtxHash)) sig m =>
-  r ->
-  (forall sig' m'. Has (State (Set CtxHash)) sig' m' => Var -> t -> m' r) ->
-  (Var -> r -> r) ->
-  (r -> r -> r) ->
-  (CtxHash -> r) ->
-  Map Var t ->
-  SCtx t ->
-  m r
-foldSCtx e sg del un sn m = runReader m . go
- where
-  go :: SCtx t -> ReaderC (Map Var t) m r
-  go (h, s) = do
-    seen <- get
-    case h `S.member` seen of
-      True -> pure $ sn h
-      False -> do
-        modify (S.insert h)
-        case s of
-          CtxEmpty -> pure e
-          CtxSingle x -> sg x (m M.! x)
-          CtxDelete x t sc -> local (M.insert x t) (go sc) <&> del x
-          CtxUnion sc1 sc2 -> un <$> go sc1 <*> go sc2
+------------------------------------------------------------
+-- Context instances
 
 -- XXX this instance will have to change!!
 instance ToJSON t => ToJSON (Ctx t) where
@@ -165,34 +151,22 @@ instance AsEmpty (Ctx t) where
       | M.null (unCtx c) = Right ()
       | otherwise = Left c
 
+------------------------------------------------------------
+-- Context operations
+
 -- | The empty context.
 empty :: Ctx t
-empty = Ctx M.empty (mempty, CtxEmpty)
-
--- | The hash for a single variable -> value binding.
-singletonHash :: Hashable t => Var -> t -> CtxHash
-singletonHash x t = CtxHash $ hashWithSalt (hash x) t
-
-singletonSCtx :: Hashable t => Var -> t -> SCtx t
-singletonSCtx x t = (singletonHash x t, CtxSingle x)
-
--- | The hash for an entire Map's with of bindings.
-mapHash :: Hashable t => Map Var t -> CtxHash
-mapHash = M.foldMapWithKey singletonHash
+empty = Ctx M.empty (CtxStruct mempty CtxEmpty)
 
 -- | A singleton context.
 singleton :: Hashable t => Var -> t -> Ctx t
-singleton x t = Ctx (M.singleton x t) (singletonSCtx x t)
+singleton x t = Ctx (M.singleton x t) (CtxStruct (singletonHash x t) (CtxSingle x t))
 
--- | Create a Ctx from a Map.
+-- | Create a 'Ctx' from a 'Map'.
 fromMap :: Hashable t => Map Var t -> Ctx t
-fromMap m = Ctx m mapStruct
- where
-  mapStruct = case NE.nonEmpty (map (uncurry singletonSCtx) (M.assocs m)) of
-    Nothing -> _
-    Just ne -> foldr1 CtxUnion ne
-
--- (M.foldrWithKey (insertSCtx Nothing) (mempty, CtxEmpty) m)
+fromMap m = case NE.nonEmpty (M.assocs m) of
+  Nothing -> empty
+  Just ne -> foldr1 union (NE.map (uncurry singleton) ne)
 
 -- | Look up a variable in a context.
 lookup :: Var -> Ctx t -> Maybe t
@@ -204,9 +178,9 @@ lookupR x = lookup x <$> ask
 
 -- | Delete a variable from a context.
 delete :: Hashable t => Var -> Ctx t -> Ctx t
-delete x c@(Ctx m s@(h, _)) = case M.lookup x m of
+delete x c@(Ctx m s@(CtxStruct h _)) = case M.lookup x m of
   Nothing -> c
-  Just t -> Ctx (M.delete x m) (h - singletonHash x t, CtxDelete x t s)
+  Just t -> Ctx (M.delete x m) (CtxStruct (h - singletonHash x t) (CtxDelete x t s))
 
 -- | Get the list of key-value associations from a context.
 assocs :: Ctx t -> [(Var, t)]
@@ -216,25 +190,20 @@ assocs = M.assocs . unCtx
 vars :: Ctx t -> [Var]
 vars = M.keys . unCtx
 
--- | XXX
-insertSCtx :: Hashable t => Maybe t -> Var -> t -> SCtx t -> SCtx t
-insertSCtx old x new s@(h, _) = (h', CtxUnion (tHash, CtxSingle x) s)
- where
-  tHash = singletonHash x new
-  h' = case old of
-    Nothing -> h + tHash
-    Just t' -> h - singletonHash x t' + tHash
-
 -- | Add a key-value binding to a context (overwriting the old one if
 --   the key is already present).
 addBinding :: Hashable t => Var -> t -> Ctx t -> Ctx t
-addBinding x t (Ctx m s) = Ctx (M.insert x t m) (insertSCtx (M.lookup x m) x t s)
-
--- XXX how to encode structure of unioned context when there are overlapping variables?
+addBinding x t (Ctx m s@(CtxStruct h _)) = Ctx (M.insert x t m) s'
+ where
+  s' = CtxStruct h' (CtxUnion (CtxStruct tHash (CtxSingle x t)) s)
+  tHash = singletonHash x t
+  h' = case M.lookup x m of
+    Nothing -> h + tHash
+    Just t' -> h - singletonHash x t' + tHash
 
 -- | /Right/-biased union of contexts.
 union :: Hashable t => Ctx t -> Ctx t -> Ctx t
-union (Ctx m1 s1@(h1, _)) (Ctx m2 s2@(h2, _)) = Ctx (m2 `M.union` m1) (h', CtxUnion s1 s2)
+union (Ctx m1 s1@(CtxStruct h1 _)) (Ctx m2 s2@(CtxStruct h2 _)) = Ctx (m2 `M.union` m1) (CtxStruct h' (CtxUnion s1 s2))
  where
   -- `Data.Map.intersection l r` returns a map with common keys, but values from `l`
   overwritten = M.intersection m1 m2
@@ -248,3 +217,33 @@ withBinding x ty = local (addBinding x ty)
 --   bindings.
 withBindings :: (Has (Reader (Ctx t)) sig m, Hashable t) => Ctx t -> m a -> m a
 withBindings ctx = local (`union` ctx)
+
+------------------------------------------------------------
+-- Context serializing/deserializing
+
+type CtxMap f t = Map CtxHash (CtxF f t)
+
+toCtxMap :: Ctx t -> CtxMap CtxStruct t
+toCtxMap (Ctx m s) = run $ execState M.empty (buildCtxMap m s)
+
+buildCtxMap :: forall t m sig. Has (State (CtxMap CtxStruct t)) sig m => Map Var t -> CtxStruct t -> m ()
+buildCtxMap m (CtxStruct h s) = do
+  cm <- get @(CtxMap CtxStruct t)
+  case h `M.member` cm of
+    True -> pure ()
+    False -> do
+      modify (M.insert h s)
+      case s of
+        CtxEmpty -> pure ()
+        CtxSingle {} -> pure ()
+        CtxDelete x t s1 -> buildCtxMap (M.insert x t m) s1
+        CtxUnion s1 s2 -> buildCtxMap m s1 *> buildCtxMap m s2
+
+dessicate :: CtxMap CtxStruct t -> CtxMap (Const CtxHash) t
+dessicate = M.map (restructure (\(CtxStruct h1 _) -> Const h1))
+
+rehydrate :: forall t. CtxMap (Const CtxHash) t -> CtxMap CtxStruct t
+rehydrate m = m'
+ where
+  m' :: CtxMap CtxStruct t
+  m' = M.map (restructure (\(Const h) -> CtxStruct h (m' M.! h))) m

From 2383071c62ff778f7e70d98b465c36ff052305d5 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Mon, 4 Nov 2024 17:21:50 -0600
Subject: [PATCH 08/17] some additional comments

---
 src/swarm-lang/Swarm/Language/Context.hs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index be5b11bd3..35d584931 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -221,11 +221,21 @@ withBindings ctx = local (`union` ctx)
 ------------------------------------------------------------
 -- Context serializing/deserializing
 
+-- | A 'CtxMap' maps context hashes to context structures.  Those
+--   structures could either be complete context trees, or just a
+--   single level of structure containing more hashes.
 type CtxMap f t = Map CtxHash (CtxF f t)
 
+-- | Turn a context into a context map containing every subtree of its
+--   structure.
 toCtxMap :: Ctx t -> CtxMap CtxStruct t
 toCtxMap (Ctx m s) = run $ execState M.empty (buildCtxMap m s)
 
+-- | Build a context map by keeping track of the incrementally built
+--   map in a state effect, and traverse the given context structure
+--   to add all subtrees to the map---but, of course, stopping without
+--   recursing further whenever we see a hash that is already in the
+--   map.
 buildCtxMap :: forall t m sig. Has (State (CtxMap CtxStruct t)) sig m => Map Var t -> CtxStruct t -> m ()
 buildCtxMap m (CtxStruct h s) = do
   cm <- get @(CtxMap CtxStruct t)
@@ -239,9 +249,20 @@ buildCtxMap m (CtxStruct h s) = do
         CtxDelete x t s1 -> buildCtxMap (M.insert x t m) s1
         CtxUnion s1 s2 -> buildCtxMap m s1 *> buildCtxMap m s2
 
+-- | "Dessicate" a context map by replacing the actual context trees
+--   with single-layers containing only hashes.  A dessicated context
+--   map is very suitable for serializing, since it makes sharing
+--   completely explicit---even if a given context is referenced
+--   multiple times, the references are simply hash values, and the
+--   context is stored only once, under its hash.
 dessicate :: CtxMap CtxStruct t -> CtxMap (Const CtxHash) t
 dessicate = M.map (restructure (\(CtxStruct h1 _) -> Const h1))
 
+-- | "Rehydrate" a dessicated context map by replacing every hash with
+--   an actual context structure.  We do this by building the result
+--   as a lazy, recursive map, replacing each hash by the result we
+--   get when looking it up in the map being built.  A context which
+--   is referenced multiple times will thus be shared in memory.
 rehydrate :: forall t. CtxMap (Const CtxHash) t -> CtxMap CtxStruct t
 rehydrate m = m'
  where

From 64ea43f1ade8d9f35647a528c956a578cf67e6f4 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Mon, 4 Nov 2024 17:31:30 -0600
Subject: [PATCH 09/17] small improvements

---
 src/swarm-lang/Swarm/Language/Context.hs | 79 ++++++++++++++----------
 1 file changed, 48 insertions(+), 31 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index 35d584931..9bc15fae1 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -17,6 +17,7 @@ import Control.Lens.Empty (AsEmpty (..))
 import Control.Lens.Prism (prism)
 import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
 import Data.Data (Data)
+import Data.Function (on)
 import Data.Functor.Const
 import Data.Hashable
 import Data.List.NonEmpty qualified as NE
@@ -83,38 +84,44 @@ restructure _ (CtxSingle x t) = CtxSingle x t
 restructure h (CtxDelete x t f1) = CtxDelete x t (h f1)
 restructure h (CtxUnion f1 f2) = CtxUnion (h f1) (h f2)
 
--- | A "context structure" is one possible representation of a
---   context, consisting of a structured record of the process by
---   which a context was constructed.  This representation would be
---   terrible for doing efficient variable lookups, but it can be used
---   to efficiently destruct/serialize the context while recovering
+-- | A 'CtxTree' is one possible representation of a context,
+--   consisting of a structured record of the process by which a
+--   context was constructed.  This representation would be terrible
+--   for doing efficient variable lookups, but it can be used to
+--   efficiently serialize/deserialize the context while recovering
 --   sharing.
 --
 --   It stores a top-level hash of the context, along with a recursive
 --   tree built via 'CtxF'.
-data CtxStruct t = CtxStruct CtxHash (CtxF CtxStruct t)
+data CtxTree t = CtxTree CtxHash (CtxF CtxTree t)
   deriving (Eq, Functor, Foldable, Traversable, Data, Generic, ToJSON, FromJSON, Show)
 
 -- | A 'CtxNode' is just a single level of structure for a context,
 --   with any recursive contexts replaced by their hash.
+--
+--   For example, a 'CtxNode' could look something like @CtxUnion
+--   (Const 0fe5b299) (Const abcdef12)@.
 type CtxNode t = CtxF (Const CtxHash) t
 
 ------------------------------------------------------------
 -- Contexts
 
 -- | A context is a mapping from variable names to things.  We store
---   both a 'Map' (for efficient lookup) as well as a 'CtxStruct' for
+--   both a 'Map' (for efficient lookup) as well as a 'CtxTree' for
 --   sharing-aware serializing/deserializing of contexts.
-data Ctx t = Ctx {unCtx :: Map Var t, ctxStruct :: CtxStruct t}
-  deriving (Eq, Functor, Traversable, Data, Generic)
+data Ctx t = Ctx {unCtx :: Map Var t, ctxStruct :: CtxTree t}
+  deriving (Functor, Traversable, Data, Generic)
 
 -- | Get the top-level hash of a context.
 ctxHash :: Ctx t -> CtxHash
-ctxHash (Ctx _ (CtxStruct h _)) = h
+ctxHash (Ctx _ (CtxTree h _)) = h
 
 instance Show (Ctx t) where
   show _ = "<Ctx>"
 
+instance Eq (Ctx t) where
+  (==) = (==) `on` ctxHash
+
 instance Hashable t => Hashable (Ctx t) where
   hash = getCtxHash . ctxHash
   hashWithSalt s = hashWithSalt s . getCtxHash . ctxHash
@@ -122,6 +129,16 @@ instance Hashable t => Hashable (Ctx t) where
 instance Foldable Ctx where
   foldMap f = foldMap f . unCtx
 
+-- | Rebuild a complete 'Ctx' from a 'CtxTree'.
+ctxFromTree :: CtxTree t -> Ctx t
+ctxFromTree tree = Ctx (varMap tree) tree
+ where
+  varMap (CtxTree _ s) = case s of
+    CtxEmpty -> M.empty
+    CtxSingle x t -> M.singleton x t
+    CtxDelete x _ s1 -> M.delete x (varMap s1)
+    CtxUnion s1 s2 -> varMap s2 `M.union` varMap s1
+
 ------------------------------------------------------------
 -- Context instances
 
@@ -156,11 +173,11 @@ instance AsEmpty (Ctx t) where
 
 -- | The empty context.
 empty :: Ctx t
-empty = Ctx M.empty (CtxStruct mempty CtxEmpty)
+empty = Ctx M.empty (CtxTree mempty CtxEmpty)
 
 -- | A singleton context.
 singleton :: Hashable t => Var -> t -> Ctx t
-singleton x t = Ctx (M.singleton x t) (CtxStruct (singletonHash x t) (CtxSingle x t))
+singleton x t = Ctx (M.singleton x t) (CtxTree (singletonHash x t) (CtxSingle x t))
 
 -- | Create a 'Ctx' from a 'Map'.
 fromMap :: Hashable t => Map Var t -> Ctx t
@@ -178,9 +195,9 @@ lookupR x = lookup x <$> ask
 
 -- | Delete a variable from a context.
 delete :: Hashable t => Var -> Ctx t -> Ctx t
-delete x c@(Ctx m s@(CtxStruct h _)) = case M.lookup x m of
+delete x c@(Ctx m s@(CtxTree h _)) = case M.lookup x m of
   Nothing -> c
-  Just t -> Ctx (M.delete x m) (CtxStruct (h - singletonHash x t) (CtxDelete x t s))
+  Just t -> Ctx (M.delete x m) (CtxTree (h - singletonHash x t) (CtxDelete x t s))
 
 -- | Get the list of key-value associations from a context.
 assocs :: Ctx t -> [(Var, t)]
@@ -193,9 +210,9 @@ vars = M.keys . unCtx
 -- | Add a key-value binding to a context (overwriting the old one if
 --   the key is already present).
 addBinding :: Hashable t => Var -> t -> Ctx t -> Ctx t
-addBinding x t (Ctx m s@(CtxStruct h _)) = Ctx (M.insert x t m) s'
+addBinding x t (Ctx m s@(CtxTree h _)) = Ctx (M.insert x t m) s'
  where
-  s' = CtxStruct h' (CtxUnion (CtxStruct tHash (CtxSingle x t)) s)
+  s' = CtxTree h' (CtxUnion (CtxTree tHash (CtxSingle x t)) s)
   tHash = singletonHash x t
   h' = case M.lookup x m of
     Nothing -> h + tHash
@@ -203,7 +220,7 @@ addBinding x t (Ctx m s@(CtxStruct h _)) = Ctx (M.insert x t m) s'
 
 -- | /Right/-biased union of contexts.
 union :: Hashable t => Ctx t -> Ctx t -> Ctx t
-union (Ctx m1 s1@(CtxStruct h1 _)) (Ctx m2 s2@(CtxStruct h2 _)) = Ctx (m2 `M.union` m1) (CtxStruct h' (CtxUnion s1 s2))
+union (Ctx m1 s1@(CtxTree h1 _)) (Ctx m2 s2@(CtxTree h2 _)) = Ctx (m2 `M.union` m1) (CtxTree h' (CtxUnion s1 s2))
  where
   -- `Data.Map.intersection l r` returns a map with common keys, but values from `l`
   overwritten = M.intersection m1 m2
@@ -228,7 +245,7 @@ type CtxMap f t = Map CtxHash (CtxF f t)
 
 -- | Turn a context into a context map containing every subtree of its
 --   structure.
-toCtxMap :: Ctx t -> CtxMap CtxStruct t
+toCtxMap :: Ctx t -> CtxMap CtxTree t
 toCtxMap (Ctx m s) = run $ execState M.empty (buildCtxMap m s)
 
 -- | Build a context map by keeping track of the incrementally built
@@ -236,9 +253,9 @@ toCtxMap (Ctx m s) = run $ execState M.empty (buildCtxMap m s)
 --   to add all subtrees to the map---but, of course, stopping without
 --   recursing further whenever we see a hash that is already in the
 --   map.
-buildCtxMap :: forall t m sig. Has (State (CtxMap CtxStruct t)) sig m => Map Var t -> CtxStruct t -> m ()
-buildCtxMap m (CtxStruct h s) = do
-  cm <- get @(CtxMap CtxStruct t)
+buildCtxMap :: forall t m sig. Has (State (CtxMap CtxTree t)) sig m => Map Var t -> CtxTree t -> m ()
+buildCtxMap m (CtxTree h s) = do
+  cm <- get @(CtxMap CtxTree t)
   case h `M.member` cm of
     True -> pure ()
     False -> do
@@ -250,21 +267,21 @@ buildCtxMap m (CtxStruct h s) = do
         CtxUnion s1 s2 -> buildCtxMap m s1 *> buildCtxMap m s2
 
 -- | "Dessicate" a context map by replacing the actual context trees
---   with single-layers containing only hashes.  A dessicated context
---   map is very suitable for serializing, since it makes sharing
---   completely explicit---even if a given context is referenced
---   multiple times, the references are simply hash values, and the
---   context is stored only once, under its hash.
-dessicate :: CtxMap CtxStruct t -> CtxMap (Const CtxHash) t
-dessicate = M.map (restructure (\(CtxStruct h1 _) -> Const h1))
+--   with single structure layers containing only hashes.  A
+--   dessicated context map is very suitable for serializing, since it
+--   makes sharing completely explicit---even if a given context is
+--   referenced multiple times, the references are simply hash values,
+--   and the context is stored only once, under its hash.
+dessicate :: CtxMap CtxTree t -> CtxMap (Const CtxHash) t
+dessicate = M.map (restructure (\(CtxTree h1 _) -> Const h1))
 
 -- | "Rehydrate" a dessicated context map by replacing every hash with
 --   an actual context structure.  We do this by building the result
 --   as a lazy, recursive map, replacing each hash by the result we
 --   get when looking it up in the map being built.  A context which
 --   is referenced multiple times will thus be shared in memory.
-rehydrate :: forall t. CtxMap (Const CtxHash) t -> CtxMap CtxStruct t
+rehydrate :: forall t. CtxMap (Const CtxHash) t -> CtxMap CtxTree t
 rehydrate m = m'
  where
-  m' :: CtxMap CtxStruct t
-  m' = M.map (restructure (\(Const h) -> CtxStruct h (m' M.! h))) m
+  m' :: CtxMap CtxTree t
+  m' = M.map (restructure (\(Const h) -> CtxTree h (m' M.! h))) m

From 9be83f0c4be57ae40e33d6959c34f79af089f7c9 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Mon, 4 Nov 2024 17:54:13 -0600
Subject: [PATCH 10/17] ToJSON and FromJSONE instances for Ctx

---
 src/swarm-lang/Swarm/Language/Context.hs | 30 ++++++++++++++++++------
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index 9bc15fae1..a18c04e4e 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -27,7 +27,9 @@ import Data.Semigroup (Sum (..))
 import Data.Text (Text)
 import GHC.Generics (Generic)
 import Swarm.Pretty (PrettyPrec (..))
+import Swarm.Util (failT, showT)
 import Swarm.Util.JSON (optionsUnwrapUnary)
+import Swarm.Util.Yaml (FromJSONE, getE, liftE, parseJSONE)
 import Text.Printf (printf)
 import Prelude hiding (lookup)
 
@@ -142,13 +144,21 @@ ctxFromTree tree = Ctx (varMap tree) tree
 ------------------------------------------------------------
 -- Context instances
 
--- XXX this instance will have to change!!
-instance ToJSON t => ToJSON (Ctx t) where
-  toJSON = genericToJSON optionsUnwrapUnary
-
--- XXX this instance will have to change!!
-instance FromJSON t => FromJSON (Ctx t) where
-  parseJSON = genericParseJSON optionsUnwrapUnary
+-- | Serialize a context simply as its hash; we assume that a
+--   top-level CtxMap has been seralized somewhere, from which we can
+--   recover this context by looking it up.
+instance ToJSON (Ctx t) where
+  toJSON = toJSON . ctxHash
+
+-- | Deserialize a context.  We expect to see a hash, and look it up
+--   in the provided CtxMap.
+instance FromJSONE (CtxMap CtxTree t) (Ctx t) where
+  parseJSONE v = do
+    h <- liftE $ parseJSON @CtxHash v
+    m <- getE
+    case getCtx h m of
+      Nothing -> failT ["Encountered unknown context hash", showT h]
+      Just ctx -> pure ctx
 
 instance (PrettyPrec t) => PrettyPrec (Ctx t) where
   prettyPrec _ _ = "<Ctx>"
@@ -243,6 +253,12 @@ withBindings ctx = local (`union` ctx)
 --   single level of structure containing more hashes.
 type CtxMap f t = Map CtxHash (CtxF f t)
 
+-- | Read a context from a context map.
+getCtx :: CtxHash -> CtxMap CtxTree t -> Maybe (Ctx t)
+getCtx h m = case M.lookup h m of
+  Nothing -> Nothing
+  Just tree -> Just $ ctxFromTree (CtxTree h tree)
+
 -- | Turn a context into a context map containing every subtree of its
 --   structure.
 toCtxMap :: Ctx t -> CtxMap CtxTree t

From a110a32990fafeb370e65ca63897ecd5b56c9dfd Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Mon, 4 Nov 2024 18:07:36 -0600
Subject: [PATCH 11/17] add some unit tests

---
 src/swarm-lang/Swarm/Language/Context.hs | 10 ++--
 src/swarm-lang/Swarm/Language/JSON.hs    | 14 +++---
 swarm.cabal                              |  1 +
 test/unit/Main.hs                        |  2 +
 test/unit/TestContext.hs                 | 63 ++++++++++++++++++++++++
 5 files changed, 79 insertions(+), 11 deletions(-)
 create mode 100644 test/unit/TestContext.hs

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index a18c04e4e..85b33b0aa 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -282,16 +282,16 @@ buildCtxMap m (CtxTree h s) = do
         CtxDelete x t s1 -> buildCtxMap (M.insert x t m) s1
         CtxUnion s1 s2 -> buildCtxMap m s1 *> buildCtxMap m s2
 
--- | "Dessicate" a context map by replacing the actual context trees
+-- | "Dehydrate" a context map by replacing the actual context trees
 --   with single structure layers containing only hashes.  A
---   dessicated context map is very suitable for serializing, since it
+--   dehydrated context map is very suitable for serializing, since it
 --   makes sharing completely explicit---even if a given context is
 --   referenced multiple times, the references are simply hash values,
 --   and the context is stored only once, under its hash.
-dessicate :: CtxMap CtxTree t -> CtxMap (Const CtxHash) t
-dessicate = M.map (restructure (\(CtxTree h1 _) -> Const h1))
+dehydrate :: CtxMap CtxTree t -> CtxMap (Const CtxHash) t
+dehydrate = M.map (restructure (\(CtxTree h1 _) -> Const h1))
 
--- | "Rehydrate" a dessicated context map by replacing every hash with
+-- | "Rehydrate" a dehydrated context map by replacing every hash with
 --   an actual context structure.  We do this by building the result
 --   as a lazy, recursive map, replacing each hash by the result we
 --   get when looking it up in the map being built.  A context which
diff --git a/src/swarm-lang/Swarm/Language/JSON.hs b/src/swarm-lang/Swarm/Language/JSON.hs
index f906ab25d..0e8998449 100644
--- a/src/swarm-lang/Swarm/Language/JSON.hs
+++ b/src/swarm-lang/Swarm/Language/JSON.hs
@@ -8,14 +8,13 @@
 -- to put them all here to avoid circular module dependencies.
 module Swarm.Language.JSON where
 
-import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON, withText)
+import Data.Aeson (FromJSON (..), ToJSON (..), withText)
 import Data.Aeson qualified as Ae
 import Swarm.Language.Pipeline (processTermEither)
 import Swarm.Language.Syntax (Term)
 import Swarm.Language.Syntax.Pattern (Syntax, TSyntax)
 import Swarm.Language.Value (Env, Value)
 import Swarm.Pretty (prettyText)
-import Swarm.Util.JSON (optionsMinimize)
 import Witch (into)
 
 instance FromJSON TSyntax where
@@ -30,10 +29,13 @@ instance ToJSON Term
 instance ToJSON Syntax
 
 instance ToJSON Value where
-  toJSON = genericToJSON optionsMinimize
+  toJSON = undefined
 
 instance FromJSON Value where
-  parseJSON = genericParseJSON optionsMinimize
+  parseJSON = undefined
 
-deriving instance FromJSON Env
-deriving instance ToJSON Env
+instance ToJSON Env where
+  toJSON = undefined
+
+instance FromJSON Env where
+  parseJSON = undefined
diff --git a/swarm.cabal b/swarm.cabal
index b26427638..0920c4b96 100644
--- a/swarm.cabal
+++ b/swarm.cabal
@@ -1223,6 +1223,7 @@ test-suite swarm-unit
   other-modules:
     TestBoolExpr
     TestCommand
+    TestContext
     TestEval
     TestInventory
     TestLSP
diff --git a/test/unit/Main.hs b/test/unit/Main.hs
index a2566a171..dcfc6f6a1 100644
--- a/test/unit/Main.hs
+++ b/test/unit/Main.hs
@@ -28,6 +28,7 @@ import Test.Tasty.QuickCheck (
  )
 import TestBoolExpr (testBoolExpr)
 import TestCommand (testCommands)
+import TestContext (testContext)
 import TestEval (testEval)
 import TestInventory (testInventory)
 import TestLSP (testLSP)
@@ -82,6 +83,7 @@ statelessTests =
     , testPrettyConst
     , testBoolExpr
     , testCommands
+    , testContext
     , testHighScores
     , testRepl
     , testRequirements
diff --git a/test/unit/TestContext.hs b/test/unit/TestContext.hs
new file mode 100644
index 000000000..0d89be8c0
--- /dev/null
+++ b/test/unit/TestContext.hs
@@ -0,0 +1,63 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+-- |
+-- SPDX-License-Identifier: BSD-3-Clause
+--
+-- Swarm unit tests for contexts
+module TestContext where
+
+import Data.Map qualified as M
+import Swarm.Language.Context
+import Swarm.Util (showT)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, testCase)
+
+testContext :: TestTree
+testContext =
+  testGroup
+    "Contexts"
+    [ testGroup
+        "Context equality"
+        [ testCase "idempotence 1" $ ctxsEqual ctx1 (ctx1 <> ctx1)
+        , testCase "idempotence 2" $ ctxsEqual ctx2 (ctx2 <> ctx2)
+        , testCase "deletion" $ ctxsEqual ctx1 (delete "z" ctx2)
+        , testCase "empty/delete" $ ctxsEqual empty (delete "x" ctx1)
+        , testCase "fromMap" $ ctxsEqual ctx2 (fromMap (M.fromList [("x", 3), ("z", 6)]))
+        , testCase "right bias" $ ctxsEqual ctx4 (ctx2 <> ctx3)
+        , testCase "commutativity" $ ctxsEqual (ctx1 <> ctx5) (ctx5 <> ctx1)
+        ]
+    , testGroup
+        "de/rehydrate round-trip"
+        [ testCase "empty" $ serializeRoundTrip empty
+        , testCase "ctx1" $ serializeRoundTrip ctx1
+        , testCase "ctx2" $ serializeRoundTrip ctx2
+        , testCase "ctx3" $ serializeRoundTrip ctx3
+        , testCase "ctx4" $ serializeRoundTrip ctx4
+        , testCase "ctx5" $ serializeRoundTrip ctx5
+        , testCase "large" $ serializeRoundTrip bigCtx
+        , testCase "delete" $ serializeRoundTrip (delete "y" ctx4)
+        ]
+    ]
+ where
+  ctx1 = singleton "x" 3
+  ctx2 = singleton "x" 3 <> singleton "z" 6
+  ctx3 = singleton "x" 5 <> singleton "y" 7
+  ctx4 = singleton "x" 5 <> singleton "y" 7 <> singleton "z" 6
+  ctx5 = singleton "y" 10
+  bigCtx = fromMap . M.fromList $ zip (map (("x" <>) . showT) [1 :: Int ..]) [1 .. 10000]
+
+ctxsEqual :: Ctx Int -> Ctx Int -> Assertion
+ctxsEqual ctx1 ctx2 = do
+  -- Contexts are compared by hash for equality
+  assertEqual "hash equality" ctx1 ctx2
+
+  -- Make sure they are also structurally equal
+  assertBool "structural equality" (ctxStructEqual ctx1 ctx2)
+ where
+  ctxStructEqual (Ctx m1 _) (Ctx m2 _) = m1 == m2
+
+serializeRoundTrip :: Ctx Int -> Assertion
+serializeRoundTrip ctx = do
+  case getCtx (ctxHash ctx) (rehydrate (dehydrate (toCtxMap ctx))) of
+    Nothing -> fail "Failed to reconstitute dehydrated context"
+    Just ctx' -> ctxsEqual ctx ctx'

From e04fb1daedd88e41a1a8a79f643f3da5eb4226cb Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Mon, 4 Nov 2024 21:43:51 -0600
Subject: [PATCH 12/17] new `rollCtx` function collects up all hash
 computations in one place

---
 src/swarm-lang/Swarm/Language/Context.hs | 52 ++++++++++++++----------
 1 file changed, 31 insertions(+), 21 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index 85b33b0aa..34c51b398 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -15,7 +15,7 @@ import Control.Effect.Reader (Reader, ask, local)
 import Control.Effect.State (State, get, modify)
 import Control.Lens.Empty (AsEmpty (..))
 import Control.Lens.Prism (prism)
-import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON)
+import Data.Aeson (FromJSON (..), ToJSON (..))
 import Data.Data (Data)
 import Data.Function (on)
 import Data.Functor.Const
@@ -28,7 +28,6 @@ import Data.Text (Text)
 import GHC.Generics (Generic)
 import Swarm.Pretty (PrettyPrec (..))
 import Swarm.Util (failT, showT)
-import Swarm.Util.JSON (optionsUnwrapUnary)
 import Swarm.Util.Yaml (FromJSONE, getE, liftE, parseJSONE)
 import Text.Printf (printf)
 import Prelude hiding (lookup)
@@ -105,6 +104,26 @@ data CtxTree t = CtxTree CtxHash (CtxF CtxTree t)
 --   (Const 0fe5b299) (Const abcdef12)@.
 type CtxNode t = CtxF (Const CtxHash) t
 
+-- | Roll up one level of context structure while building a new
+--   top-level Map and computing an appropriate top-level hash.
+rollCtx :: Hashable t => CtxF Ctx t -> Ctx t
+rollCtx s = Ctx m (CtxTree h (restructure ctxStruct s))
+ where
+  (m, h) = case s of
+    CtxEmpty -> (M.empty, 0)
+    CtxSingle x t -> (M.singleton x t, singletonHash x t)
+    CtxDelete x _ (Ctx m1 (CtxTree h1 _)) -> case M.lookup x m1 of
+      Nothing -> (m1, h1)
+      Just t' -> (M.delete x m1, h1 - singletonHash x t')
+    CtxUnion (Ctx m1 (CtxTree h1 _)) (Ctx m2 (CtxTree h2 _)) -> (m2 `M.union` m1, h')
+     where
+      -- `Data.Map.intersection l r` returns a map with common keys,
+      -- but values from `l`.  The values in m1 are the ones we want
+      -- to subtract from the hash, since they are the ones that will
+      -- be overwritten.
+      overwritten = M.intersection m1 m2
+      h' = h1 + h2 - mapHash overwritten
+
 ------------------------------------------------------------
 -- Contexts
 
@@ -183,11 +202,13 @@ instance AsEmpty (Ctx t) where
 
 -- | The empty context.
 empty :: Ctx t
+-- We could also define empty = rollCtx CtxEmpty but that would introduce an
+-- unnecessary Hashable t constraint.
 empty = Ctx M.empty (CtxTree mempty CtxEmpty)
 
 -- | A singleton context.
 singleton :: Hashable t => Var -> t -> Ctx t
-singleton x t = Ctx (M.singleton x t) (CtxTree (singletonHash x t) (CtxSingle x t))
+singleton x t = rollCtx $ CtxSingle x t
 
 -- | Create a 'Ctx' from a 'Map'.
 fromMap :: Hashable t => Map Var t -> Ctx t
@@ -197,7 +218,7 @@ fromMap m = case NE.nonEmpty (M.assocs m) of
 
 -- | Look up a variable in a context.
 lookup :: Var -> Ctx t -> Maybe t
-lookup x (Ctx m _) = M.lookup x m
+lookup x = M.lookup x . unCtx
 
 -- | Look up a variable in a context in an ambient Reader effect.
 lookupR :: Has (Reader (Ctx t)) sig m => Var -> m (Maybe t)
@@ -205,9 +226,9 @@ lookupR x = lookup x <$> ask
 
 -- | Delete a variable from a context.
 delete :: Hashable t => Var -> Ctx t -> Ctx t
-delete x c@(Ctx m s@(CtxTree h _)) = case M.lookup x m of
-  Nothing -> c
-  Just t -> Ctx (M.delete x m) (CtxTree (h - singletonHash x t) (CtxDelete x t s))
+delete x ctx@(Ctx m _) = case M.lookup x m of
+  Nothing -> ctx
+  Just t -> rollCtx $ CtxDelete x t ctx
 
 -- | Get the list of key-value associations from a context.
 assocs :: Ctx t -> [(Var, t)]
@@ -220,21 +241,11 @@ vars = M.keys . unCtx
 -- | Add a key-value binding to a context (overwriting the old one if
 --   the key is already present).
 addBinding :: Hashable t => Var -> t -> Ctx t -> Ctx t
-addBinding x t (Ctx m s@(CtxTree h _)) = Ctx (M.insert x t m) s'
- where
-  s' = CtxTree h' (CtxUnion (CtxTree tHash (CtxSingle x t)) s)
-  tHash = singletonHash x t
-  h' = case M.lookup x m of
-    Nothing -> h + tHash
-    Just t' -> h - singletonHash x t' + tHash
+addBinding x t ctx = ctx `union` singleton x t
 
 -- | /Right/-biased union of contexts.
 union :: Hashable t => Ctx t -> Ctx t -> Ctx t
-union (Ctx m1 s1@(CtxTree h1 _)) (Ctx m2 s2@(CtxTree h2 _)) = Ctx (m2 `M.union` m1) (CtxTree h' (CtxUnion s1 s2))
- where
-  -- `Data.Map.intersection l r` returns a map with common keys, but values from `l`
-  overwritten = M.intersection m1 m2
-  h' = h1 + h2 - mapHash overwritten
+union ctx1 ctx2 = rollCtx $ CtxUnion ctx1 ctx2
 
 -- | Locally extend the context with an additional binding.
 withBinding :: (Has (Reader (Ctx t)) sig m, Hashable t) => Var -> t -> m a -> m a
@@ -296,8 +307,7 @@ dehydrate = M.map (restructure (\(CtxTree h1 _) -> Const h1))
 --   as a lazy, recursive map, replacing each hash by the result we
 --   get when looking it up in the map being built.  A context which
 --   is referenced multiple times will thus be shared in memory.
-rehydrate :: forall t. CtxMap (Const CtxHash) t -> CtxMap CtxTree t
+rehydrate :: CtxMap (Const CtxHash) t -> CtxMap CtxTree t
 rehydrate m = m'
  where
-  m' :: CtxMap CtxTree t
   m' = M.map (restructure (\(Const h) -> CtxTree h (m' M.! h))) m

From 0409bfc34cbf649cc343641baeaf38777b8ba676 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Tue, 5 Nov 2024 17:45:58 -0600
Subject: [PATCH 13/17] QuickCheck tests for context hash collisions

---
 swarm.cabal              |  4 ++++
 test/unit/TestContext.hs | 26 ++++++++++++++++++++++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/swarm.cabal b/swarm.cabal
index 0920c4b96..00d35094c 100644
--- a/swarm.cabal
+++ b/swarm.cabal
@@ -284,6 +284,9 @@ common parser-combinators
 common prettyprinter
   build-depends: prettyprinter >=1.7.0 && <1.8
 
+common quickcheck-instances
+  build-depends: quickcheck-instances >= 0.3.17 && < 0.4
+
 common random
   build-depends: random >=1.2.0 && <1.3
 
@@ -1210,6 +1213,7 @@ test-suite swarm-unit
     lens,
     megaparsec,
     mtl,
+    quickcheck-instances,
     tasty,
     tasty-hunit,
     tasty-quickcheck,
diff --git a/test/unit/TestContext.hs b/test/unit/TestContext.hs
index 0d89be8c0..fa4373100 100644
--- a/test/unit/TestContext.hs
+++ b/test/unit/TestContext.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE OverloadedStrings #-}
+{-# OPTIONS_GHC -fno-warn-orphans #-}
 
 -- |
 -- SPDX-License-Identifier: BSD-3-Clause
@@ -6,11 +7,16 @@
 -- Swarm unit tests for contexts
 module TestContext where
 
+import Control.Monad (replicateM)
+import Data.Hashable
+import Data.List (nub)
 import Data.Map qualified as M
 import Swarm.Language.Context
 import Swarm.Util (showT)
+import Test.QuickCheck.Instances.Text ()
 import Test.Tasty (TestTree, testGroup)
 import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, testCase)
+import Test.Tasty.QuickCheck (testProperty, Gen, generate, withMaxSuccess, Arbitrary(..))
 
 testContext :: TestTree
 testContext =
@@ -37,6 +43,15 @@ testContext =
         , testCase "large" $ serializeRoundTrip bigCtx
         , testCase "delete" $ serializeRoundTrip (delete "y" ctx4)
         ]
+    , testProperty
+        "no paired hash collisions"
+        (withMaxSuccess 10000 (hashConsistent @Int))
+    , testCase
+        "no hash collisions in a large pool" $
+        do
+          ctxs <- generate (replicateM 100000 (arbitrary :: Gen (Ctx Int)))
+          let m = M.fromListWith (++) (map (\ctx -> (ctxHash ctx, [unCtx ctx])) ctxs)
+          assertBool "foo" $ all ((==1) . length . nub) m
     ]
  where
   ctx1 = singleton "x" 3
@@ -46,6 +61,12 @@ testContext =
   ctx5 = singleton "y" 10
   bigCtx = fromMap . M.fromList $ zip (map (("x" <>) . showT) [1 :: Int ..]) [1 .. 10000]
 
+instance (Hashable a, Arbitrary a) => Arbitrary (Ctx a) where
+  arbitrary = fromMap <$> arbitrary
+
+hashConsistent :: Eq a => Ctx a -> Ctx a -> Bool
+hashConsistent ctx1 ctx2 = (ctx1 == ctx2) == (ctx1 `ctxStructEqual` ctx2)
+
 ctxsEqual :: Ctx Int -> Ctx Int -> Assertion
 ctxsEqual ctx1 ctx2 = do
   -- Contexts are compared by hash for equality
@@ -53,8 +74,9 @@ ctxsEqual ctx1 ctx2 = do
 
   -- Make sure they are also structurally equal
   assertBool "structural equality" (ctxStructEqual ctx1 ctx2)
- where
-  ctxStructEqual (Ctx m1 _) (Ctx m2 _) = m1 == m2
+
+ctxStructEqual :: Eq a => Ctx a -> Ctx a -> Bool
+ctxStructEqual (Ctx m1 _) (Ctx m2 _) = m1 == m2
 
 serializeRoundTrip :: Ctx Int -> Assertion
 serializeRoundTrip ctx = do

From d5e6ae59eca05db38947e0852797c0736ea37b4a Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Tue, 5 Nov 2024 17:50:51 -0600
Subject: [PATCH 14/17] disable warning for orphan Hashable Free instance

---
 src/swarm-lang/Swarm/Language/Types.hs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/swarm-lang/Swarm/Language/Types.hs b/src/swarm-lang/Swarm/Language/Types.hs
index 289b5e2c6..3ef4e05ae 100644
--- a/src/swarm-lang/Swarm/Language/Types.hs
+++ b/src/swarm-lang/Swarm/Language/Types.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# OPTIONS_GHC -fno-warn-orphans #-}
 
 -- |
 -- SPDX-License-Identifier: BSD-3-Clause
@@ -313,7 +314,7 @@ instance PrettyPrec IntVar where
 --   working with 'UType' as if it were defined directly.
 type UType = Free TypeF IntVar
 
--- XXX orphan instance
+-- orphan instance
 instance (Eq1 f, Hashable x, Hashable (f (Free f x))) => Hashable (Free f x)
 
 -- | A generic /fold/ for things defined via 'Free' (including, in

From a6412522b16109496738e90ff9a26d29c80c55c0 Mon Sep 17 00:00:00 2001
From: "github-actions[bot]"
 <41898282+github-actions[bot]@users.noreply.github.com>
Date: Tue, 5 Nov 2024 17:52:08 -0600
Subject: [PATCH 15/17] Restyled by fourmolu (#2203)

Co-authored-by: Restyled.io <commits@restyled.io>
---
 test/unit/TestContext.hs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/unit/TestContext.hs b/test/unit/TestContext.hs
index fa4373100..694035aaf 100644
--- a/test/unit/TestContext.hs
+++ b/test/unit/TestContext.hs
@@ -16,7 +16,7 @@ import Swarm.Util (showT)
 import Test.QuickCheck.Instances.Text ()
 import Test.Tasty (TestTree, testGroup)
 import Test.Tasty.HUnit (Assertion, assertBool, assertEqual, testCase)
-import Test.Tasty.QuickCheck (testProperty, Gen, generate, withMaxSuccess, Arbitrary(..))
+import Test.Tasty.QuickCheck (Arbitrary (..), Gen, generate, testProperty, withMaxSuccess)
 
 testContext :: TestTree
 testContext =
@@ -47,11 +47,11 @@ testContext =
         "no paired hash collisions"
         (withMaxSuccess 10000 (hashConsistent @Int))
     , testCase
-        "no hash collisions in a large pool" $
-        do
+        "no hash collisions in a large pool"
+        $ do
           ctxs <- generate (replicateM 100000 (arbitrary :: Gen (Ctx Int)))
           let m = M.fromListWith (++) (map (\ctx -> (ctxHash ctx, [unCtx ctx])) ctxs)
-          assertBool "foo" $ all ((==1) . length . nub) m
+          assertBool "foo" $ all ((== 1) . length . nub) m
     ]
  where
   ctx1 = singleton "x" 3

From 85dab031d91215d533b752c9cb4e9160c0d5e6bd Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Tue, 5 Nov 2024 17:53:30 -0600
Subject: [PATCH 16/17] normalize cabal formatting

---
 swarm.cabal | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/swarm.cabal b/swarm.cabal
index 00d35094c..d912d51aa 100644
--- a/swarm.cabal
+++ b/swarm.cabal
@@ -285,7 +285,7 @@ common prettyprinter
   build-depends: prettyprinter >=1.7.0 && <1.8
 
 common quickcheck-instances
-  build-depends: quickcheck-instances >= 0.3.17 && < 0.4
+  build-depends: quickcheck-instances >=0.3.17 && <0.4
 
 common random
   build-depends: random >=1.2.0 && <1.3

From 2a3b795c054e76d6853f0603c81cbafefb1ea559 Mon Sep 17 00:00:00 2001
From: Brent Yorgey <byorgey@gmail.com>
Date: Wed, 6 Nov 2024 07:30:17 -0600
Subject: [PATCH 17/17] To/FromJSON instances for CtxHash and CtxF

---
 src/swarm-lang/Swarm/Language/Context.hs | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/swarm-lang/Swarm/Language/Context.hs b/src/swarm-lang/Swarm/Language/Context.hs
index 34c51b398..a88aa2faf 100644
--- a/src/swarm-lang/Swarm/Language/Context.hs
+++ b/src/swarm-lang/Swarm/Language/Context.hs
@@ -15,7 +15,7 @@ import Control.Effect.Reader (Reader, ask, local)
 import Control.Effect.State (State, get, modify)
 import Control.Lens.Empty (AsEmpty (..))
 import Control.Lens.Prism (prism)
-import Data.Aeson (FromJSON (..), ToJSON (..))
+import Data.Aeson (FromJSON (..), FromJSONKey, ToJSON (..), ToJSONKey, genericParseJSON, genericToJSON, withText)
 import Data.Data (Data)
 import Data.Function (on)
 import Data.Functor.Const
@@ -25,11 +25,14 @@ import Data.Map (Map)
 import Data.Map qualified as M
 import Data.Semigroup (Sum (..))
 import Data.Text (Text)
+import Data.Text qualified as T
 import GHC.Generics (Generic)
 import Swarm.Pretty (PrettyPrec (..))
 import Swarm.Util (failT, showT)
+import Swarm.Util.JSON (optionsMinimize)
 import Swarm.Util.Yaml (FromJSONE, getE, liftE, parseJSONE)
 import Text.Printf (printf)
+import Text.Read (readMaybe)
 import Prelude hiding (lookup)
 
 -- | We use 'Text' values to represent variables.
@@ -51,13 +54,21 @@ type Var = Text
 --   functional, without the need to thread around some kind of
 --   globally unique ID generation effect.
 newtype CtxHash = CtxHash {getCtxHash :: Int}
-  deriving (Eq, Ord, Data, Generic, ToJSON, FromJSON)
+  deriving (Eq, Ord, Data, Generic, ToJSONKey, FromJSONKey)
   deriving (Semigroup, Monoid) via Sum Int
   deriving (Num) via Int
 
 instance Show CtxHash where
   show (CtxHash h) = printf "%016x" h
 
+instance ToJSON CtxHash where
+  toJSON h = toJSON (show h)
+
+instance FromJSON CtxHash where
+  parseJSON = withText "hash" $ \t -> case readMaybe ("0x" ++ T.unpack t) of
+    Nothing -> fail "Could not parse CtxHash"
+    Just h -> pure (CtxHash h)
+
 -- | The hash for a single variable -> value binding.
 singletonHash :: Hashable t => Var -> t -> CtxHash
 singletonHash x t = CtxHash $ hashWithSalt (hash x) t
@@ -76,7 +87,13 @@ data CtxF f t
   | CtxSingle Var t
   | CtxDelete Var t (f t)
   | CtxUnion (f t) (f t)
-  deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic, ToJSON, FromJSON)
+  deriving (Eq, Show, Functor, Foldable, Traversable, Data, Generic)
+
+instance (ToJSON t, ToJSON (f t)) => ToJSON (CtxF f t) where
+  toJSON = genericToJSON optionsMinimize
+
+instance (FromJSON t, FromJSON (f t)) => FromJSON (CtxF f t) where
+  parseJSON = genericParseJSON optionsMinimize
 
 -- | Map over the recursive structure stored in a 'CtxF'.
 restructure :: (f t -> g t) -> CtxF f t -> CtxF g t