Fix variable capture bug during context flattening

Includes proper test case result

Author: dolio

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

@@ -693,6 +693,13 @@ data CTE v s
 pattern ST1 :: Direction Word16 -> v -> Mem -> s -> CTE v s
 pattern ST1 d v m s = ST d [v] [m] s
 
+-- All variables, both bound and free occurring in a CTE. This is
+-- useful for avoiding both free and bound variables when
+-- freshening.
+cteVars :: Ord v => Cte v -> Set v
+cteVars (ST _ vs _ e) = Set.fromList vs `Set.union` ABTN.freeVars e
+cteVars (LZ v r as) = Set.fromList (either (const id) (:) r $ v:as)
+
 data ANormalF v e
   = ALet (Direction Word16) [Mem] e e
   | AName (Either Reference v) [v] e
@@ -1706,8 +1713,16 @@ tru = TCon Ty.booleanRef 1 []
 -- binding during ANF translation. Renames a variable in a
 -- context, and returns an indication of whether the varible
 -- was shadowed by one of the context bindings.
+--
+-- Note: this assumes that `u` is not bound by any of the context
+-- entries, as no effort is made to rename them to avoid capturing
+-- `u`.
 renameCtx :: (Var v) => v -> v -> Ctx v -> (Ctx v, Bool)
-renameCtx v u (d, ctx) | (ctx, b) <- rn [] ctx = ((d, ctx), b)
+renameCtx v u (d, ctx) | (ctx, b) <- renameCtes v u ctx = ((d, ctx), b)
+
+-- As above, but without the Direction.
+renameCtes :: Var v => v -> v -> [Cte v] -> ([Cte v], Bool)
+renameCtes v u = rn []
   where
     swap w
       | w == v = u
@@ -1725,7 +1740,92 @@ renameCtx v u (d, ctx) | (ctx, b) <- rn [] ctx = ((d, ctx), b)
       where
         e = LZ w (swap <$> f) (swap <$> as)
 
-anfBlock :: (Var v) => Term v a -> ANFM v (Ctx v, DNormal v)
+-- Simultaneously renames variables in a list of context entries.
+--
+-- Assumes that the variables being renamed to are not bound by the
+-- context entries, so that it is unnecessary to rename them.
+renamesCtes :: Var v => Map v v -> [Cte v] -> [Cte v]
+renamesCtes rn = map f
+  where
+    swap w
+      | Just u <- Map.lookup w rn = u
+      | otherwise = w
+
+    f (ST d vs ccs b) = ST d vs ccs (ABTN.renames rn b)
+    f (LZ v r as) = LZ v (second swap r) (map swap as)
+
+-- Calculates the free variables occurring in a context. This
+-- consists of the free variables in the expressions being bound,
+-- but with previously bound variables subtracted.
+freeVarsCtx :: Ord v => Ctx v -> Set v
+freeVarsCtx = freeVarsCte . snd
+
+freeVarsCte :: Ord v => [Cte v] -> Set v
+freeVarsCte = foldr m Set.empty
+  where
+    m (ST _ vs _ bn) rest =
+      ABTN.freeVars bn `Set.union` (rest Set.\\ Set.fromList vs)
+    m (LZ v r as) rest =
+      Set.fromList (either (const id) (:) r as)
+        `Set.union` Set.delete v rest
+
+-- Conditionally freshens a list of variables. The predicate
+-- argument selects which variables to freshen, and the set is a set
+-- of variables to avoid for freshness. The process ensures that the
+-- result is mutually fresh, and returns a new set of variables to
+-- avoid, which includes the freshened variables.
+--
+-- Presumably any variables selected by the predicate should be
+-- included in the set, but the set may contain additional variables
+-- to avoid, when freshening.
+freshens :: Var v => (v -> Bool) -> Set v -> [v] -> (Set v, [v])
+freshens p avoid0 vs =
+  mapAccumL f (Set.union avoid0 (Set.fromList vs)) vs
+  where
+    f avoid v
+      | p v, u <- Var.freshIn avoid v = (Set.insert u avoid, u)
+      | otherwise = (avoid, v)
+
+-- Freshens the variable bindings in a context to avoid a set of
+-- variables. Returns the renaming necessary for anything that was
+-- bound in the freshened context.
+--
+-- Note: this only freshens if it's necessary to avoid variables in
+-- the _original_ set. We need to keep track of other variables to
+-- avoid when making up new names for those, but it it isn't
+-- necessary to freshen variables to remove shadowing _within_ the
+-- context, since it is presumably already correctly formed.
+freshenCtx :: (Var v) => Set v -> Ctx v -> (Map v v, Ctx v)
+freshenCtx avoid0 (d, ctx) =
+  case go lavoid Map.empty [] $ reverse ctx of
+    (rn, ctx) -> (rn, (d, ctx))
+  where
+    -- precalculate all variable occurrences in the context to just
+    -- completely avoid those as well.
+    lavoid =
+      foldl (flip $ Set.union . cteVars) avoid0 ctx
+
+    go _     rns fresh [] = (rns, fresh)
+    go avoid rns fresh (bn : bns) = case bn of
+      LZ v r as
+        | v `Set.member` avoid0,
+          u <- Var.freshIn avoid v,
+          (fresh, _) <- renameCtes v u fresh,
+          avoid <- Set.insert u avoid,
+          rns <- Map.alter (Just . fromMaybe u) v rns ->
+          go avoid rns (LZ u r as : fresh) bns
+      ST d vs ccs expr
+        | (avoid, us) <- freshens (`Set.member` avoid0) avoid vs,
+          rn <- Map.fromList (filter (uncurry (/=)) $ zip vs us),
+          not (Map.null rn),
+          fresh <- renamesCtes rn fresh,
+          -- Note: rns union left-biased, so inner contexts take
+          -- priority.
+          rns <- Map.union rns rn ->
+          go avoid rns (ST d us ccs expr : fresh) bns
+      _ -> go avoid rns (bn : fresh) bns
+
+anfBlock :: (Ord v, Var v) => Term v a -> ANFM v (Ctx v, DNormal v)
 anfBlock (Var' v) = pure (mempty, pure $ TVar v)
 anfBlock (If' c t f) = do
   (cctx, cc) <- anfBlock c
@@ -1875,14 +1975,25 @@ anfBlock (Let1Named' v b e) =
   anfBlock b >>= \case
     (bctx, (Direct, TVar u)) -> do
       (ectx, ce) <- anfBlock e
+      (brn, bctx) <- fixupBctx bctx ectx ce
+      u <- pure $ Map.findWithDefault u u brn
       (ectx, shaded) <- pure $ renameCtx v u ectx
       ce <- pure $ if shaded then ce else ABTN.rename v u <$> ce
       pure (bctx <> ectx, ce)
     (bctx, (d0, cb)) -> bindLocal [v] $ do
       (ectx, ce) <- anfBlock e
       d <- bindDirection d0
+      (brn, bctx) <- fixupBctx bctx ectx ce
+      cb <- pure $ ABTN.renames brn cb
       let octx = bctx <> directed [ST1 d v BX cb] <> ectx
       pure (octx, ce)
+  where
+  fixupBctx bctx ectx (_, ce) =
+    pure $ freshenCtx (Set.union ecfvs efvs) bctx
+    where
+      ecfvs = freeVarsCtx ectx
+      efvs = ABTN.freeVars ce
+
 anfBlock (Apps' (Blank' b) args) = do
   nm <- fresh
   (actx, cas) <- anfArgs args

unison-src/transcripts/fix5419.output.md

@@ -0,0 +1,75 @@
+Below is an example of variable capture occuring from pattern matching.
+
+``` unison
+
+foo w = match (5, w) with
+  x ->
+    y = toText x
+    match 99 with _ -> ()
+    z = toText x
+    (y,z)
+
+> foo 8
+```
+
+``` ucm
+
+  Loading changes detected in scratch.u.
+
+  I found and typechecked these definitions in scratch.u. If you
+  do an `add` or `update`, here's how your codebase would
+  change:
+  
+    ⍟ These new definitions are ok to `add`:
+    
+      foo : w
+            -> ( Optional (Either Text Text),
+              Optional (Either Text Text))
+  
+  Now evaluating any watch expressions (lines starting with
+  `>`)... Ctrl+C cancels.
+
+    9 | > foo 8
+          ⧩
+          (Some (Right "(5, 8)"), Some (Right "(5, 8)"))
+
+```
+Arguably, the root cause is flattening of nested lets like this one.
+
+``` unison
+
+bar x =
+  -- argument here
+  y = Debug.toText x
+  let
+    x = 5
+    ()
+  -- 5 here, before fix
+  z = Debug.toText x
+  (y, z)
+
+> bar 3
+```
+
+``` ucm
+
+  Loading changes detected in scratch.u.
+
+  I found and typechecked these definitions in scratch.u. If you
+  do an `add` or `update`, here's how your codebase would
+  change:
+  
+    ⍟ These new definitions are ok to `add`:
+    
+      bar : x
+            -> ( Optional (Either Text Text),
+              Optional (Either Text Text))
+  
+  Now evaluating any watch expressions (lines starting with
+  `>`)... Ctrl+C cancels.
+
+    12 | > bar 3
+           ⧩
+           (Some (Right "3"), Some (Right "3"))
+
+```