Look through newtype wrappers (Trac #16254)

exprIsConApp_maybe could detect that I# 10 is a constructor application,
but not that Size (I# 10) is, because it was an application with a
nontrivial argument.

7 files changed, 81 insertions, 8 deletions

diff --git a/compiler/basicTypes/Id.hs b/compiler/basicTypes/Id.hs
index 01b648ee89..199842ceb1 100644
--- a/compiler/basicTypes/Id.hs
+++ b/compiler/basicTypes/Id.hs
@@ -66,7 +66,9 @@ module Id (
         isClassOpId_maybe, isDFunId,
         isPrimOpId, isPrimOpId_maybe,
         isFCallId, isFCallId_maybe,
-        isDataConWorkId, isDataConWorkId_maybe, isDataConWrapId, isDataConId_maybe,
+        isDataConWorkId, isDataConWorkId_maybe,
+        isDataConWrapId, isDataConWrapId_maybe,
+        isDataConId_maybe,
         idDataCon,
         isConLikeId, isBottomingId, idIsFrom,
         hasNoBinding,
@@ -427,6 +429,7 @@ isClassOpId_maybe       :: Id -> Maybe Class
 isPrimOpId_maybe        :: Id -> Maybe PrimOp
 isFCallId_maybe         :: Id -> Maybe ForeignCall
 isDataConWorkId_maybe   :: Id -> Maybe DataCon
+isDataConWrapId_maybe   :: Id -> Maybe DataCon
 
 isRecordSelector id = case Var.idDetails id of
                         RecSelId {}     -> True
@@ -480,6 +483,10 @@ isDataConWrapId id = case Var.idDetails id of
                        DataConWrapId _ -> True
                        _               -> False
 
+isDataConWrapId_maybe id = case Var.idDetails id of
+                        DataConWrapId con -> Just con
+                        _                 -> Nothing
+
 isDataConId_maybe :: Id -> Maybe DataCon
 isDataConId_maybe id = case Var.idDetails id of
                          DataConWorkId con -> Just con
diff --git a/compiler/coreSyn/CoreOpt.hs b/compiler/coreSyn/CoreOpt.hs
index dc74acf8f0..548b5de269 100644
--- a/compiler/coreSyn/CoreOpt.hs
+++ b/compiler/coreSyn/CoreOpt.hs
@@ -28,7 +28,7 @@ import CoreSyn
 import CoreSubst
 import CoreUtils
 import CoreFVs
-import MkCore ( FloatBind(..) )
+import MkCore ( FloatBind(..), mkCoreLet )
 import PprCore  ( pprCoreBindings, pprRules )
 import OccurAnal( occurAnalyseExpr, occurAnalysePgm )
 import Literal  ( Literal(LitString) )
@@ -42,7 +42,7 @@ import OptCoercion ( optCoercion )
 import Type     hiding ( substTy, extendTvSubst, extendCvSubst, extendTvSubstList
                        , isInScope, substTyVarBndr, cloneTyVarBndr )
 import Coercion hiding ( substCo, substCoVarBndr )
-import TyCon        ( tyConArity )
+import TyCon        ( tyConArity, isNewTyCon )
 import TysWiredIn
 import PrelNames
 import BasicTypes
@@ -783,7 +783,7 @@ Here's how exprIsConApp_maybe achieves this:
       scrutinee = (\n. case n of n' -> MkT n') e
 
 2.  Beta-reduce the application, generating a floated 'let'.
-    See Note [beta-reduction in exprIsConApp_maybe] below.  Now we have
+    See Note [Special case for newtype wrappers] below.  Now we have
       scrutinee = case n of n' -> MkT n'
       with floats {Let n = e}
 
@@ -796,9 +796,8 @@ And now we have a known-constructor MkT that we can return.
 Notice that both (2) and (3) require exprIsConApp_maybe to gather and return
 a bunch of floats, both let and case bindings.
 
-Note [beta-reduction in exprIsConApp_maybe]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
+Note [Special case for newtype wrappers]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 The unfolding a definition (_e.g._ a let-bound variable or a datacon wrapper) is
 typically a function. For instance, take the wrapper for MkT in Note
 [exprIsConApp_maybe on data constructors with wrappers]:
@@ -829,6 +828,40 @@ Is transformed into
 Which, effectively, means emitting a float `let x = arg` and recursively
 analysing the body.
 
+This strategy requires a special case for newtypes. Suppose we have
+   newtype T a b where
+     MkT :: a -> T b a   -- Note args swapped
+
+This defines a worker function MkT, a wrapper function $WMkT, and an axT:
+   $WMkT :: forall a b. a -> T b a
+   $WMkT = /\b a. \(x:a). MkT a b x    -- A real binding
+
+   MkT :: forall a b. a -> T a b
+   MkT = /\a b. \(x:a). x |> (ax a b)  -- A compulsory unfolding
+
+   axiom axT :: a ~R# T a b
+
+Now we are optimising
+   case $WMkT (I# 3) |> sym axT of I# y -> ...
+we clearly want to simplify this.  The danger is that we'll end up with
+   let a = I#3 in case a of I# y -> ...
+because in general, we do this on-the-fly beta-reduction
+   (\x. e) blah  -->  let x = blah in e
+and then float the the let.  (Substitution would risk duplicating 'blah'.)
+
+But if the case-of-known-constructor doesn't actually fire (i.e.
+exprIsConApp_maybe does not return Just) then nothing happens, and nothing
+will happen the next time either.
+
+For newtype wrappers we know for sure that the argument of the beta-redex
+is used exactly once, so we can substitute aggressively rather than use a let.
+Hence the special case, implemented in dealWithNewtypeWrapper.
+(It's sound for any beta-redex where the argument is used once, of course.)
+
+dealWithNewtypeWrapper is recursive since newtypes can have
+multiple type arguments.
+
+See test T16254, which checks the behavior of newtypes.
 -}
 
 data ConCont = CC [CoreExpr] Coercion
@@ -871,7 +904,9 @@ exprIsConApp_maybe (in_scope, id_unf) expr
     go subst floats (Lam var body) (CC (arg:args) co)
        | exprIsTrivial arg          -- Don't duplicate stuff!
        = go (extend subst var arg) floats body (CC args co)
-    go subst floats (Let bndr@(NonRec b _) expr) cont
+    go subst floats (Lam var body) (CC (arg:args) co)
+       = go subst floats (mkCoreLet (NonRec var arg) body) (CC args co)
+    go subst floats (Let bndr@(NonRec _ _) expr) cont
        = let (subst', bndr') = subst_bind subst bndr in
            go subst' (FloatLet bndr' : floats) expr cont
     go subst floats (Case scrut b _ [(con, vars, expr)]) cont
@@ -892,6 +927,12 @@ exprIsConApp_maybe (in_scope, id_unf) expr
         , count isValArg args == idArity fun
         = pushFloats floats $ pushCoDataCon con args co
 
+        -- See Note [Special case for newtype wrappers]
+        | Just a <- isDataConWrapId_maybe fun
+        , isNewTyCon (dataConTyCon a)
+        , let rhs = uf_tmpl (realIdUnfolding fun)
+        = dealWithNewtypeWrapper (Left in_scope) floats rhs cont
+
         -- Look through data constructor wrappers: they inline late (See Note
         -- [Activation for data constructor wrappers]) but we want to do
         -- case-of-known-constructor optimisation eagerly.
@@ -932,6 +973,9 @@ exprIsConApp_maybe (in_scope, id_unf) expr
       (c, tys, args) <- x
       return (floats, c, tys, args)
 
+    dealWithNewtypeWrapper scope floats (Lam v body) (CC (arg:args) co) =
+      dealWithNewtypeWrapper (extend scope v arg) floats body (CC args co)
+    dealWithNewtypeWrapper scope floats expr args = go scope floats expr args
     ----------------------------
     -- Operations on the (Either InScopeSet CoreSubst)
     -- The Left case is wildly dominant
diff --git a/compiler/coreSyn/CoreUtils.hs b/compiler/coreSyn/CoreUtils.hs
index 9c425e72f0..4602dfa065 100644
--- a/compiler/coreSyn/CoreUtils.hs
+++ b/compiler/coreSyn/CoreUtils.hs
@@ -1360,6 +1360,7 @@ isExpandableApp fn n_val_args
   | otherwise
   = case idDetails fn of
       DataConWorkId {} -> True  -- Actually handled by isWorkFreeApp
+      DataConWrapId {} -> True  -- See Note [Special case for newtype wrappers]
       RecSelId {}      -> n_val_args == 1  -- See Note [Record selection]
       ClassOpId {}     -> n_val_args == 1
       PrimOpId {}      -> False
diff --git a/testsuite/tests/simplCore/should_compile/Makefile b/testsuite/tests/simplCore/should_compile/Makefile
index 277a5a664b..8577dea2a7 100644
--- a/testsuite/tests/simplCore/should_compile/Makefile
+++ b/testsuite/tests/simplCore/should_compile/Makefile
@@ -139,6 +139,11 @@ T5327:
 	$(RM) -f T5327.hi T5327.o
 	'$(TEST_HC)' $(TEST_HC_OPTS) -c T5327.hs -O -ddump-simpl | grep -c '># 34# '
 
+.PHONY: T16254
+T16254:
+	$(RM) -f T16254.hi T16254.o
+	'$(TEST_HC)' $(TEST_HC_OPTS) -c T16254.hs -O -ddump-simpl | grep -c '># 34# '
+
 .PHONY: T5623
 T5623:
 	$(RM) -f T5623.hi T5623.o
diff --git a/testsuite/tests/simplCore/should_compile/T16254.hs b/testsuite/tests/simplCore/should_compile/T16254.hs
new file mode 100644
index 0000000000..3c1490c17c
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T16254.hs
@@ -0,0 +1,14 @@
+-- variant of T5327, where we force the newtype to have a wrapper
+{-# LANGUAGE GADTs, ExplicitForAll #-}
+module T16254 where
+
+newtype Size a b where
+  Size :: forall b a. Int -> Size a b
+
+{-# INLINABLE val2 #-}
+val2 = Size 17
+
+-- In the core, we should see a comparison against 34#, i.e. constant
+-- folding should have happened. We actually see it twice: Once in f's
+-- definition, and once in its unfolding.
+f n = case val2 of Size s -> s + s > n
diff --git a/testsuite/tests/simplCore/should_compile/T16254.stdout b/testsuite/tests/simplCore/should_compile/T16254.stdout
new file mode 100644
index 0000000000..0cfbf08886
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T16254.stdout
@@ -0,0 +1 @@
+2
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 779b09175e..6e1979c5e6 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -113,6 +113,7 @@ test('T5359b', normal, compile, [''])  # Lint error with -O (OccurAnal)
 test('T5458', normal, compile, [''])
 test('simpl021', [extra_files(['Simpl021A.hs', 'Simpl021B.hs'])], makefile_test, ['simpl021'])
 test('T5327', normal, makefile_test, ['T5327'])
+test('T16254', normal, makefile_test, ['T16254'])
 test('T5615', normal, makefile_test, ['T5615'])
 test('T5623', normal, makefile_test, ['T5623'])
 test('T13155', normal, makefile_test, ['T13155'])