Ensure floated dictionaries are in scope (again)wip/T21689

In the Specialiser, we missed one more call to
bringFloatedDictsIntoScope (see #21391).

This omission led to #21689. The problem is that the call
to `rewriteClassOps` needs to have in scope any dictionaries
floated out of the arguments we have just specialised.

Easy fix.

4 files changed, 240 insertions, 2 deletions

diff --git a/compiler/GHC/Core/Opt/Specialise.hs b/compiler/GHC/Core/Opt/Specialise.hs
index 74ee8d1f5f..2e18049dd7 100644
--- a/compiler/GHC/Core/Opt/Specialise.hs
+++ b/compiler/GHC/Core/Opt/Specialise.hs
@@ -1114,7 +1114,10 @@ specExpr env (Tick tickish body)
 specExpr env expr@(App {})
   = do { let (fun_in, args_in) = collectArgs expr
        ; (args_out, uds_args) <- mapAndCombineSM (specExpr env) args_in
-       ; let (fun_in', args_out') = rewriteClassOps env fun_in args_out
+       ; let env_args = env `bringFloatedDictsIntoScope` ud_binds uds_args
+                -- Some dicts may have floated out of args_in;
+                -- they should be in scope for rewriteClassOps (#21689)
+             (fun_in', args_out') = rewriteClassOps env_args fun_in args_out
        ; (fun_out', uds_fun) <- specExpr env fun_in'
        ; let uds_call = mkCallUDs env fun_out' args_out'
        ; return (fun_out' `mkApps` args_out', uds_fun `thenUDs` uds_call `thenUDs` uds_args) }
@@ -1488,7 +1491,8 @@ specCalls spec_imp env dict_binds existing_rules calls_for_me fn rhs
 
     already_covered :: [CoreRule] -> [CoreExpr] -> Bool
     already_covered new_rules args      -- Note [Specialisations already covered]
-       = isJust (specLookupRule env fn args (new_rules ++ existing_rules))
+       = isJust (specLookupRule env_with_dict_bndrs fn args
+                                (new_rules ++ existing_rules))
          -- NB: we look both in the new_rules (generated by this invocation
          --     of specCalls), and in existing_rules (passed in to specCalls)
 
diff --git a/testsuite/tests/simplCore/should_compile/T21689.hs b/testsuite/tests/simplCore/should_compile/T21689.hs
new file mode 100644
index 0000000000..44ddba5944
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T21689.hs
@@ -0,0 +1,41 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE TypeFamilies #-}
+module Unboxed (Vec, MVec, Unbox) where
+
+import Data.Data (Data(..), Constr, DataType, Fixity(..), mkConstr, mkDataType)
+import Data.Typeable (Typeable)
+import GHC.TypeLits (Nat)
+
+import T21689a
+
+data family Vec  (n :: Nat) a
+data family MVec (n :: Nat) s a
+
+class (Arity n, IVector (Vec n) a, MVector (MVec n) a) => Unbox n a
+
+type instance Mutable (Vec n) = MVec n
+
+type instance Dim  (Vec  n) = n
+type instance DimM (MVec n) = n
+
+instance (Unbox n a) => Vector (Vec n) a where
+  construct  = constructVec
+  inspect    = inspectVec
+  {-# INLINE construct  #-}
+  {-# INLINE inspect    #-}
+
+instance (Typeable n, Unbox n a, Data a) => Data (Vec n a) where
+  gfoldl       = gfoldl'
+  gunfold      = gunfold'
+  toConstr   _ = con_Vec
+  dataTypeOf _ = ty_Vec
+
+ty_Vec :: DataType
+ty_Vec  = mkDataType "Data.Vector.Fixed.Unboxed.Vec" [con_Vec]
+
+con_Vec :: Constr
+con_Vec = mkConstr ty_Vec "Vec" [] Prefix
diff --git a/testsuite/tests/simplCore/should_compile/T21689a.hs b/testsuite/tests/simplCore/should_compile/T21689a.hs
new file mode 100644
index 0000000000..e95f4343d5
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T21689a.hs
@@ -0,0 +1,192 @@
+{-# LANGUAGE ConstraintKinds       #-}
+{-# LANGUAGE DataKinds             #-}
+{-# LANGUAGE DeriveDataTypeable    #-}
+{-# LANGUAGE FlexibleContexts      #-}
+{-# LANGUAGE FlexibleInstances     #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE PolyKinds             #-}
+{-# LANGUAGE RankNTypes            #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
+{-# LANGUAGE TypeFamilies          #-}
+{-# LANGUAGE TypeOperators         #-}
+{-# LANGUAGE UndecidableInstances  #-}
+module T21689a
+  ( Arity
+  , Dim
+  , DimM
+  , IVector
+  , Mutable
+  , MVector
+  , Vector(..)
+  , constructVec
+  , inspectVec
+  , gfoldl'
+  , gunfold'
+  ) where
+
+import Control.Monad.ST (ST, runST)
+import Data.Coerce (coerce)
+import Data.Data (Data)
+import Data.Functor.Const (Const(..))
+import Data.Kind (Type)
+import GHC.TypeLits (KnownNat, Nat, type (+), type (-))
+
+-----
+-- Data.Vector.Fixed.Cont
+-----
+
+data PeanoNum = Z
+              | S PeanoNum
+
+type family Peano (n :: Nat) :: PeanoNum where
+  Peano 0 = 'Z
+  Peano n = 'S (Peano (n - 1))
+
+type family Fn (n :: PeanoNum) (a :: Type) (b :: Type) where
+  Fn 'Z     a b = b
+  Fn ('S n) a b = a -> Fn n a b
+
+newtype Fun n a b = Fun { unFun :: Fn n a b }
+
+type family Dim (v :: Type -> Type) :: Nat
+
+class Arity (Dim v) => Vector v a where
+  construct :: Fun (Peano (Dim v)) a (v a)
+
+  inspect   :: v a -> Fun (Peano (Dim v)) a b -> b
+
+type Arity n = ( ArityPeano (Peano n)
+               , KnownNat n
+               , Peano (n+1) ~ 'S (Peano n)
+               )
+
+class ArityPeano n where
+  accum :: (forall k. t ('S k) -> a -> t k)
+        -> (t 'Z -> b)
+        -> t n
+        -> Fun n a b
+
+  applyFun :: (forall k. t ('S k) -> (a, t k))
+           -> t n
+           -> (CVecPeano n a, t 'Z)
+
+  gunfoldF :: (Data a)
+           => (forall b x. Data b => c (b -> x) -> c x)
+           -> T_gunfold c r a n -> c r
+
+newtype T_gunfold c r a n = T_gunfold (c (Fn n a r))
+
+gfoldl' :: forall c v a. (Vector v a, Data a)
+        => (forall x y. Data x => c (x -> y) -> x -> c y)
+        -> (forall x  . x -> c x)
+        -> v a -> c (v a)
+gfoldl' f inj v
+  = inspect v
+  $ gfoldlF f (inj $ unFun (construct :: Fun (Peano (Dim v)) a (v a)))
+
+gunfold' :: forall con c v a. (Vector v a, Data a)
+         => (forall b r. Data b => c (b -> r) -> c r)
+         -> (forall r. r -> c r)
+         -> con -> c (v a)
+gunfold' f inj _
+  = gunfoldF f gun
+  where
+    con = construct                   :: Fun (Peano (Dim v)) a (v a)
+    gun = T_gunfold (inj $ unFun con) :: T_gunfold c (v a) a (Peano (Dim v))
+
+gfoldlF :: (ArityPeano n, Data a)
+        => (forall x y. Data x => c (x -> y) -> x -> c y)
+        -> c (Fn n a r) -> Fun n a (c r)
+gfoldlF f c0 = accum
+  (\(T_gfoldl c) x -> T_gfoldl (f c x))
+  (\(T_gfoldl c)   -> c)
+  (T_gfoldl   c0)
+
+newtype T_gfoldl c r a n = T_gfoldl (c (Fn n a r))
+
+newtype ContVec n a = ContVec (forall r. Fun (Peano n) a r -> r)
+
+type instance Dim (ContVec n) = n
+
+instance Arity n => Vector (ContVec n) a where
+  construct = accum
+    (\(T_mkN f) a -> T_mkN (f . consPeano a))
+    (\(T_mkN f)   -> toContVec $ f (CVecPeano unFun))
+    (T_mkN id)
+  inspect (ContVec c) f = c f
+  {-# INLINE construct #-}
+  {-# INLINE inspect   #-}
+
+newtype T_mkN n_tot a n = T_mkN (CVecPeano n a -> CVecPeano n_tot a)
+
+toContVec :: CVecPeano (Peano n) a -> ContVec n a
+toContVec = coerce
+
+newtype CVecPeano n a = CVecPeano (forall r. Fun n a r -> r)
+
+consPeano :: a -> CVecPeano n a -> CVecPeano ('S n) a
+consPeano a (CVecPeano cont) = CVecPeano $ \f -> cont $ curryFirst f a
+{-# INLINE consPeano #-}
+
+curryFirst :: Fun ('S n) a b -> a -> Fun n a b
+curryFirst = coerce
+{-# INLINE curryFirst #-}
+
+apply :: Arity n
+      => (forall k. t ('S k) -> (a, t k))
+      -> t (Peano n)
+      -> ContVec n a
+{-# INLINE apply #-}
+apply step' z = toContVec $ fst (applyFun step' z)
+
+-----
+-- Data.Vector.Fixed.Mutable
+-----
+
+type family Mutable (v :: Type -> Type) :: Type -> Type -> Type
+
+type family DimM (v :: Type -> Type -> Type) :: Nat
+
+class (Arity (DimM v)) => MVector v a where
+  new   :: PrimMonad m => m (v (PrimState m) a)
+
+  unsafeWrite :: PrimMonad m => v (PrimState m) a -> Int -> a -> m ()
+
+class (Dim v ~ DimM (Mutable v), MVector (Mutable v) a) => IVector v a where
+  unsafeFreeze :: PrimMonad m => Mutable v (PrimState m) a -> m (v a)
+
+  unsafeIndex :: v a -> Int -> a
+
+inspectVec :: forall v a b. (Arity (Dim v), IVector v a) => v a -> Fun (Peano (Dim v)) a b -> b
+{-# INLINE inspectVec #-}
+inspectVec v
+  = inspect cv
+  where
+    cv :: ContVec (Dim v) a
+    cv = apply (\(Const i) -> (unsafeIndex v i, Const (i+1)))
+               (Const 0 :: Const Int (Peano (Dim v)))
+
+constructVec :: forall v a. (Arity (Dim v), IVector v a) => Fun (Peano (Dim v)) a (v a)
+{-# INLINE constructVec #-}
+constructVec =
+  accum step
+        (\(T_new _ st) -> runST $ unsafeFreeze =<< st :: v a)
+        (T_new 0 new :: T_new v a (Peano (Dim v)))
+
+data T_new v a n = T_new Int (forall s. ST s (Mutable v s a))
+
+step :: (IVector v a) => T_new v a ('S n) -> a -> T_new v a n
+step (T_new i st) x = T_new (i+1) $ do
+  mv <- st
+  unsafeWrite mv i x
+  return mv
+
+-----
+-- Control.Monad.Primitive
+-----
+
+class Monad m => PrimMonad m where
+  type PrimState m
+
+instance PrimMonad (ST s) where
+  type PrimState (ST s) = s
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 5a018cdb2d..b9b1956f51 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -414,3 +414,4 @@ test('T21391', normal, compile, ['-O -dcore-lint'])
 test('T21391a', normal, compile, ['-O -dcore-lint'])
 # We don't want to see a thunk allocation for the insertBy expression after CorePrep.
 test('T21392', [ grep_errmsg(r'sat.* :: \[\(.*Unique, .*Int\)\]'), expect_broken(21392) ], compile, ['-O -ddump-prep -dno-typeable-binds -dsuppress-uniques'])
+test('T21689', [extra_files(['T21689a.hs'])], multimod_compile, ['T21689', '-v0 -O'])