fix: EvalFull causes capture in BETA reduction

We add a unit test to avoid regressions in the future. This test is
rather detailed, testing multiple intermediate steps of the reduction to
illustrate the slightly subtle process.

primer/src/Primer/EvalFull.hs

@@ -98,6 +98,8 @@ import Primer.Typecheck (instantiateValCons', lookupConstructor, mkTAppCon)
 import Primer.Zipper (
   ExprZ,
   TypeZ,
+  bindersBelow,
+  bindersBelowTy,
   down,
   focus,
   focusType,
@@ -126,8 +128,10 @@ data Redex
     ElideLet SomeLocal Expr
   | -- (λx.t : S -> T) s  ~>  let x = s:S in t : T
     Beta LVarName Expr Type Type Expr
-  | -- (Λa.t : ∀b.T) S  ~>  lettype b = S in (lettype a = S in t) : T
+  | -- (Λa.t : ∀b.T) S  ~>  lettype b = S in (lettype a = S in t) : T for b not free in S,t
     BETA TyVarName Expr TyVarName Type Type
+  | -- (Λa.t : ∀b.T) S  ~> letType c = b in letType b = c in (Λa.t : ∀b.T) S  for b free in t or S, and fresh c
+    RenameBETA TyVarName Expr (Set Name)
   | -- case C as : T of ... ; C xs -> e ; ...   ~>  let xs=as:As in e for constructor C of type T, where args have types As
     -- also the non-annotated case, as we consider constructors to be synthesisable
     -- case C as of ... ; C xs -> e ; ...   ~>  let xs=as:As in e for constructor C of type T, where args have types As
@@ -359,7 +363,21 @@ viewRedex tydefs globals dir = \case
   Var _ (GlobalVarRef x) | Just (DefAST y) <- x `M.lookup` globals -> pure $ pure $ InlineGlobal x y
   App _ (Ann _ (Lam _ x t) (TFun _ src tgt)) s -> pure $ pure $ Beta x t src tgt s
   e@App{} -> pure . ApplyPrimFun . thd3 <$> tryPrimFun (M.mapMaybe defPrim globals) e
-  APP _ (Ann _ (LAM _ a t) (TForall _ b _ ty1)) ty2 -> pure $ pure $ BETA a t b ty1 ty2
+  e@(APP _ (Ann _ (LAM _ a t) (TForall _ b _ ty1)) ty2) ->
+    -- We would like to say (Λa.t : ∀b.T) S  ~> (letType a = S in t) : (letType b = S in T)
+    -- but we do not have letTypes inside types, so the best we can do is
+    -- (Λa.t : ∀b.T) S  ~> letType b = S in ((letType a = S in t) : T)
+    -- We need to be careful if a /= b: as this can capture a 'b' inside 'S' or 't'.
+    -- Thus if necessary we do some renaming
+    -- (Λa.t : ∀b.T) S  ~> letType c = b in letType b = c in (Λa.t : ∀b.T) S  for b free in t or S, and fresh c
+    -- We then ensure the delicate property that we reduce the b=c first, then the BETA, then the c=b
+    let fvs = freeVars t <> S.map unLocalName (freeVarsTy ty2)
+        -- we only really need to avoid free things, but avoiding bound
+        -- things means we do not need to do any further renaming
+        bvs = bindersBelow (focus t) <> S.map unLocalName (bindersBelowTy $ focus ty2)
+     in if a /= b && S.member (unLocalName b) fvs
+          then pure $ pure $ RenameBETA b e (fvs <> bvs)
+          else pure $ pure $ BETA a t b ty1 ty2
   e | Just r <- viewCaseRedex tydefs e -> Just r
   Ann _ t ty | Chk <- dir, concreteTy ty -> pure $ pure $ Upsilon t ty
   _ -> Nothing
@@ -432,6 +450,14 @@ findRedex tydefs globals dir = go . focus
         -- This case should have caught by the TC: a term var is bound by a lettype
         LLetType _ _ -> Nothing
       -- We have found something like
+      --   letType c=b in (Λa.t : ∀b.T) S
+      -- where inlining 'c' would block the BETA redex. Thus we do the BETA first
+      APP _ (Ann _ (LAM _ a t) (TForall _ b1 _ ty1)) ty2
+        | LLetType c (TVar _ b2) <- l
+        , b1 == b2
+        , S.member (unLocalName c) (freeVars t <> S.map unLocalName (freeVarsTy ty2)) ->
+            pure $ RExpr ez $ BETA a t b1 ty1 ty2
+      -- We have found something like
       --   let x=y in let y=z in t
       -- to substitute the 'x' inside 't' we would need to rename the 'let y'
       -- binding, but that is implemented in terms of let:
@@ -493,9 +519,14 @@ runRedex = \case
   -- (λx.t : S -> T) s  ~>  let x = s:S in t : T
   Beta x t tyS tyT s -> let_ x (pure s `ann` pure tyS) (pure t) `ann` pure tyT
   -- (Λa.t : ∀b.T) S  ~>  lettype b = S in (lettype a = S in t) : T
+  --  if b is not free in t or S
   BETA a t b tyT tyS
     | a == b -> letType a (pure tyS) $ pure t `ann` pure tyT
     | otherwise -> letType b (regenerateTypeIDs tyS) $ letType a (pure tyS) (pure t) `ann` pure tyT
+  -- (Λa.t : ∀b.T) S  ~> letType c = b in letType b = c in (Λa.t : ∀b.T) S  for b free in t or S, and fresh c
+  RenameBETA b beta avoid -> do
+    c <- freshLocalName' avoid
+    letType c (tvar b) $ letType b (tvar c) $ pure beta
   -- case C as : T of ... ; C xs -> e ; ...   ~>  let xs=as:As in e for constructor C of type T, where args have types As
   -- (and also the non-annotated-constructor case)
   CaseRedex _ as _ xs e -> foldrM (\(x, (a, tyA)) t -> let_ x (pure a `ann` pure tyA) (pure t)) e (zip xs as)

primer/test/Tests/EvalFull.hs

@@ -4,6 +4,7 @@ import Foreword hiding (unlines)
 
 import Data.Generics.Uniplate.Data (universe)
 import Data.List (span, (\\))
+import qualified Data.List.NonEmpty as NE
 import qualified Data.Map as M
 import qualified Data.Map as Map
 import qualified Data.Set as S
@@ -53,6 +54,7 @@ import Primer.Primitives (primitiveGVar, primitiveModule, tChar, tInt)
 import Primer.Typecheck (
   SmartHoles (NoSmartHoles),
   check,
+  extendGlobalCxt,
   typeDefs,
  )
 import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, (@?=))
@@ -61,7 +63,7 @@ import TestUtils (exprIDs, withPrimDefs, zeroIDs)
 import Tests.Action.Prog (runAppTestM)
 import Tests.Eval ((~=))
 import Tests.Gen.Core.Typed (checkTest)
-import Tests.Typecheck (runTypecheckTestM)
+import Tests.Typecheck (runTypecheckTestM, runTypecheckTestMWithPrims)
 
 unit_1 :: Assertion
 unit_1 =
@@ -470,6 +472,86 @@ unit_type_preservation_case_regression_ty =
         s1 <~==> Left (TimedOut expected1)
         s2 <~==> Left (TimedOut expected2)
 
+-- Previously EvalFull reducing a BETA expression could result in variable
+-- capture. We would reduce (Λa.t : ∀b.T) S to
+-- let b = S in (let a = S in t) : T
+-- The outer let binding could capture references within S or t.
+unit_type_preservation_BETA_regression :: Assertion
+unit_type_preservation_BETA_regression =
+  let (((exprA, expectedAs), (exprB, expectedBs)), maxID) = create $ do
+        let n = "a145"
+        -- The 'A' sequence previously captured in the type "S" above
+        let eA' b =
+              ( lAM "a" (lam "c" $ emptyHole `ann` tvar "a")
+                  `ann` tforall "b" KType (tcon tNat `tfun` tvar "b")
+              )
+                `aPP` (tvar b `tapp` tcon tBool)
+        eA <- lAM "b" $ eA' "b"
+        -- Do some renaming to set up
+        expectA1 <- lAM "b" $ letType n (tvar "b") $ letType "b" (tvar n) $ eA' "b"
+        -- Resolve the renaming
+        expectA3 <- lAM "b" $ letType n (tvar "b") $ eA' n
+        -- Do the BETA step
+        expectA4 <-
+          lAM "b" $
+            letType n (tvar "b") $
+              letType "b" (tvar n `tapp` tcon tBool) $
+                letType
+                  "a"
+                  (tvar n `tapp` tcon tBool)
+                  (lam "c" $ emptyHole `ann` tvar "a")
+                  `ann` (tcon tNat `tfun` tvar "b")
+        -- Resolve all the letTypes
+        expectA11 <-
+          lAM "b" $
+            lam "c" (emptyHole `ann` (tvar "b" `tapp` tcon tBool))
+              `ann` (tcon tNat `tfun` (tvar "b" `tapp` tcon tBool))
+        -- The 'B' sequence previously captured in the term "t" above
+        let eB' b =
+              ( lAM "a" (gvar foo `aPP` (tvar b `tapp` tcon tBool))
+                  `ann` tforall "b" KType (tcon tNat)
+              )
+                `aPP` tcon tChar
+        eB <- lAM "b" $ eB' "b"
+        -- Do some renaming to set up
+        expectB1 <- lAM "b" $ letType n (tvar "b") $ letType "b" (tvar n) $ eB' "b"
+        -- Resolve the renaming
+        expectB3 <- lAM "b" $ letType n (tvar "b") $ eB' n
+        -- Do the BETA step
+        expectB4 <-
+          lAM "b" $
+            letType n (tvar "b") $
+              letType "b" (tcon tChar) $
+                letType "a" (tcon tChar) (gvar foo `aPP` (tvar n `tapp` tcon tBool))
+                  `ann` tcon tNat
+        -- Resolve all the letTypes (and elide an annotation)
+        expectB9 <- lAM "b" $ gvar foo `aPP` (tvar "b" `tapp` tcon tBool)
+        -- Note that the reduction of eA and eB take slightly
+        -- different paths: we do not remove the annotation in eA
+        -- because it has an occurrence of a type variable and is thus
+        -- not "concrete"
+        pure
+          ( (eA, [(1, expectA1), (3, expectA3), (4, expectA4), (11, expectA11)])
+          , (eB, [(1, expectB1), (3, expectB3), (4, expectB4), (9, expectB9)])
+          )
+      sA n = evalFullTest maxID builtinTypes mempty n Chk exprA
+      sB n = evalFullTest maxID builtinTypes mempty n Chk exprB
+      tyA = TForall () "c" (KFun KType KType) $ TFun () (TCon () tNat) (TApp () (TVar () "c") (TCon () tBool))
+      tyB = TForall () "c" (KFun KType KType) $ TCon () tNat
+      foo = qualifyName (ModuleName ["M"]) "foo"
+      fooTy = TForall () "d" KType $ TCon () tNat
+      tmp ty e = case runTypecheckTestMWithPrims NoSmartHoles $
+        local (extendGlobalCxt [(foo, fooTy)]) $ check ty e of
+        Left err -> assertFailure $ show err
+        Right _ -> pure ()
+   in do
+        tmp tyA exprA
+        for_ expectedAs $ \(n, e) -> sA n <~==> Left (TimedOut e)
+        tmp tyA $ snd $ NE.last expectedAs
+        tmp tyB exprB
+        for_ expectedBs $ \(n, e) -> sB n <~==> Left (TimedOut e)
+        tmp tyB $ snd $ NE.last expectedBs
+
 -- Previously EvalFull reducing a let expression could result in variable
 -- capture. We would reduce 'Λx. let x = _ :: x in x'
 -- to (eliding annotations) 'Λx. let x = _ :: x in _ :: x', where the