GitLab

Marge Bot pushed to branch master at Glasgow Haskell Compiler / GHC

Commits:

1d4a1229

by sheaf at 2025-11-27T17:58:02-05:00

SimpleOpt: don't subst in pushCoercionIntoLambda

It was noticed in #26589 that the change in 15b311be was incorrect:
the simple optimiser carries two different substitution-like pieces of
information: 'soe_subst' (from InVar to OutExpr) and 'soe_inl'
(from InId to InExpr). It is thus incorrect to have 'pushCoercionIntoLambda'
apply the substitution from 'soe_subst' while discarding 'soe_inl'
entirely, which is what was done in 15b311be.

Instead, we change back pushCoercionIntoLambda to take an InScopeSet,
and optimise the lambda before calling 'pushCoercionIntoLambda' to avoid
mixing InExpr with OutExpr, or mixing two InExpr with different
environments. We can then call 'soeZapSubst' without problems.

Fixes #26588 #26589

6 changed files:

compiler/GHC/Core/Opt/Arity.hs
compiler/GHC/Core/SimpleOpt.hs
compiler/GHC/Tc/Gen/App.hs
+ testsuite/tests/simplCore/should_compile/T26588.hs
+ testsuite/tests/simplCore/should_compile/T26589.hs
testsuite/tests/simplCore/should_compile/all.T

Changes:

compiler/GHC/Core/Opt/Arity.hs

@@ -2993,12 +2993,12 @@ pushCoValArg co
      Pair tyL tyR = coercionKind co
  pushCoercionIntoLambda
 -    :: HasDebugCallStack => Subst -> InVar -> InExpr -> OutCoercionR -> Maybe (OutVar, OutExpr)
 +    :: HasDebugCallStack => InScopeSet -> Var -> CoreExpr -> CoercionR -> Maybe (Var, CoreExpr)
  -- This implements the Push rule from the paper on coercions
  --    (\x. e) |> co
  -- ===>
  --    (\x'. e |> co')
 -pushCoercionIntoLambda subst x e co
 +pushCoercionIntoLambda in_scope x e co
      | assert (not (isTyVar x) && not (isCoVar x)) True
      , Pair s1s2 t1t2 <- coercionKind co
      , Just {}              <- splitFunTy_maybe s1s2
@@ -3011,9 +3011,9 @@ pushCoercionIntoLambda subst x e co
            -- Should we optimize the coercions here?
            -- Otherwise they might not match too well
            x' = x `setIdType` t1 `setIdMult` w1
 -          in_scope' = substInScopeSet subst `extendInScopeSet` x'
 +          in_scope' = in_scope `extendInScopeSet` x'
            subst' =
 -            extendIdSubst (setInScope subst in_scope')
 +            extendIdSubst (setInScope emptySubst in_scope')
+               x
                (mkCast (Var x') (mkSymCo co1))
              -- We substitute x' for x, except we need to preserve types.

compiler/GHC/Core/SimpleOpt.hs

@@ -393,12 +393,19 @@ simple_app env e0@(Lam {}) as0@(_:_)
        = wrapLet mb_pr $ do_beta env'' body as
        where (env', b') = subst_opt_bndr env b
 -    do_beta env e@(Lam b body) as@(CastIt co:rest)
 -      -- See Note [Desugaring unlifted newtypes]
 +    -- See Note [Eliminate casts in function position]
 +    do_beta env e@(Lam b _) as@(CastIt out_co:rest)
        | isNonCoVarId b
 -      , Just (b', body') <- pushCoercionIntoLambda (soe_subst env) b body co
 +      -- Optimise the inner lambda to make it an 'OutExpr', which makes it
 +      -- possible to call 'pushCoercionIntoLambda' with the 'OutCoercion' 'co'.
 +      -- This is kind of horrible, as for nested casted lambdas with a big body,
 +      -- we will repeatedly optimise the body (once for each binder). However,
 +      -- we need to do this to avoid mixing 'InExpr' and 'OutExpr', or two
 +      -- 'InExpr' with different environments (getting this wrong caused #26588 & #26589.)
 +      , Lam out_b out_body <- simple_app env e []
 +      , Just (b', body') <- pushCoercionIntoLambda (soeInScope env) out_b out_body out_co
        = do_beta (soeZapSubst env) (Lam b' body') rest
 -        -- soeZapSubst: pushCoercionIntoLambda applies the substitution
 +        -- soeZapSubst: we've already optimised everything (the lambda and 'rest') by now.
        | otherwise
        = rebuild_app env (simple_opt_expr env e) as
@@ -511,7 +518,31 @@ TL;DR: To avoid the rest of the compiler pipeline seeing these bad lambas, we
  rely on the simple optimiser to both inline the newtype unfolding and
  subsequently deal with the resulting lambdas (either beta-reducing them
  altogether or pushing coercions into them so that they satisfy the
 -representation-polymorphism invariants).
 +representation-polymorphism invariants). See Note [Eliminate casts in function position].
++
 +[Alternative approach] (GHC ticket #26608)
++
 +  We could instead, in the typechecker, emit a special form (a new constructor
 +  of XXExprGhcTc) for instantiations of representation-polymorphic unlifted
 +  newtypes (whether applied to a value argument or not):
++
 +    UnliftedNT :: DataCon -> [Type] -> Coercion -> XXExprGhcTc
++
 +  where "UnliftedNT nt_con [ty1, ...] co" represents the expression:
++
 +    ( nt_con @ty1 ... ) |> co
++
 +  The desugarer would then turn these AST nodes into appropriate Core, doing
 +  what the simple optimiser does today:
 +    - inline the compulsory unfolding of the newtype constructor
 +    - apply it to its type arguments and beta reduce
 +    - push the coercion into the resulting lambda
++
 +  This would have several advantages:
 +    - the desugarer would never produce "invalid" Core that needs to be
 +      tidied up by the simple optimiser,
 +    - the ugly and inefficient implementation described in
 +      Note [Eliminate casts in function position] could be removed.
  Wrinkle [Unlifted newtypes with wrappers]
@@ -717,50 +748,49 @@ rhss here.
  Note [Eliminate casts in function position]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 -Consider the following program:
 +Due to the current implementation strategy for representation-polymorphic
 +unlifted newtypes, as described in Note [Desugaring unlifted newtypes], we rely
 +on the simple optimiser to push coercions into lambdas, such as in the following
 +example:
    type R :: Type -> RuntimeRep
 -  type family R a where { R Float = FloatRep; R Double = DoubleRep }
 -  type F :: forall (a :: Type) -> TYPE (R a)
 -  type family F a where { F Float = Float#  ; F Double = Double# }
 +  type family R a where { R Int = IntRep }
 +  type F :: forall a -> TYPE (R a)
 +  type family F a where { F Int = Int# }
 -  type N :: forall (a :: Type) -> TYPE (R a)
    newtype N a = MkN (F a)
 -As MkN is a newtype, its unfolding is a lambda which wraps its argument
 -in a cast:
+-
 -  MkN :: forall (a :: Type). F a -> N a
 -  MkN = /\a \(x::F a). x |> co_ax
 -    -- recall that F a :: TYPE (R a)
+-
 -This is a representation-polymorphic lambda, in which the binder has an unknown
 -representation (R a). We can't compile such a lambda on its own, but we can
 -compile instantiations, such as `MkN @Float` or `MkN @Double`.
 +Now, an instantiated occurrence of 'MkN', such as 'MkN @Int' (whether applied
 +to a value argument or not) will lead, after inlining the compulsory unfolding
 +of 'MkN', to a lambda fo the form:
 -Our strategy to avoid running afoul of the representation-polymorphism
 -invariants of Note [Representation polymorphism invariants] in GHC.Core is thus:
 +  ( \ ( x :: F Int ) -> body ) |> co
 -  1. Give the newtype a compulsory unfolding (it has no binding, as we can't
 -     define lambdas with representation-polymorphic value binders in source Haskell).
 -  2. Rely on the optimiser to beta-reduce away any representation-polymorphic
 -     value binders.
 +    where
 +      co :: ( F Int -> res ) ~# ( Int# -> res )
 -For example, consider the application
 +The problem is that we now have a lambda abstraction whose binder does not have a
 +fixed RuntimeRep in the sense of Note [Fixed RuntimeRep] in GHC.Tc.Utils.Concrete.
 -    MkN @Float 34.0#
 +However, if we use 'pushCoercionIntoLambda', we end up with:
 -After inlining MkN we'll get
 +  ( \ ( x' :: Int# ) -> body' )
 -   ((/\a \(x:F a). x |> co_ax) @Float) |> co 34#
 +which satisfies the representation-polymorphism invariants of
 +Note [Representation polymorphism invariants] in GHC.Core.
 -where co :: (F Float -> N Float) ~ (Float# ~ N Float)
 +In conclusion:
 -But to actually beta-reduce that lambda, we need to push the 'co'
 -inside the `\x` with pushCoercionIntoLambda.  Hence the extra
 -equation for Cast-of-Lam in simple_app.
 +  1. The simple optimiser must push casts into lambdas.
 +  2. It must also deal with a situation such as (MkN @Int) |> co, where we first
 +     inline the compulsory unfolding of N. This means the simple optimiser must
 +     "peel off" the casts and optimise the inner expression first, to determine
 +     whether it is a lambda abstraction or not.
 -This is regrettably delicate.
 +This is regrettably delicate. If we could make sure the typechecker/desugarer
 +did not produce these bad lambdas in the first place (as described in
 +[Alternative approach] in Note [Desugaring unlifted newtypes]), we could
 +get rid of this ugly logic.
  Note [Preserve join-binding arity]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1673,7 +1703,7 @@ exprIsLambda_maybe ise@(ISE in_scope_set _) (Cast casted_e co)
      -- this implies that x is not in scope in gamma (makes this code simpler)
      , not (isTyVar x) && not (isCoVar x)
      , assert (not $ x `elemVarSet` tyCoVarsOfCo co) True
 -    , Just (x',e') <- pushCoercionIntoLambda (mkEmptySubst in_scope_set) x e co
 +    , Just (x',e') <- pushCoercionIntoLambda in_scope_set x e co
      , let res = Just (x',e',ts)
      = --pprTrace "exprIsLambda_maybe:Cast" (vcat [ppr casted_e,ppr co,ppr res)])
        res

compiler/GHC/Tc/Gen/App.hs

@@ -749,13 +749,13 @@ tcInstFun do_ql inst_final (tc_fun, fun_ctxt) fun_sigma rn_args
      go1 _pos acc fun_ty []
         | XExpr (ConLikeTc (RealDataCon dc)) <- tc_fun
         , isNewDataCon dc
 -       , [Scaled _ arg_ty] <- dataConOrigArgTys dc
 +       , [Scaled _ orig_arg_ty] <- dataConOrigArgTys dc
         , n_val_args == 0
         -- If we're dealing with an unsaturated representation-polymorphic
         -- UnliftedNewype, then perform a representation-polymorphism check.
         -- See Note [Representation-polymorphism checks for unsaturated unlifted newtypes]
         -- in GHC.Tc.Utils.Concrete.
 -       , not $ typeHasFixedRuntimeRep arg_ty
 +       , not $ typeHasFixedRuntimeRep orig_arg_ty
         = do { (wrap_co, arg_ty, res_ty) <-
                    matchActualFunTy (FRRRepPolyUnliftedNewtype dc)
                      (Just $ HsExprTcThing tc_fun)

testsuite/tests/simplCore/should_compile/T26588.hs

 +module T26588 ( getOptionSettingFromText ) where
++
 +import           Control.Applicative ( Const(..) )
 +import           Data.Map (Map)
 +import qualified Data.Map.Strict as Map
++
 +------------------------------------------------------------------------
 +-- ConfigState
++
 +data ConfigLeaf
 +data ConfigTrie = ConfigTrie !(Maybe ConfigLeaf) !ConfigMap
++
 +type ConfigMap = Map Int ConfigTrie
++
 +freshLeaf :: [Int] -> ConfigLeaf -> ConfigTrie
 +freshLeaf [] l     = ConfigTrie (Just l) mempty
 +freshLeaf (a:as) l = ConfigTrie Nothing (Map.singleton a (freshLeaf as l))
++
 +adjustConfigTrie :: Functor t => [Int] -> (Maybe ConfigLeaf -> t (Maybe ConfigLeaf)) -> Maybe (ConfigTrie) -> t (Maybe ConfigTrie)
 +adjustConfigTrie     as f Nothing                 = fmap (freshLeaf as) <$> f Nothing
 +adjustConfigTrie (a:as) f (Just (ConfigTrie x m)) = Just . ConfigTrie x <$> adjustConfigMap a as f m
 +adjustConfigTrie     [] f (Just (ConfigTrie x m)) = g <$> f x
 +  where g Nothing | Map.null m = Nothing
 +        g x' = Just (ConfigTrie x' m)
++
 +adjustConfigMap :: Functor t => Int -> [Int] -> (Maybe ConfigLeaf -> t (Maybe ConfigLeaf)) -> ConfigMap -> t ConfigMap
 +adjustConfigMap a as f = Map.alterF (adjustConfigTrie as f) a
++
 +getOptionSettingFromText :: Int -> [Int] -> ConfigMap -> IO ()
 +getOptionSettingFromText p ps = getConst . adjustConfigMap p ps f
 +  where
 +    f _ = Const (return ())

testsuite/tests/simplCore/should_compile/T26589.hs

 +module T26589 ( executeTest ) where
++
 +-- base
 +import Data.Coerce ( coerce )
 +import Data.Foldable ( foldMap )
++
 +--------------------------------------------------------------------------------
++
 +newtype Traversal f = Traversal { getTraversal :: f () }
++
 +instance Applicative f => Semigroup (Traversal f) where
 +  Traversal f1 <> Traversal f2 = Traversal $ f1 *> f2
 +instance Applicative f => Monoid (Traversal f) where
 +  mempty = Traversal $ pure ()
++
 +newtype Seq a = Seq (FingerTree (Elem a))
 +newtype Elem a = Elem { getElem :: a }
++
 +data FingerTree a
 +    = EmptyT
 +    | Deep !a (FingerTree a) !a
++
 +executeTest :: Seq () -> IO ()
 +executeTest fins = destroyResources
 +  where
 +    destroyResources :: IO ()
 +    destroyResources =
 +      getTraversal $
 +        flip foldMap1 fins $ \ _ ->
 +          Traversal $ return ()
++
 +foldMap1 :: forall m a. Monoid m => (a -> m) -> Seq a -> m
 +foldMap1 = coerce (foldMap2 :: (Elem a -> m) -> FingerTree (Elem a) -> m)
++
 +foldMap2 :: Monoid m => (Elem a -> m) -> FingerTree (Elem a) -> m
 +foldMap2 _ EmptyT = mempty
 +foldMap2 f' (Deep pr' m' sf') = f' pr' <> foldMapTree f' m' <> f' sf'
 +      where
 +        foldMapTree :: Monoid m => (a -> m) -> FingerTree a -> m
 +        foldMapTree _ EmptyT = mempty
 +        foldMapTree f (Deep pr m sf) =
 +            f pr <>
 +            foldMapTree f m <>
 +            f sf

testsuite/tests/simplCore/should_compile/all.T

@@ -544,6 +544,9 @@ test('T25883b', normal, compile_grep_core, [''])
  test('T25883c', normal, compile_grep_core, [''])
  test('T25883d', [extra_files(['T25883d_import.hs'])], multimod_compile_filter, ['T25883d', '-O -ddump-simpl -dno-typeable-binds -dsuppress-all -dsuppress-uniques', r'grep -e "y ="'])
 +test('T26588', normal, compile, ['-package containers -O'])
 +test('T26589', normal, compile, ['-O'])
++
  test('T25976', [grep_errmsg('Dead Code')], compile, ['-O -ddump-simpl -dsuppress-uniques -dno-typeable-binds'])
  test('T25965', normal, compile, ['-O'])