From dbbd142418cf37b70c9fb40325062ed6954e8ed7 Mon Sep 17 00:00:00 2001
From: Leon Schoorl <leon@qbaylogic.com>
Date: Tue, 27 Feb 2024 11:23:26 +0100
Subject: [PATCH] Make applyDebug's free variable check take the context into
 account (#2624)

Fixes #2623

(cherry picked from commit 0aa341ab831bae1769ea1d9c2175e9bffde4e36b)
---
 clash-lib/src/Clash/Rewrite/Util.hs        | 29 +++++++++++++++++-----
 tests/Main.hs                              |  1 +
 tests/shouldwork/Issues/T2623CaseConFVs.hs | 19 ++++++++++++++
 3 files changed, 43 insertions(+), 6 deletions(-)
 create mode 100644 tests/shouldwork/Issues/T2623CaseConFVs.hs

diff --git a/clash-lib/src/Clash/Rewrite/Util.hs b/clash-lib/src/Clash/Rewrite/Util.hs
index 420fa29574..15c3bc6d1c 100644
--- a/clash-lib/src/Clash/Rewrite/Util.hs
+++ b/clash-lib/src/Clash/Rewrite/Util.hs
@@ -2,7 +2,7 @@
   Copyright  :  (C) 2012-2016, University of Twente,
                     2016     , Myrtle Software Ltd,
                     2017     , Google Inc.,
-                    2021-2022, QBayLogic B.V.
+                    2021-2023, QBayLogic B.V.
   License    :  BSD2 (see the file LICENSE)
   Maintainer :  QBayLogic B.V. <devops@qbaylogic.com>
 
@@ -77,7 +77,7 @@ import           Clash.Core.Type             (Type (..), normalizeType)
 import           Clash.Core.Var
   (Id, IdScope (..), TyVar, Var (..), mkGlobalId, mkLocalId, mkTyVar)
 import           Clash.Core.VarEnv
-  (InScopeSet, extendInScopeSet, extendInScopeSetList, mkInScopeSet,
+  (InScopeSet, extendInScopeSet, extendInScopeSetList, mkInScopeSet, notElemInScopeSet,
    uniqAway, uniqAway', mapVarEnv, eltsVarEnv, unitVarSet, emptyVarEnv,
    mkVarEnv, eltsVarSet, elemVarEnv, lookupVarEnv, extendVarEnv, elemVarSet,
    differenceVarEnv)
@@ -173,12 +173,13 @@ apply = \s rewrite ctx expr0 -> do
     return ()
 
   if isDebugging opts
-    then applyDebug s expr0 hasChanged expr1
+    then applyDebug ctx s expr0 hasChanged expr1
     else return expr1
 {-# INLINE apply #-}
 
 applyDebug
-  :: String
+  :: TransformContext
+  -> String
   -- ^ Name of the transformation
   -> Term
   -- ^ Original expression
@@ -187,7 +188,7 @@ applyDebug
   -> Term
   -- ^ New expression
   -> RewriteMonad extra Term
-applyDebug name exprOld hasChanged exprNew = do
+applyDebug ctx name exprOld hasChanged exprNew = do
   nTrans <- Lens.use transformCounter
   opts <- Lens.view debugOpts
 
@@ -212,9 +213,14 @@ applyDebug name exprOld hasChanged exprNew = do
       let beforeTy          = inferCoreTypeOf tcm exprOld
           beforeFV          = Lens.setOf freeLocalVars exprOld
           afterTy           = inferCoreTypeOf tcm exprNew
-          afterFV           = Lens.setOf freeLocalVars exprNew
+          afterFV           = filterFVs (Lens.setOf freeLocalVars exprNew)
           newFV             = not (afterFV `Set.isSubsetOf` beforeFV)
           accidentalShadows = findAccidentialShadows exprNew
+          -- see NOTE [Filter free variables]
+          allowNewFVsFromCtx = name == "caseCon"
+          filterFVs | allowNewFVsFromCtx = Set.filter notInCtx
+                    | otherwise = id
+          notInCtx v = notElemInScopeSet v (tfInScope ctx)
 
       Monad.when newFV $
               error ( concat [ $(curLoc)
@@ -266,6 +272,17 @@ applyDebug name exprOld hasChanged exprNew = do
     before = showPpr exprOld
     after  = showPpr exprNew
 
+-- NOTE [Filter free variables]
+-- Since [Give evaluator acces to inscope let-bindings #2571](https://github.com/clash-lang/clash-compiler/pull/2571)
+-- the evaluator can rewrite expressions using let bindings from the 'TransformContext',
+-- these bindings may reference other things bound in the context which weren't
+-- in the expression before, and in doing so introduces new free variables and
+-- fails this check for new free variables.
+-- To prevent this we filter out all variables from bound in the context.
+-- But only during a caseCon transformation, to not weaken this check unnecessarily.
+
+
+
 -- | Perform a transformation on a Term
 runRewrite
   :: String
diff --git a/tests/Main.hs b/tests/Main.hs
index 178c980464..6d5a7a0fbb 100755
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -789,6 +789,7 @@ runClashTest = defaultMain $ clashTestRoot
 #endif
         , outputTest "T2542" def{hdlTargets=[VHDL]}
         , runTest "T2593" def{hdlSim=[]}
+        , runTest "T2623CaseConFVs" def{hdlLoad=[],hdlSim=[],hdlTargets=[VHDL]}
         ] <>
         if compiledWith == Cabal then
           -- This tests fails without environment files present, which are only
diff --git a/tests/shouldwork/Issues/T2623CaseConFVs.hs b/tests/shouldwork/Issues/T2623CaseConFVs.hs
new file mode 100644
index 0000000000..1a53546f93
--- /dev/null
+++ b/tests/shouldwork/Issues/T2623CaseConFVs.hs
@@ -0,0 +1,19 @@
+module T2623CaseConFVs where
+import Clash.Prelude
+
+topEntity = foo @System
+
+foo :: forall dom. Signal dom (Vec 1 (Signed 2)) -> Signal dom Bool
+foo = \input ->
+    let
+      scs :: Signal dom (Vec 1 Bool)
+      scs = bundle $ map f $ unbundle input
+    in
+      fmap bar scs
+
+
+bar :: (KnownNat n) => Vec (n+1) Bool -> Bool
+bar = fold (&&) . map (id)
+
+f :: Signal dom a -> Signal dom Bool
+f = const $ pure $ True