Merge pull request #5757 from unisonweb/topic/affine-improvements

Fix some variable capture potential in the optimizations

Author: pchiusano

unison-core/src/Unison/ABT/Normalized.hs

@@ -15,8 +15,18 @@ module Unison.ABT.Normalized
     Align (..),
     alpha,
     freshen,
+    Renaming (..),
+    isEmptyRenaming,
+    freshenBinder,
+    freshenBinders,
+    pruneRenaming,
+    renameVar,
+    mapping,
+    avoiding,
+    mappingAndAvoiding,
     renames,
     renamesAvoiding,
+    renamesAndFreshen0,
     rename,
     transform,
     visit,
@@ -26,13 +36,13 @@ where
 
 import Data.Bifoldable
 import Data.Bifunctor
-import Data.Foldable (toList)
 import Data.Functor.Identity (Identity (..))
 import Data.Map.Strict (Map)
 import Data.Map.Strict qualified as Map
 import Data.Maybe (fromMaybe)
 import Data.Set (Set)
 import Data.Set qualified as Set
+import Data.Traversable (mapAccumL)
 import Unison.ABT (Var (..))
 
 -- ABTs with support for 'normalized' structure where only variables
@@ -107,9 +117,9 @@ class (Bifoldable f, Bifunctor f) => Align f where
 
 alphaErr ::
   (Align f) => (Var v) => Map v v -> Term f v -> Term f v -> Either (Term f v, Term f v) a
-alphaErr un tml tmr = Left (tml, renamesAndFreshen0 count un tmr)
+alphaErr un tml tmr = Left (tml, renamesAndFreshen0 rn tmr)
   where
-    count = Map.fromListWith (+) . flip zip [1, 1 ..] $ toList un
+    rn = mapping un
 
 -- Checks if two terms are equal up to a given variable renaming. The
 -- renaming should map variables in the right hand term to the
@@ -137,6 +147,93 @@ pattern TAbss vs bd <-
 
 {-# COMPLETE TAbss #-}
 
+-- Renaming data.
+--
+-- `renames` contains an actual mapping from old variables to new
+-- variables, which is used for the actual substitution.
+--
+-- `conflicts` stores information about which variables should be
+-- avoided. The value at a variable `u` should _at least_ be the
+-- number of variables in `renames` that are being substituted _to_
+-- `u`, because substituting under a binder for `u` will capture those
+-- substitutions. However, `conflicts` can be bootstrapped with extra
+-- counts to avoid ambient variables as well, so that it is possible
+-- to substitute/rewrite expressions without introducing variable
+-- captures.
+data Renaming v = RN
+  { conflicts :: Map v Int,
+    renamings :: Map v v
+  }
+
+-- Creates a Renaming that will avoid the given set of variables.
+avoiding :: Set v -> Renaming v
+avoiding avoid = RN (Map.fromSet (const 1) avoid) Map.empty
+
+mapping :: (Var v) => Map v v -> Renaming v
+mapping rn = RN cf rn
+  where
+    cf = Map.fromListWith (+) . fmap (,1) $ Map.elems rn
+
+mappingAndAvoiding :: (Var v) => Map v v -> Set v -> Renaming v
+mappingAndAvoiding rn avoid = RN cf rn
+  where
+    cf =
+      Map.unionWith
+        (+)
+        (Map.fromSet (const 1) avoid)
+        (Map.fromListWith (+) . fmap (,1) $ Map.elems rn)
+
+-- Adjusts a renaming with respect to a remaining set of free
+-- variables. Unnecessary renamings are discarded.
+pruneRenaming :: (Var v) => Set v -> Renaming v -> Renaming v
+pruneRenaming fvs (RN cf rn) =
+  RN
+    { renamings = Map.restrictKeys rn fvs,
+      conflicts = Map.foldl' decrement cf $ Map.withoutKeys rn fvs
+    }
+  where
+    decrement sv v = Map.update drop v sv
+    drop n
+      | n <= 1 = Nothing
+      | otherwise = Just (n - 1)
+
+renameVar :: (Var v) => Renaming v -> v -> v
+renameVar (RN _ rn) u = Map.findWithDefault u u rn
+
+-- Tests if the renaming is empty in the sense that it will never
+-- cause bound variables to be renamed. This is _not_ just a test of
+-- whether the substitutions are empty, because the conflicts can
+-- cause variables to need freshening even without variable
+-- substitutions.
+isEmptyRenaming :: Renaming v -> Bool
+isEmptyRenaming = null . conflicts
+
+-- Freshens a bound variable with regard to a renaming, yielding the
+-- fresh variable and a renaming appropriate for the term within the
+-- binder. The `Set` should be the free variables of the expression
+-- within the binder, for proper freshening.
+freshenBinder :: (Var v) => Set v -> Renaming v -> v -> (Renaming v, v)
+freshenBinder fvs rn0@(RN cf rn) u = (rn', u')
+  where
+    -- if u conflicts with the renaming, freshen it
+    u'
+      | u `Map.member` cf = freshIn (fvs `Set.union` Map.keysSet cf) u
+      | otherwise = u
+
+    -- if u needs to be renamed, and it actually occurs in the body,
+    -- add it to the Renaming.
+    rn'
+      | u /= u' && u `Set.member` fvs =
+          RN
+            { conflicts = Map.insertWith (+) u' 1 cf,
+              renamings = Map.alter (const $ Just u') u rn
+            }
+      | otherwise = rn0
+
+freshenBinders ::
+  (Var v) => Set v -> Renaming v -> [v] -> (Renaming v, [v])
+freshenBinders fvs = mapAccumL (freshenBinder fvs)
+
 -- Simultaneous variable renaming and freshening implementation.
 --
 -- subvs0 is a count of the number of conflicts associated with a
@@ -153,48 +250,22 @@ pattern TAbss vs bd <-
 -- rnv0 is the variable renaming map.
 renamesAndFreshen0 ::
   (Var v, Bifunctor f, Bifoldable f) =>
-  Map v Int ->
-  Map v v ->
+  Renaming v ->
   Term f v ->
   Term f v
-renamesAndFreshen0 subvs0 rnv0 tm = case tm of
+renamesAndFreshen0 rn0 tm = case tm of
   TAbs u body
-    | not $ Map.null subvs' ->
-        TAbs u' (renamesAndFreshen0 subvs' rnv' body)
-    where
-      rnv' = Map.alter (const $ adjustment) u rnv
-      bfvs = freeVars body
-      -- if u is in the set of variables we're substituting in, it
-      -- needs to be renamed to avoid capturing things.
-      u'
-        | u `Map.member` subvs = freshIn (bfvs `Set.union` Map.keysSet subvs) u
-        | otherwise = u
-
-      -- if u needs to be renamed to avoid capturing subvs
-      -- and u actually occurs in the body, then add it to
-      -- the substitutions
-      (adjustment, subvs')
-        | u /= u' && u `Set.member` bfvs =
-            (Just u', Map.insertWith (+) u' 1 subvs)
-        | otherwise = (Nothing, subvs)
+    | (rn, u') <- freshenBinder (freeVars body) rn u,
+      u /= u' || not (isEmptyRenaming rn) ->
+        TAbs u' (renamesAndFreshen0 rn body)
   TTm body
-    | not $ Map.null subvs ->
-        TTm $ bimap lkup (renamesAndFreshen0 subvs rnv) body
+    | not $ isEmptyRenaming rn ->
+        TTm $ bimap (renameVar rn) (renamesAndFreshen0 rn) body
   _ -> tm
   where
     fvs = freeVars tm
 
-    -- throw out irrelevant renamings
-    rnv = Map.restrictKeys rnv0 fvs
-
-    lkup u = Map.findWithDefault u u rnv
-
-    -- decrement the variable usage counts for the renamings we threw away
-    subvs = Map.foldl' decrement subvs0 $ Map.withoutKeys rnv0 fvs
-    decrement sv v = Map.update drop v sv
-    drop n
-      | n <= 1 = Nothing
-      | otherwise = Just (n - 1)
+    rn = pruneRenaming fvs rn0
 
 -- Freshens the bound variables in a term to avoid capturing variables
 -- in the set.
@@ -203,9 +274,7 @@ freshen ::
   Set v ->
   Term f v ->
   Term f v
-freshen avoid = renamesAndFreshen0 subvs Map.empty
-  where
-    subvs = Map.fromSet (const 1) avoid
+freshen avoid = renamesAndFreshen0 (avoiding avoid)
 
 -- Renames some variables while also avoiding a given set of variables
 -- for any bindings in the term.
@@ -215,30 +284,24 @@ renamesAvoiding ::
   Map v v ->
   Term f v ->
   Term f v
-renamesAvoiding avoid rnv = renamesAndFreshen0 subvs rnv
-  where
-    suba = Map.fromSet (const 1) avoid
-    subr = Map.fromListWith (+) . fmap (,1) $ Map.elems rnv
-
-    subvs = Map.unionWith (+) suba subr
+renamesAvoiding avoid rnv =
+  renamesAndFreshen0 (mappingAndAvoiding rnv avoid)
 
 -- Simultaneous variable renaming.
 renames ::
   (Var v, Bifunctor f, Bifoldable f) =>
   Map v v ->
   Term f v ->
   Term f v
-renames rnv tm = renamesAndFreshen0 subvs rnv tm
-  where
-    subvs = Map.fromListWith (+) . fmap (,1) $ Map.elems rnv
+renames rnv tm = renamesAndFreshen0 (mapping rnv) tm
 
 rename ::
   (Var v, Ord v, Bifunctor f, Bifoldable f) =>
   v ->
   v ->
   Term f v ->
   Term f v
-rename old new = renamesAndFreshen0 (Map.singleton new 1) (Map.singleton old new)
+rename old new = renamesAndFreshen0 (mapping $ Map.singleton old new)
 
 transform ::
   (Var v, Bifunctor g, Bifoldable f, Bifoldable g) =>

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

@@ -489,47 +489,57 @@ pattern HandlerResume lz f as lh h bs rs <-
 matchHandledThunk ::
   (Var v) => ANormal v -> Maybe (Reference, Int, Bool, ANormal v)
 matchHandledThunk (TLet _ th _ (TCom r vs) bd) =
-  final <$> runWriterT (prefix bd)
+  final <$> runWriterT (prefix (ABTN.avoiding . Set.fromList $ th : vs) bd)
   where
     none = WriterT Nothing
 
     final (bd, All lazy) = (r, length vs, lazy, bd)
 
-    prefix (TName v g bs bd)
+    prefix rn (TName v g bs bd)
       | v /= th,
         all (/= th) g,
-        all (/= th) bs =
-          TName v g bs <$> prefix bd
-    prefix (TLets d vs ccs bn bd)
+        all (/= th) bs,
+        g <- ABTN.renameVar rn <$> g,
+        bs <- ABTN.renameVar rn <$> bs,
+        (rn, v) <- ABTN.freshenBinder (ABTN.freeVars bd) rn v =
+          TName v g bs <$> prefix rn bd
+    prefix rn (TLets d vs ccs bn bd)
       | all (/= th) vs,
-        th `Set.notMember` ABTN.freeVars bn =
-          TLets d vs ccs bn <$> prefix bd <* tell (All False)
-    prefix (TMatch sc bs) =
-      TMatch sc <$> traverse under bs
+        th `Set.notMember` ABTN.freeVars bn,
+        bn <- ABTN.renamesAndFreshen0 rn bn,
+        (rn, vs) <- ABTN.freshenBinders (ABTN.freeVars bd) rn vs =
+          TLets d vs ccs bn <$> prefix rn bd <* tell (All False)
+    prefix rn (TMatch sc bs) =
+      TMatch (ABTN.renameVar rn sc) <$> traverse under bs
       where
         under (ABTN.TAbss vs bd)
-          | all (/= th) vs =
-              ABTN.TAbss vs <$> prefix bd
+          | all (/= th) vs,
+            (rn, vs) <- ABTN.freshenBinders (ABTN.freeVars bd) rn vs =
+              ABTN.TAbss vs <$> prefix rn bd
         under _ = none
-    prefix (THnd rs nh ah bd)
+    prefix rn (THnd rs nh ah bd)
       | nh /= th,
-        all (/= th) ah =
-          THnd rs nh ah <$> suffix bd
-    prefix _ = none
+        all (/= th) ah,
+        ah <- ABTN.renameVar rn <$> ah,
+        nh <- ABTN.renameVar rn nh =
+          THnd rs nh ah <$> suffix rn bd
+    prefix _ _ = none
 
     -- Some values may be bound before the thunk call as long as
     -- they're 'direct' calls that can't capture stacks and reveal
     -- that we've changed the convention.
-    suffix (TLets d vs ccs bn bd)
+    suffix rn (TLets d vs ccs bn bd)
       | all (/= th) vs,
-        th `Set.notMember` ABTN.freeVars bn =
-          TLets d vs ccs bn <$> suffix bd <* tell (All $ d == Direct)
+        th `Set.notMember` ABTN.freeVars bn,
+        bn <- ABTN.renamesAndFreshen0 rn bn,
+        (rn, vs) <- ABTN.freshenBinders (ABTN.freeVars bd) rn vs =
+          TLets d vs ccs bn <$> suffix rn bd <* tell (All $ d == Direct)
     -- final expression in handle body is a call to the thunk.
-    suffix (TApv h us)
+    suffix rn (TApv h us)
       | h == th,
         all (/= th) us =
-          pure . TCom r $ vs ++ us
-    suffix _ = none
+          pure . TCom r $ vs ++ fmap (ABTN.renameVar rn) us
+    suffix _ _ = none
 matchHandledThunk _ = Nothing
 
 --  th = f <vs> -- undersaturated
@@ -748,9 +758,10 @@ affineHandlerCase opts self vs rec br
   | ABTN.TAbss us body <- br,
     TShift _ kf0 body <- body,
     TName kf (Left (Builtin "jumpCont")) [kf1] body <- body,
+    bound <- Set.fromList (kf0 : kf : us),
     kf0 == kf1 =
       ABTN.TAbss us
-        <$> affinePreBranch opts self Set.empty vs rec ar kf body
+        <$> affinePreBranch opts self bound vs rec ar kf body
   | otherwise = Nothing
   where
     ar = freshAff 2
@@ -850,7 +861,8 @@ linearTail opts self vs bound rec ar kf0 tm
     rh /= kf0, -- no shadowing or non-linearity
     THnd _rs hh Nothing bd <- tm,
     rh == hh, -- handle recursively
-    bd <- replaceLinearBody opts bd,
+    avoid <- Set.insert rh bound,
+    bd <- replaceLinearBody opts avoid bd,
     SimpleBody pre ind shad free kf1 result <- bd, -- simple enough body
     kf0 `Set.notMember` shad, -- kf is not shadowed in body
     kf0 `Set.notMember` free, -- kf is not free in `pre`
@@ -888,17 +900,18 @@ linearTail opts self vs bound rec ar kf0 tm
 -- continuations aren't captured, so we don't actually need a correct
 -- numbering. If this is ever changed, then the numbering here must be
 -- adjusted.
-replaceLinearBody :: (Var v) => OptInfos v -> ANormal v -> ANormal v
-replaceLinearBody opts@(arities, inls) bd
+replaceLinearBody ::
+  (Var v) => OptInfos v -> Set v -> ANormal v -> ANormal v
+replaceLinearBody opts@(arities, inls) avoid bd
   | TLetD v cc bn bd <- bd =
-      TLetD v cc bn $ replaceLinearBody opts bd
+      TLetD v cc bn $ replaceLinearBody opts (Set.insert v avoid) bd
   | TCom r vs <- bd,
     Just n <- Map.lookup r arities,
     length vs == n,
     Just (InlInfo _ (ABTN.TAbss us expr)) <- Map.lookup r inls,
     rn <- Map.fromList (zip us vs) =
-      ABTN.renames rn expr
-replaceLinearBody _ bd = bd
+      ABTN.renamesAvoiding avoid rn expr
+replaceLinearBody _ _ bd = bd
 
 parseSimpleHandlerBody ::
   (Var v) =>