Merge pull request #5778 from unisonweb/fix/opt-captures

Fix some variable naming problems with new optimizations

Author: aryairani

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

@@ -230,9 +230,16 @@ freshenBinder fvs rn0@(RN cf rn) u = (rn', u')
             }
       | otherwise = rn0
 
+-- Simultaneously freshens some binders. This ensures not just that
+-- they're fresh with respect to the given set of variables, but
+-- mutually distinct.
 freshenBinders ::
   (Var v) => Set v -> Renaming v -> [v] -> (Renaming v, [v])
-freshenBinders fvs = mapAccumL (freshenBinder fvs)
+freshenBinders fvs rn0 = first snd . mapAccumL f (Set.empty, rn0)
+  where
+    f (avoid, rn) u
+      | (rn, v) <- freshenBinder (Set.union avoid fvs) rn u =
+          ((Set.insert v avoid, rn), v)
 
 -- Simultaneous variable renaming and freshening implementation.
 --
@@ -254,10 +261,10 @@ renamesAndFreshen0 ::
   Term f v ->
   Term f v
 renamesAndFreshen0 rn0 tm = case tm of
-  TAbs u body
-    | (rn, u') <- freshenBinder (freeVars body) rn u,
-      u /= u' || not (isEmptyRenaming rn) ->
-        TAbs u' (renamesAndFreshen0 rn body)
+  TAbs u (TAbss us body)
+    | (rn, vs) <- freshenBinders (freeVars body) rn (u : us),
+      u : us /= vs || not (isEmptyRenaming rn) ->
+        TAbss vs (renamesAndFreshen0 rn body)
   TTm body
     | not $ isEmptyRenaming rn ->
         TTm $ bimap (renameVar rn) (renamesAndFreshen0 rn) body

unison-runtime/package.yaml

@@ -17,6 +17,12 @@ flags:
     manual: true
     default: false
 
+  # Run code through the serializer to test code loading as if it were
+  # remote.
+  codeserialchecks:
+    manual: true
+    default: false
+
   # Dumps core for debugging to unison-runtime/.stack-work/dist/<arch>/ghc-x.y.z/build/
   dumpcore:
     manual: true
@@ -32,6 +38,8 @@ when:
     cpp-options: -DOPT_CHECK
     dependencies:
       - inspection-testing
+  - condition: flag(codeserialchecks)
+    ghc-options: -DCODE_SERIAL_CHECK
   - condition: flag(dumpcore)
     ghc-options: -ddump-simpl -ddump-stg-final -ddump-to-file -dsuppress-coercions -dsuppress-idinfo -dsuppress-module-prefixes -ddump-str-signatures -ddump-simpl-stats # -dsuppress-type-applications -dsuppress-type-signatures
 

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

@@ -2538,10 +2538,10 @@ prettyBranches ind bs = case bs of
   MatchText bs df ->
     maybe id (\e -> prettyCase ind (showString "_") e id) df
       . foldr (uncurry $ prettyCase ind . shows) id (Map.toList bs)
-  MatchData _ bs df ->
+  MatchData r bs df ->
     maybe id (\e -> prettyCase ind (showString "_") e id) df
       . foldr
-        (uncurry $ prettyCase ind . shows)
+        (uncurry $ prettyCase ind . prettyTag r)
         id
         (mapToList $ snd <$> bs)
   MatchRequest bs df ->
@@ -2572,6 +2572,13 @@ prettyBranches ind bs = case bs of
         . shows c
         . showString ")"
 
+    prettyTag r c =
+      showString "CON("
+        . showsShort r
+        . showString ","
+        . shows c
+        . showString ")"
+
 prettyCase :: (Var v) => Int -> ShowS -> ANormal v -> ShowS -> ShowS
 prettyCase ind sc (ABTN.TAbss vs e) r =
   showString "\n"

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

@@ -165,46 +165,52 @@ whenChanged f act = do
 
 descend ::
   (Memo m, Var v) =>
-  (Bool -> ANormal v -> m (ANormal v)) ->
+  (Bool -> Set v -> ANormal v -> m (ANormal v)) ->
   Bool ->
+  Set v ->
   ANormal v ->
   m (ANormal v)
-descend rec tail tm = memo tm $ case tm of
+descend rec tail bound tm = memo tm $ case tm of
   TLets d vs ccs bn bd ->
-    TLets d vs ccs <$> rec False bn <*> rec tail bd
+    TLets d vs ccs <$> rec False bound bn <*> rec tail bnd bd
+    where
+      bnd = Set.union (Set.fromList vs) bound
   TName v f vs bd ->
-    TName v f vs <$> rec tail bd
+    TName v f vs <$> rec tail (Set.insert v bound) bd
   TMatch v bs ->
-    TMatch v <$> traverse (rec tail) bs
+    TMatch v <$> traverse (rec tail bound) bs
   TShift r v bd ->
-    TShift r v <$> rec tail bd
+    TShift r v <$> rec tail (Set.insert v bound) bd
   THnd rs hn ha bd ->
-    THnd rs hn ha <$> rec tail bd
+    THnd rs hn ha <$> rec tail bound bd
   TLocal v bd ->
-    TLocal v <$> rec tail bd
+    TLocal v <$> rec tail bound bd
   ABTN.TAbs v (ABTN.TAbss vs bd) ->
-    ABTN.TAbss (v : vs) <$> rec tail bd
+    ABTN.TAbss (v : vs) <$> rec tail bnd bd
+    where
+      bnd = Set.union (Set.fromList $ v : vs) bound
   _ -> pure tm
 
 -- Rewrites a term from the top down, first applying the step
 -- transform given, then descending to children.
 rewriteDown ::
   (Memo m, Var v) =>
-  (Bool -> ANormal v -> m (ANormal v)) ->
+  (Bool -> Set v -> ANormal v -> m (ANormal v)) ->
   ANormal v ->
   m (ANormal v)
-rewriteDown step = go True
+rewriteDown step = go True Set.empty
   where
-    go tail tm = step tail tm >>= descend go tail
+    go tail bound tm = step tail bound tm >>= descend go tail bound
 
 rewriteUp ::
   (Memo m, Var v) =>
-  (Bool -> ANormal v -> m (ANormal v)) ->
+  (Bool -> Set v -> ANormal v -> m (ANormal v)) ->
   ANormal v ->
   m (ANormal v)
-rewriteUp step = go True
+rewriteUp step = go True Set.empty
   where
-    go tail tm = memo tm (descend go tail tm) >>= step tail
+    go tail bound tm =
+      memo tm (descend go tail bound tm) >>= step tail bound
 
 -- Performs inlining on a `SuperGroup` using the inlining information
 -- in the map. The map can be created from typical `SuperGroup` data
@@ -225,15 +231,15 @@ inline avoid (arities, inls) n0 = memo n0 $ go (30 :: Int) n0
       | n <= 0 = pure tm
       | otherwise = rewriteUp (step n) tm
 
-    step n tail (TApp (FComb r) args)
-      | Just new <- findInline tail r args =
+    step n tail bound (TApp (FComb r) args)
+      | Just new <- findInline tail bound r args =
           dirty *> go (n - 1) new
-    step _ _tail tm = pure tm
+    step _ _tail _bound tm = pure tm
 
-    findInline tail r args = do
+    findInline tail bound r args = do
       info <- Map.lookup r inls
       arity <- Map.lookup r arities
-      tweak tail args arity info
+      tweak tail bound args arity info
 
     don'tInline Don'tInl _ = True
     don'tInline TailInl isTail = not isTail
@@ -245,17 +251,17 @@ inline avoid (arities, inls) n0 = memo n0 $ go (30 :: Int) n0
     -- multiple inlining steps, so we freshen anything else we inline
     -- to not be capable of capturing the variables from the entry
     -- code.
-    tweak isTail args arity (InlInfo clazz (ABTN.TAbss vs body))
+    tweak isTail bound args arity (InlInfo clazz (ABTN.TAbss vs body))
       | don'tInline clazz isTail = Nothing
       -- exactly saturated
       | length args == arity,
         rn <- Map.fromList (zip vs args) =
-          Just $ ABTN.renamesAvoiding avoid rn body
+          Just $ ABTN.renamesAvoiding (avoid `Set.union` bound) rn body
       -- oversaturated, only makes sense if body is a call
       | length args > arity,
         (pre, post) <- splitAt arity args,
         rn <- Map.fromList (zip vs pre),
-        TApp f pre <- ABTN.renamesAvoiding avoid rn body =
+        TApp f pre <- ABTN.renamesAvoiding (avoid `Set.union` bound) rn body =
           Just $ TApp f (pre ++ post)
       | otherwise = Nothing
 
@@ -276,7 +282,7 @@ peephole arities affine n0 = memo n0 $ go (30 :: Int) n0
   where
     go 0 = pure
     go n =
-      whenChanged (go $ n - 1) . rewriteDown \tail -> \case
+      whenChanged (go $ n - 1) . rewriteDown \tail _bound -> \case
         -- eliminate `v = u` bindings in affine contexts
         TLet _ v _ (TVar u) bd
           | affine -> ABTN.rename v u bd <$ dirty
@@ -338,9 +344,11 @@ optSuper ::
   Bool ->
   SuperNormal v ->
   m (SuperNormal v)
-optSuper opts avoid affine sn@(Lambda ccs (ABTN.TAbss vs bd)) =
+optSuper opts avoid0 affine sn@(Lambda ccs (ABTN.TAbss vs bd)) =
   memo sn $
     Lambda ccs . ABTN.TAbss vs <$> optNormal opts avoid affine bd
+  where
+    avoid = Set.union (Set.fromList vs) avoid0
 
 -- Optimizes a single group
 optGroup ::
@@ -716,6 +724,7 @@ translateHandlerMatch ::
   (Var v) => OptInfos v -> v -> v -> SuperNormal v -> Maybe (SuperNormal v)
 translateHandlerMatch opts self ah (Lambda ccs (ABTN.TAbss args body))
   | v : vs <- shiftArgs args,
+    bound <- Set.fromList (self : args),
     TMatch u branches <- body,
     u == v,
     MatchRequest cs df <- branches,
@@ -725,7 +734,7 @@ translateHandlerMatch opts self ah (Lambda ccs (ABTN.TAbss args body))
         . ABTN.TAbss args
         . TMatch u
         . flip MatchRequest df
-        <$> traverse3 (affineHandlerCase opts self vs ah) cs
+        <$> traverse3 (affineHandlerCase opts self bound vs ah) cs
   | otherwise = Nothing
   where
     ar = freshAff 2
@@ -753,12 +762,12 @@ augmentHandlerEntry thunk0 mv0 ah body
 -- Recognizes an affine handler case, yielding a translated efficient
 -- version if it is one.
 affineHandlerCase ::
-  (Var v) => OptInfos v -> v -> [v] -> v -> ANormal v -> Maybe (ANormal v)
-affineHandlerCase opts self vs rec br
+  (Var v) => OptInfos v -> v -> Set v -> [v] -> v -> ANormal v -> Maybe (ANormal v)
+affineHandlerCase opts self bound 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),
+    bound <- Set.union bound (Set.fromList (kf0 : kf : us)),
     kf0 == kf1 =
       ABTN.TAbss us
         <$> affinePreBranch opts self bound vs rec ar kf body

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

@@ -65,6 +65,9 @@ import Unison.Runtime.ANF as ANF
   )
 import Unison.Runtime.ANF qualified as ANF
 import Unison.Runtime.ANF.Optimize qualified as ANF
+#ifdef CODE_SERIAL_CHECK
+import Unison.Runtime.ANF.Serialize (serializeCode, deserializeCode)
+#endif
 import Unison.Runtime.Array as PA
 import Unison.Runtime.Builtin hiding (unitValue)
 import Unison.Runtime.Exception hiding (die)
@@ -1233,13 +1236,35 @@ addRefs vfrsh vfrom vto rs = do
 evaluateSTM :: a -> STM a
 evaluateSTM x = unsafeIOToSTM (evaluate x)
 
+-- If this flag is set, all code is run through serialization before
+-- loading. This renames variables, and it's possible a problem would
+-- only be visible with the renamed variables. This allows testing
+-- these cases just by rebuilding ucm, rather than actually concocting
+-- a test that involves remote code loading.
+#if defined(CODE_SERIAL_CHECK)
+
+normalizeCode :: Code -> Code
+normalizeCode co = case deserializeCode (serializeCode False co) of
+  Left _ -> error "normalizeCode: impossible"
+  Right co -> co
+
+normalizeCodes :: [(Reference, Code)] -> [(Reference, Code)]
+normalizeCodes = fmap $ second normalizeCode
+
+#else
+
+normalizeCodes :: [(Reference, Code)] -> [(Reference, Code)]
+normalizeCodes = id
+
+#endif
+
 cacheAdd0 ::
   S.Set Reference ->
   [(Reference, Code)] ->
   [(Reference, Set Reference)] ->
   CCache ->
   IO ()
-cacheAdd0 ntys0 termSuperGroups sands cc = do
+cacheAdd0 ntys0 (normalizeCodes -> termSuperGroups) sands cc = do
   let toAdd = M.fromList (termSuperGroups <&> second codeGroup)
   (unresolvedCacheableCombs, unresolvedNonCacheableCombs) <- atomically $ do
     have <- readTVar (intermed cc)

unison-runtime/unison-runtime.cabal

@@ -29,6 +29,10 @@ flag optchecks
   manual: True
   default: False
 
+flag codeserialchecks
+  manual: True
+  default: False
+
 flag stackchecks
   manual: True
   default: False
@@ -161,6 +165,8 @@ library
     cpp-options: -DOPT_CHECK
     build-depends:
         inspection-testing
+  if flag(codeserialchecks)
+    cpp-options: -DCODE_SERIAL_CHECK
   if flag(dumpcore)
     ghc-options: -ddump-simpl -ddump-stg-final -ddump-to-file -dsuppress-coercions -dsuppress-idinfo -dsuppress-module-prefixes -ddump-str-signatures -ddump-simpl-stats