Eta expand through CallStacks

This patch fixes #20103, by treating HasCallStack constraints as
cheap when eta-expanding.

See Note [Eta expanding through CallStacks] in GHC.Core.Opt.Arity

6 files changed, 217 insertions, 24 deletions

diff --git a/compiler/GHC/Core/Opt/Arity.hs b/compiler/GHC/Core/Opt/Arity.hs
index e5e63aca26..b2476d39f0 100644
--- a/compiler/GHC/Core/Opt/Arity.hs
+++ b/compiler/GHC/Core/Opt/Arity.hs
@@ -39,7 +39,7 @@ import GHC.Core.Utils
 import GHC.Core.DataCon
 import GHC.Core.TyCon     ( tyConArity )
 import GHC.Core.TyCon.RecWalk     ( initRecTc, checkRecTc )
-import GHC.Core.Predicate ( isDictTy )
+import GHC.Core.Predicate ( isDictTy, isCallStackPredTy )
 import GHC.Core.Multiplicity
 
 -- We have two sorts of substitution:
@@ -835,6 +835,17 @@ topDiv.
 
 Historical note: long ago, we unconditionally switched to topDiv when we
 encountered a cast, but that is far too conservative: see #5475
+
+Note [Eta expanding through CallStacks]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Just as it's good to eta-expand through dictionaries, so it is good to
+do so through CallStacks.  #20103 is a case in point, where we got
+  foo :: HasCallStack => Int -> Int
+  foo = \(d::CallStack). let d2 = pushCallStack blah d in
+        \(x:Int). blah
+
+We really want to eta-expand this!  #20103 is quite convincing!
+We do this regardless of -fdicts-cheap; it's not really a dictionary.
 -}
 
 ---------------------------
@@ -963,7 +974,13 @@ myExprIsCheap AE{ae_mode = mode} e mb_ty = case mode of
   BotStrictness -> False
   _             -> cheap_dict || cheap_fun e
     where
-      cheap_dict = am_dicts_cheap mode && fmap isDictTy mb_ty == Just True
+      cheap_dict = case mb_ty of
+                     Nothing -> False
+                     Just ty -> (am_dicts_cheap mode && isDictTy ty)
+                                || isCallStackPredTy ty
+        -- See Note [Eta expanding through dictionaries]
+        -- See Note [Eta expanding through CallStacks]
+
       cheap_fun e = case mode of
 #if __GLASGOW_HASKELL__ <= 900
         BotStrictness                -> panic "impossible"
diff --git a/compiler/GHC/Core/Predicate.hs b/compiler/GHC/Core/Predicate.hs
index 43f52b9b5c..9601a92138 100644
--- a/compiler/GHC/Core/Predicate.hs
+++ b/compiler/GHC/Core/Predicate.hs
@@ -24,6 +24,7 @@ module GHC.Core.Predicate (
 
   -- Implicit parameters
   isIPLikePred, hasIPSuperClasses, isIPTyCon, isIPClass,
+  isCallStackTy, isCallStackPred, isCallStackPredTy,
 
   -- Evidence variables
   DictId, isEvVar, isDictId
@@ -44,6 +45,7 @@ import GHC.Builtin.Names
 import GHC.Utils.Outputable
 import GHC.Utils.Misc
 import GHC.Utils.Panic
+import GHC.Data.FastString( FastString )
 
 
 -- | A predicate in the solver. The solver tries to prove Wanted predicates
@@ -257,6 +259,39 @@ has_ip_super_classes rec_clss cls tys
 initIPRecTc :: RecTcChecker
 initIPRecTc = setRecTcMaxBound 1 initRecTc
 
+-- --------------------- CallStack predicates ---------------------------------
+
+isCallStackPredTy :: Type -> Bool
+-- True of HasCallStack, or IP "blah" CallStack
+isCallStackPredTy ty
+  | Just (tc, tys) <- splitTyConApp_maybe ty
+  , Just cls <- tyConClass_maybe tc
+  , Just {} <- isCallStackPred cls tys
+  = True
+  | otherwise
+  = False
+
+-- | Is a 'PredType' a 'CallStack' implicit parameter?
+--
+-- If so, return the name of the parameter.
+isCallStackPred :: Class -> [Type] -> Maybe FastString
+isCallStackPred cls tys
+  | [ty1, ty2] <- tys
+  , isIPClass cls
+  , isCallStackTy ty2
+  = isStrLitTy ty1
+  | otherwise
+  = Nothing
+
+-- | Is a type a 'CallStack'?
+isCallStackTy :: Type -> Bool
+isCallStackTy ty
+  | Just tc <- tyConAppTyCon_maybe ty
+  = tc `hasKey` callStackTyConKey
+  | otherwise
+  = False
+
+
 {- Note [Local implicit parameters]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 The function isIPLikePred tells if this predicate, or any of its
diff --git a/compiler/GHC/Tc/Utils/TcType.hs b/compiler/GHC/Tc/Utils/TcType.hs
index 9a67143892..367922e3e5 100644
--- a/compiler/GHC/Tc/Utils/TcType.hs
+++ b/compiler/GHC/Tc/Utils/TcType.hs
@@ -80,7 +80,7 @@ module GHC.Tc.Utils.TcType (
   isSigmaTy, isRhoTy, isRhoExpTy, isOverloadedTy,
   isFloatingTy, isDoubleTy, isFloatTy, isIntTy, isWordTy, isStringTy,
   isIntegerTy, isNaturalTy,
-  isBoolTy, isUnitTy, isCharTy, isCallStackTy, isCallStackPred,
+  isBoolTy, isUnitTy, isCharTy,
   isTauTy, isTauTyCon, tcIsTyVarTy, tcIsForAllTy,
   isPredTy, isTyVarClassPred,
   checkValidClsArgs, hasTyVarHead,
@@ -228,7 +228,6 @@ import GHC.Data.List.SetOps ( getNth, findDupsEq )
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Utils.Panic.Plain
-import GHC.Data.FastString
 import GHC.Utils.Error( Validity(..), isValid )
 import qualified GHC.LanguageExtensions as LangExt
 
@@ -2120,26 +2119,6 @@ isStringTy ty
       Just (tc, [arg_ty]) -> tc == listTyCon && isCharTy arg_ty
       _                   -> False
 
--- | Is a type a 'CallStack'?
-isCallStackTy :: Type -> Bool
-isCallStackTy ty
-  | Just tc <- tyConAppTyCon_maybe ty
-  = tc `hasKey` callStackTyConKey
-  | otherwise
-  = False
-
--- | Is a 'PredType' a 'CallStack' implicit parameter?
---
--- If so, return the name of the parameter.
-isCallStackPred :: Class -> [Type] -> Maybe FastString
-isCallStackPred cls tys
-  | [ty1, ty2] <- tys
-  , isIPClass cls
-  , isCallStackTy ty2
-  = isStrLitTy ty1
-  | otherwise
-  = Nothing
-
 is_tc :: Unique -> Type -> Bool
 -- Newtypes are opaque to this
 is_tc uniq ty = case tcSplitTyConApp_maybe ty of
diff --git a/testsuite/tests/simplCore/should_compile/T20103.hs b/testsuite/tests/simplCore/should_compile/T20103.hs
new file mode 100644
index 0000000000..e19ca5d127
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T20103.hs
@@ -0,0 +1,8 @@
+module T20103 where
+
+import GHC.Stack
+
+foo :: HasCallStack => Int -> Int
+foo 0 = length . fst . head $ getCallStack callStack
+foo n = foo (n-1)
+
diff --git a/testsuite/tests/simplCore/should_compile/T20103.stderr b/testsuite/tests/simplCore/should_compile/T20103.stderr
new file mode 100644
index 0000000000..251a3e60fa
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T20103.stderr
@@ -0,0 +1,153 @@
+
+==================== Tidy Core ====================
+Result size of Tidy Core
+  = {terms: 98, types: 59, coercions: 14, joins: 0/0}
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl :: Int
+[GblId, Unf=OtherCon []]
+lvl = GHC.Types.I# 12#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl1 :: Int
+[GblId, Unf=OtherCon []]
+lvl1 = GHC.Types.I# 7#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl2 :: Int
+[GblId, Unf=OtherCon []]
+lvl2 = GHC.Types.I# 9#
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+lvl3 :: GHC.Prim.Addr#
+[GblId, Unf=OtherCon []]
+lvl3 = "T20103.hs"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl4 :: [Char]
+[GblId]
+lvl4 = GHC.CString.unpackCString# lvl3
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T20103.$trModule2 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 30 0}]
+T20103.$trModule2 = "T20103"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl5 :: [Char]
+[GblId]
+lvl5 = GHC.CString.unpackCString# T20103.$trModule2
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T20103.$trModule4 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 20 0}]
+T20103.$trModule4 = "main"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl6 :: [Char]
+[GblId]
+lvl6 = GHC.CString.unpackCString# T20103.$trModule4
+
+-- RHS size: {terms: 8, types: 0, coercions: 0, joins: 0/0}
+lvl7 :: SrcLoc
+[GblId, Unf=OtherCon []]
+lvl7 = GHC.Stack.Types.SrcLoc lvl6 lvl5 lvl4 lvl1 lvl2 lvl1 lvl
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+lvl8 :: GHC.Prim.Addr#
+[GblId, Unf=OtherCon []]
+lvl8 = "foo"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl9 :: [Char]
+[GblId]
+lvl9 = GHC.CString.unpackCString# lvl8
+
+Rec {
+-- RHS size: {terms: 36, types: 29, coercions: 14, joins: 0/0}
+T20103.$wfoo [InlPrag=[2], Occ=LoopBreaker]
+  :: HasCallStack => GHC.Prim.Int# -> GHC.Prim.Int#
+[GblId, Arity=2, Str=<1L><1L>, Unf=OtherCon []]
+T20103.$wfoo
+  = \ ($dIP :: HasCallStack) (ww :: GHC.Prim.Int#) ->
+      case ww of ds {
+        __DEFAULT ->
+          case $dIP
+               `cast` (GHC.Classes.N:IP[0] <"callStack">_N <CallStack>_N
+                       :: (?callStack::CallStack) ~R# CallStack)
+          of wild1 {
+            __DEFAULT ->
+              T20103.$wfoo
+                ((GHC.Stack.Types.PushCallStack lvl9 lvl7 wild1)
+                 `cast` (Sym (GHC.Classes.N:IP[0] <"callStack">_N <CallStack>_N)
+                         :: CallStack ~R# (?callStack::CallStack)))
+                (GHC.Prim.-# ds 1#);
+            GHC.Stack.Types.FreezeCallStack ds1 ->
+              T20103.$wfoo
+                (wild1
+                 `cast` (Sym (GHC.Classes.N:IP[0] <"callStack">_N <CallStack>_N)
+                         :: CallStack ~R# (?callStack::CallStack)))
+                (GHC.Prim.-# ds 1#)
+          };
+        0# ->
+          case getCallStack
+                 ($dIP
+                  `cast` (GHC.Classes.N:IP[0] <"callStack">_N <CallStack>_N
+                          :: (?callStack::CallStack) ~R# CallStack))
+          of {
+            [] -> case GHC.List.badHead of wild1 { };
+            : x ds1 -> case x of { (x1, ds2) -> GHC.List.$wlenAcc @Char x1 0# }
+          }
+      }
+end Rec }
+
+-- RHS size: {terms: 12, types: 5, coercions: 0, joins: 0/0}
+foo [InlPrag=[2]] :: HasCallStack => Int -> Int
+[GblId,
+ Arity=2,
+ Str=<1L><1P(1L)>,
+ Cpr=1,
+ Unf=Unf{Src=InlineStable, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True,
+         Guidance=ALWAYS_IF(arity=2,unsat_ok=True,boring_ok=False)
+         Tmpl= \ ($dIP [Occ=Once1] :: HasCallStack)
+                 (eta [Occ=Once1!] :: Int) ->
+                 case eta of { GHC.Types.I# ww [Occ=Once1] ->
+                 case T20103.$wfoo $dIP ww of ww1 [Occ=Once1] { __DEFAULT ->
+                 GHC.Types.I# ww1
+                 }
+                 }}]
+foo
+  = \ ($dIP :: HasCallStack) (eta :: Int) ->
+      case eta of { GHC.Types.I# ww ->
+      case T20103.$wfoo $dIP ww of ww1 { __DEFAULT -> GHC.Types.I# ww1 }
+      }
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T20103.$trModule3 :: GHC.Types.TrName
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+T20103.$trModule3 = GHC.Types.TrNameS T20103.$trModule4
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T20103.$trModule1 :: GHC.Types.TrName
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+T20103.$trModule1 = GHC.Types.TrNameS T20103.$trModule2
+
+-- RHS size: {terms: 3, types: 0, coercions: 0, joins: 0/0}
+T20103.$trModule :: GHC.Types.Module
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+T20103.$trModule
+  = GHC.Types.Module T20103.$trModule3 T20103.$trModule1
+
+
+
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 8bbf16627a..2fe050e1ba 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -363,6 +363,7 @@ test('T13873',  [ grep_errmsg(r'SPEC') ], compile, ['-O -ddump-rules'])
 # Look for a specialisation rule for wimwam
 test('T19672', normal, compile, ['-O2 -ddump-rules'])
 
+test('T20103',  [ grep_errmsg(r'Arity') ], compile, ['-O -ddump-simpl -dsuppress-uniques'])
 test('T19780', normal, compile, ['-O2'])
 test('T19794', normal, compile, ['-O'])
 test('T19890',  [ grep_errmsg(r'= T19890.foo1') ], compile, ['-O -ddump-simpl'])