Add -funbox-strict-primitive-fields

When enabled, this flag causes all strict fields which representation is
smaller or equal to the size of a pointer to be unboxed.

4 files changed, 87 insertions, 1 deletions

diff --git a/compiler/basicTypes/DataCon.lhs-boot b/compiler/basicTypes/DataCon.lhs-boot
index 94bf889325..716dc7e547 100644
--- a/compiler/basicTypes/DataCon.lhs-boot
+++ b/compiler/basicTypes/DataCon.lhs-boot
@@ -2,9 +2,11 @@
 module DataCon where
 import Name( Name )
 import {-# SOURCE #-} TyCon( TyCon )
+import {-# SOURCE #-} TypeRep (Type)
 
 data DataCon
 dataConName      :: DataCon -> Name
+dataConRepArgTys :: DataCon -> [Type]
 dataConTyCon     :: DataCon -> TyCon
 isVanillaDataCon :: DataCon -> Bool
 instance Eq DataCon
diff --git a/compiler/main/DynFlags.hs b/compiler/main/DynFlags.hs
index 8686e553d0..1c47d6dfad 100644
--- a/compiler/main/DynFlags.hs
+++ b/compiler/main/DynFlags.hs
@@ -269,6 +269,7 @@ data GeneralFlag
    | Opt_DoEtaReduction
    | Opt_CaseMerge
    | Opt_UnboxStrictFields
+   | Opt_UnboxStrictPrimitiveFields
    | Opt_DictsCheap
    | Opt_EnableRewriteRules             -- Apply rewrite rules during simplification
    | Opt_Vectorise
@@ -2359,6 +2360,7 @@ fFlags = [
   ( "do-eta-reduction",                 Opt_DoEtaReduction, nop ),
   ( "case-merge",                       Opt_CaseMerge, nop ),
   ( "unbox-strict-fields",              Opt_UnboxStrictFields, nop ),
+  ( "unbox-strict-primitive-fields",    Opt_UnboxStrictPrimitiveFields, nop ),
   ( "dicts-cheap",                      Opt_DictsCheap, nop ),
   ( "excess-precision",                 Opt_ExcessPrecision, nop ),
   ( "eager-blackholing",                Opt_EagerBlackHoling, nop ),
diff --git a/compiler/typecheck/TcTyClsDecls.lhs b/compiler/typecheck/TcTyClsDecls.lhs
index ffcf5c2991..a7606012ba 100644
--- a/compiler/typecheck/TcTyClsDecls.lhs
+++ b/compiler/typecheck/TcTyClsDecls.lhs
@@ -37,6 +37,7 @@ import TcClassDcl
 import TcHsType
 import TcMType
 import TcType
+import qualified TysPrim
 import TysWiredIn( unitTy )
 import Type
 import Kind
@@ -1208,6 +1209,9 @@ chooseBoxingStrategy dflags arg_ty bang
 	               HsNoBang -> HsNoBang
 	               HsStrict | gopt Opt_UnboxStrictFields dflags
                                 -> can_unbox HsStrict arg_ty
+	                        | gopt Opt_UnboxStrictPrimitiveFields dflags &&
+                                  can_unbox_prim arg_ty
+                                -> HsUnpack
                                 | otherwise -> HsStrict
                        HsNoUnpack -> HsStrict
 	               HsUnpack   -> can_unbox HsUnpackFailed arg_ty
@@ -1234,6 +1238,49 @@ chooseBoxingStrategy dflags arg_ty bang
                  else HsUnpack
 
               | otherwise -> fail_bang
+
+    -- TODO: Deal with type synonyms?
+
+    can_unbox_prim :: TcType -> Bool
+    -- We unpack any field which final unpacked size would be smaller
+    -- or equal to the size of a pointer.
+    can_unbox_prim arg_ty
+       = case splitTyConApp_maybe arg_ty of
+            Nothing -> False
+
+            Just (arg_tycon, _)
+              | isAbstractTyCon arg_tycon -> False
+                      -- See Note [Don't complain about UNPACK on abstract TyCons]
+              | isPrimTyCon arg_tycon &&
+                arg_tycon `elem` ptrSizedPrimTyCons -> True
+              -- TODO: Check that the PrimTyCon corresponds to a type
+              -- with pointer-sized representation.
+              | isEmptyDataTyCon arg_tycon -> True
+              | not (isRecursiveTyCon arg_tycon)        -- Note [Recusive unboxing]
+              , Just ty <- tyConSingleFieldDataCon_maybe arg_tycon
+              -> can_unbox_prim ty
+              | otherwise -> False
+
+ptrSizedPrimTyCons :: [TyCon]
+ptrSizedPrimTyCons =
+    [ TysPrim.addrPrimTyCon
+    , TysPrim.arrayPrimTyCon
+    , TysPrim.byteArrayPrimTyCon
+    , TysPrim.arrayArrayPrimTyCon
+    , TysPrim.charPrimTyCon
+    , TysPrim.doublePrimTyCon
+    , TysPrim.floatPrimTyCon
+    , TysPrim.intPrimTyCon
+    , TysPrim.int32PrimTyCon
+    , TysPrim.int64PrimTyCon
+    , TysPrim.mutableArrayPrimTyCon
+    , TysPrim.mutableByteArrayPrimTyCon
+    , TysPrim.mutableArrayArrayPrimTyCon
+    , TysPrim.wordPrimTyCon
+    , TysPrim.word32PrimTyCon
+    , TysPrim.word64PrimTyCon
+    ]
+
 \end{code}
 
 Note [Don't complain about UNPACK on abstract TyCons]
diff --git a/compiler/types/TyCon.lhs b/compiler/types/TyCon.lhs
index 458f5c6e20..36c52a4ab8 100644
--- a/compiler/types/TyCon.lhs
+++ b/compiler/types/TyCon.lhs
@@ -53,6 +53,7 @@ module TyCon(
         isTyConAssoc, tyConAssoc_maybe,
         isRecursiveTyCon,
         isImplicitTyCon,
+        isEmptyDataTyCon,
 
         -- ** Extracting information out of TyCons
         tyConName,
@@ -72,6 +73,7 @@ module TyCon(
         algTyConRhs,
         newTyConRhs, newTyConEtadRhs, unwrapNewTyCon_maybe,
         tupleTyConBoxity, tupleTyConSort, tupleTyConArity,
+        tyConSingleFieldDataCon_maybe,
 
         -- ** Manipulating TyCons
         tcExpandTyCon_maybe, coreExpandTyCon_maybe,
@@ -88,7 +90,7 @@ module TyCon(
 #include "HsVersions.h"
 
 import {-# SOURCE #-} TypeRep ( Kind, Type, PredType )
-import {-# SOURCE #-} DataCon ( DataCon, isVanillaDataCon )
+import {-# SOURCE #-} DataCon ( DataCon, dataConRepArgTys, isVanillaDataCon )
 
 import Var
 import Class
@@ -1074,6 +1076,18 @@ isDataTyCon (AlgTyCon {algTcRhs = rhs})
 isDataTyCon (TupleTyCon {tyConTupleSort = sort}) = isBoxed (tupleSortBoxity sort)
 isDataTyCon _ = False
 
+isEmptyDataTyCon :: TyCon -> Bool
+isEmptyDataTyCon (AlgTyCon {algTcRhs = DataTyCon { data_cons = [data_con] } })
+    = isEmptyDataCon data_con
+isEmptyDataTyCon (TupleTyCon {dataCon = data_con })
+    = isEmptyDataCon data_con
+isEmptyDataTyCon _ = False
+
+isEmptyDataCon :: DataCon -> Bool
+isEmptyDataCon data_con = case dataConRepArgTys data_con of
+    [] -> True
+    _  -> False
+
 -- | 'isDistinctTyCon' is true of 'TyCon's that are equal only to
 -- themselves, even via coercions (except for unsafeCoerce).
 -- This excludes newtypes, type functions, type synonyms.
@@ -1128,6 +1142,27 @@ isProductTyCon tc@(AlgTyCon {}) = case algTcRhs tc of
 isProductTyCon (TupleTyCon {})  = True
 isProductTyCon _                = False
 
+-- | If the given 'TyCon' has a /single/ data constructor with a /single/ field,
+-- i.e. it is a @data@ type with one alternative and one field, or a @newtype@
+-- then the type of that field is returned. If the 'TyCon' has a single
+-- constructor with more than one field, more than one constructor, or
+-- represents a primitive or function type constructor then @Nothing@ is
+-- returned. In any other case, the function panics
+tyConSingleFieldDataCon_maybe :: TyCon -> Maybe Type
+tyConSingleFieldDataCon_maybe tc@(AlgTyCon {}) = case algTcRhs tc of
+    DataTyCon{ data_cons = [data_con] }
+        | isVanillaDataCon data_con -> case dataConRepArgTys data_con of
+            [ty] -> Just ty
+            _    -> Nothing
+        | otherwise -> Nothing
+    NewTyCon { data_con = data_con }
+        ->  case dataConRepArgTys data_con of
+            [ty] -> Just ty
+            _    -> pprPanic "tyConSingleFieldDataCon_maybe"
+                    (ppr $ dataConRepArgTys data_con)
+    _           -> Nothing
+tyConSingleFieldDataCon_maybe _                = Nothing
+
 -- | Is this a 'TyCon' representing a type synonym (@type@)?
 isSynTyCon :: TyCon -> Bool
 isSynTyCon (SynTyCon {}) = True