Make isTcLevPoly more conservative with newtypes (#17360)

`isTcLevPoly` gives an approximate answer for when a type constructor
is levity polymorphic when fully applied, where `True` means
"possibly levity polymorphic" and `False` means "definitely not
levity polymorphic". `isTcLevPoly` returned `False` for newtypes,
which is incorrect in the presence of `UnliftedNewtypes`, leading
to #17360. This patch tweaks `isTcLevPoly` to return `True` for
newtypes instead.

Fixes #17360.

5 files changed, 41 insertions, 6 deletions

diff --git a/compiler/iface/BuildTyCl.hs b/compiler/iface/BuildTyCl.hs
index dc1c843889..827b89983f 100644
--- a/compiler/iface/BuildTyCl.hs
+++ b/compiler/iface/BuildTyCl.hs
@@ -54,12 +54,14 @@ mkNewTyConRhs tycon_name tycon con
         ; return (NewTyCon { data_con    = con,
                              nt_rhs      = rhs_ty,
                              nt_etad_rhs = (etad_tvs, etad_rhs),
-                             nt_co       = nt_ax } ) }
+                             nt_co       = nt_ax,
+                             nt_lev_poly = isKindLevPoly res_kind } ) }
                              -- Coreview looks through newtypes with a Nothing
                              -- for nt_co, or uses explicit coercions otherwise
   where
-    tvs    = tyConTyVars tycon
-    roles  = tyConRoles tycon
+    tvs      = tyConTyVars tycon
+    roles    = tyConRoles tycon
+    res_kind = tyConResKind tycon
     con_arg_ty = case dataConRepArgTys con of
                    [arg_ty] -> arg_ty
                    tys -> pprPanic "mkNewTyConRhs" (ppr con <+> ppr tys)
diff --git a/compiler/types/TyCon.hs b/compiler/types/TyCon.hs
index 7af2bc0ad7..31bfddce81 100644
--- a/compiler/types/TyCon.hs
+++ b/compiler/types/TyCon.hs
@@ -992,7 +992,7 @@ data AlgTyConRhs
                         -- shorter than the declared arity of the 'TyCon'.
 
                         -- See Note [Newtype eta]
-        nt_co :: CoAxiom Unbranched
+        nt_co :: CoAxiom Unbranched,
                              -- The axiom coercion that creates the @newtype@
                              -- from the representation 'Type'.
 
@@ -1001,6 +1001,16 @@ data AlgTyConRhs
                              -- See Note [Newtype eta]
                              -- Watch out!  If any newtypes become transparent
                              -- again check #1072.
+        nt_lev_poly :: Bool
+                        -- 'True' if the newtype can be levity polymorphic when
+                        -- fully applied to its arguments, 'False' otherwise.
+                        -- This can only ever be 'True' with UnliftedNewtypes.
+                        --
+                        -- Invariant: nt_lev_poly nt = isTypeLevPoly (nt_rhs nt)
+                        --
+                        -- This is cached to make it cheaper to check if a
+                        -- variable binding is levity polymorphic, as used by
+                        -- isTcLevPoly.
     }
 
 mkSumTyConRhs :: [DataCon] -> AlgTyConRhs
@@ -2241,8 +2251,13 @@ setTcTyConKind tc              _    = pprPanic "setTcTyConKind" (ppr tc)
 -- Precondition: The fully-applied TyCon has kind (TYPE blah)
 isTcLevPoly :: TyCon -> Bool
 isTcLevPoly FunTyCon{}           = False
-isTcLevPoly (AlgTyCon { algTcParent = UnboxedAlgTyCon _ }) = True
-isTcLevPoly AlgTyCon{}           = False
+isTcLevPoly (AlgTyCon { algTcParent = parent, algTcRhs = rhs })
+  | UnboxedAlgTyCon _ <- parent
+  = True
+  | NewTyCon { nt_lev_poly = lev_poly } <- rhs
+  = lev_poly -- Newtypes can be levity polymorphic with UnliftedNewtypes (#17360)
+  | otherwise
+  = False
 isTcLevPoly SynonymTyCon{}       = True
 isTcLevPoly FamilyTyCon{}        = True
 isTcLevPoly PrimTyCon{}          = False
diff --git a/testsuite/tests/typecheck/should_fail/T17360.hs b/testsuite/tests/typecheck/should_fail/T17360.hs
new file mode 100644
index 0000000000..fc82675271
--- /dev/null
+++ b/testsuite/tests/typecheck/should_fail/T17360.hs
@@ -0,0 +1,11 @@
+{-# LANGUAGE UnliftedNewtypes #-}
+{-# LANGUAGE ExplicitForAll #-}
+{-# LANGUAGE PolyKinds #-}
+module T17360 where
+
+import GHC.Exts
+
+newtype Id (a :: TYPE r) = Id a
+
+foo :: forall r (a :: TYPE r). Id a -> Id a
+foo x = x
diff --git a/testsuite/tests/typecheck/should_fail/T17360.stderr b/testsuite/tests/typecheck/should_fail/T17360.stderr
new file mode 100644
index 0000000000..9eccd7fe29
--- /dev/null
+++ b/testsuite/tests/typecheck/should_fail/T17360.stderr
@@ -0,0 +1,6 @@
+
+T17360.hs:11:5: error:
+    A levity-polymorphic type is not allowed here:
+      Type: Id a
+      Kind: TYPE r
+    In the type of binder ‘x’
diff --git a/testsuite/tests/typecheck/should_fail/all.T b/testsuite/tests/typecheck/should_fail/all.T
index 3f8dfc0abe..75c1fe6f19 100644
--- a/testsuite/tests/typecheck/should_fail/all.T
+++ b/testsuite/tests/typecheck/should_fail/all.T
@@ -547,3 +547,4 @@ test('T16512a', normal, compile_fail, [''])
 test('T16512b', normal, compile_fail, [''])
 test('T17213', [extra_files(['T17213a.hs'])], multimod_compile_fail, ['T17213', '-v0'])
 test('T17355', normal, compile_fail, [''])
+test('T17360', normal, compile_fail, [''])