From 9187d5fb1d3d38a4e607b0d61784c21447c8195b Mon Sep 17 00:00:00 2001
From: Simon Peyton Jones <simonpj@microsoft.com>
Date: Mon, 2 Apr 2018 14:55:43 +0100
Subject: [PATCH] Allow unpacking of single-data-con GADTs

Trac #14978 pointed out that single-constructor GADTs should be
unpackable without trouble.

Acutally I realise that even existentials should be unpackable
too, but that's a bit more work, so it's not part of this patch.

See Note [Unpacking GADTs and existentials] in MkId.
---
 compiler/basicTypes/MkId.hs                   | 59 ++++++++++++-------
 .../tests/simplCore/should_compile/Makefile   |  5 ++
 .../tests/simplCore/should_compile/T14978.hs  | 12 ++++
 .../simplCore/should_compile/T14978.stdout    |  2 +
 .../tests/simplCore/should_compile/all.T      |  4 ++
 5 files changed, 62 insertions(+), 20 deletions(-)
 create mode 100644 testsuite/tests/simplCore/should_compile/T14978.hs
 create mode 100644 testsuite/tests/simplCore/should_compile/T14978.stdout

diff --git a/compiler/basicTypes/MkId.hs b/compiler/basicTypes/MkId.hs
index 38c772c935..27e9f2bd29 100644
--- a/compiler/basicTypes/MkId.hs
+++ b/compiler/basicTypes/MkId.hs
@@ -855,7 +855,7 @@ dataConArgUnpack arg_ty
       -- A recursive newtype might mean that
       -- 'arg_ty' is a newtype
   , let rep_tys = dataConInstArgTys con tc_args
-  = ASSERT( isVanillaDataCon con )
+  = ASSERT( null (dataConExTyVars con) )  -- Note [Unpacking GADTs and existentials]
     ( rep_tys `zip` dataConRepStrictness con
     ,( \ arg_id ->
        do { rep_ids <- mapM newLocal rep_tys
@@ -879,31 +879,27 @@ isUnpackableType :: DynFlags -> FamInstEnvs -> Type -> Bool
 -- we encounter on the way, because otherwise we might well
 -- end up relying on ourselves!
 isUnpackableType dflags fam_envs ty
-  | Just (tc, _) <- splitTyConApp_maybe ty
-  , Just con <- tyConSingleAlgDataCon_maybe tc
-  , isVanillaDataCon con
-  = ok_con_args (unitNameSet (getName tc)) con
+  | Just data_con <- unpackable_type ty
+  = ok_con_args (unitNameSet (getName data_con)) data_con
   | otherwise
   = False
   where
-    ok_arg tcs (ty, bang) = not (attempt_unpack bang) || ok_ty tcs norm_ty
+    ok_con_args dcs con
+       = all (ok_arg dcs) (dataConOrigArgTys con `zip` dataConSrcBangs con)
+         -- NB: dataConSrcBangs gives the *user* request;
+         -- We'd get a black hole if we used dataConImplBangs
+
+    ok_arg dcs (ty, bang) = not (attempt_unpack bang) || ok_ty dcs norm_ty
         where
           norm_ty = topNormaliseType fam_envs ty
-    ok_ty tcs ty
-      | Just (tc, _) <- splitTyConApp_maybe ty
-      , let tc_name = getName tc
-      =  not (tc_name `elemNameSet` tcs)
-      && case tyConSingleAlgDataCon_maybe tc of
-            Just con | isVanillaDataCon con
-                    -> ok_con_args (tcs `extendNameSet` getName tc) con
-            _ -> True
-      | otherwise
-      = True
 
-    ok_con_args tcs con
-       = all (ok_arg tcs) (dataConOrigArgTys con `zip` dataConSrcBangs con)
-         -- NB: dataConSrcBangs gives the *user* request;
-         -- We'd get a black hole if we used dataConImplBangs
+    ok_ty dcs ty
+      | Just data_con <- unpackable_type ty
+      , let dc_name = getName data_con
+      , not (dc_name `elemNameSet` dcs)
+      = ok_con_args (dcs `extendNameSet` dc_name) data_con
+      | otherwise
+      = True        -- NB True here, in constrast to False at top level
 
     attempt_unpack (HsSrcBang _ SrcUnpack NoSrcStrict)
       = xopt LangExt.StrictData dflags
@@ -915,7 +911,30 @@ isUnpackableType dflags fam_envs ty
       = xopt LangExt.StrictData dflags -- Be conservative
     attempt_unpack _ = False
 
+    unpackable_type :: Type -> Maybe DataCon
+    -- Works just on a single level
+    unpackable_type ty
+      | Just (tc, _) <- splitTyConApp_maybe ty
+      , Just data_con <- tyConSingleAlgDataCon_maybe tc
+      , null (dataConExTyVars data_con)  -- See Note [Unpacking GADTs and existentials]
+      = Just data_con
+      | otherwise
+      = Nothing
+
 {-
+Note [Unpacking GADTs and existentials]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+There is nothing stopping us unpacking a data type with equality
+components, like
+  data Equal a b where
+    Equal :: Equal a a
+
+And it'd be fine to unpack a product type with existential components
+too, but that would require a bit more plumbing, so currently we don't.
+
+So for now we require: null (dataConExTyVars data_con)
+See Trac #14978
+
 Note [Unpack one-wide fields]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 The flag UnboxSmallStrictFields ensures that any field that can
diff --git a/testsuite/tests/simplCore/should_compile/Makefile b/testsuite/tests/simplCore/should_compile/Makefile
index 33322f38a1..1f1bf25b69 100644
--- a/testsuite/tests/simplCore/should_compile/Makefile
+++ b/testsuite/tests/simplCore/should_compile/Makefile
@@ -2,6 +2,11 @@ TOP=../../..
 include $(TOP)/mk/boilerplate.mk
 include $(TOP)/mk/test.mk
 
+T14978:
+	$(RM) -f T14978.o T14978.hi
+	-'$(TEST_HC)' $(TEST_HC_OPTS) -O -c -ddump-simpl T14978.hs -dsuppress-coercions | grep 'foo'
+        # Expecting the defn of 'foo' to apply Goof to an unboxed coercion
+
 T13468:
 	$(RM) -f T13468.o T13468.hi
 	-'$(TEST_HC)' $(TEST_HC_OPTS) -O -c -ddump-simpl T13468.hs | grep 'Error'
diff --git a/testsuite/tests/simplCore/should_compile/T14978.hs b/testsuite/tests/simplCore/should_compile/T14978.hs
new file mode 100644
index 0000000000..3fedb960d8
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T14978.hs
@@ -0,0 +1,12 @@
+{-# LANGUAGE GADTs  #-}
+module T14978 where
+
+data Equal a b where
+  Refl :: Equal a a
+
+data Goof a where
+  Goof :: {-# UNPACK #-} !(Equal a Int) -> Goof a
+
+
+foo :: Goof Int
+foo = Goof Refl
diff --git a/testsuite/tests/simplCore/should_compile/T14978.stdout b/testsuite/tests/simplCore/should_compile/T14978.stdout
new file mode 100644
index 0000000000..1b12ead299
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T14978.stdout
@@ -0,0 +1,2 @@
+foo :: Goof Int
+foo = T14978.Goof @ Int @~ <Co:1>
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 13511ee4c0..a521a1084e 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -297,3 +297,7 @@ test('T14152a', [extra_files(['T14152.hs']), pre_cmd('cp T14152.hs T14152a.hs'),
 test('T13990', normal, compile, ['-dcore-lint -O'])
 test('T14650', normal, compile, ['-O2'])
 test('T14959', normal, compile, ['-O'])
+test('T14978',
+     normal,
+     run_command,
+     ['$MAKE -s --no-print-directory T14978'])
-- 
GitLab