From 2598d72e85d3064f7c03cf321a4e40a3452b5af9 Mon Sep 17 00:00:00 2001
From: Sebastian Graf <sebastian.graf@kit.edu>
Date: Tue, 24 Aug 2021 15:37:19 +0200
Subject: SetLevels: Box expressions returning uboxed tuples so that they float

---
 compiler/GHC/Core/Make.hs          | 13 +++++++++-
 compiler/GHC/Core/Opt/SetLevels.hs | 50 +++++++++++++++++++++++++++++---------
 2 files changed, 51 insertions(+), 12 deletions(-)

diff --git a/compiler/GHC/Core/Make.hs b/compiler/GHC/Core/Make.hs
index d7a78b5888..f1338db907 100644
--- a/compiler/GHC/Core/Make.hs
+++ b/compiler/GHC/Core/Make.hs
@@ -24,7 +24,7 @@ module GHC.Core.Make (
         FloatBind(..), wrapFloat, wrapFloats, floatBindings,
 
         -- * Constructing small tuples
-        mkCoreVarTupTy, mkCoreTup, mkCoreUbxTup, mkCoreUbxSum,
+        mkCoreVarTupTy, mkCoreTup, mkCoreUbxTup, mkUbxTup, mkCoreUbxSum,
         mkCoreTupBoxity, unitExpr,
 
         -- * Constructing big tuples
@@ -431,6 +431,17 @@ mkCoreUbxTup tys exps
     mkCoreConApps (tupleDataCon Unboxed (length tys))
              (map (Type . getRuntimeRep) tys ++ map Type tys ++ exps)
 
+-- | Build a small unboxed tuple holding the specified binders.
+-- Does /not/ flatten one-tuples; see Note [Flattening one-tuples]
+mkUbxTup :: [Var] -> Expr b
+mkUbxTup arg_ids
+  = mkConApp2 (tupleDataCon Unboxed (length tys))
+              (reps ++ tys)
+              arg_ids
+  where
+    tys  = map idType arg_ids
+    reps = map getRuntimeRep tys
+
 -- | Make a core tuple of the given boxity; don't flatten 1-tuples
 mkCoreTupBoxity :: Boxity -> [CoreExpr] -> CoreExpr
 mkCoreTupBoxity Boxed   exps = mkCoreTup1 exps
diff --git a/compiler/GHC/Core/Opt/SetLevels.hs b/compiler/GHC/Core/Opt/SetLevels.hs
index 2d69e8eb04..a00e2272df 100644
--- a/compiler/GHC/Core/Opt/SetLevels.hs
+++ b/compiler/GHC/Core/Opt/SetLevels.hs
@@ -79,21 +79,16 @@ import GHC.Prelude
 
 import GHC.Core
 import GHC.Core.Opt.Monad ( FloatOutSwitches(..) )
-import GHC.Core.Utils   ( exprType, exprIsHNF
-                        , exprOkForSpeculation
-                        , exprIsTopLevelBindable
-                        , isExprLevPoly
-                        , collectMakeStaticArgs
-                        , mkLamTypes
-                        )
+import GHC.Core.Utils
 import GHC.Core.Opt.Arity   ( exprBotStrictness_maybe )
 import GHC.Core.FVs     -- all of it
 import GHC.Core.Subst
-import GHC.Core.Make    ( sortQuantVars )
+import GHC.Core.DataCon
+import GHC.Core.Make    ( sortQuantVars, mkUbxTup )
 import GHC.Core.Type    ( Type, splitTyConApp_maybe, tyCoVarsOfType
                         , mightBeUnliftedType, closeOverKindsDSet )
-import GHC.Core.Multiplicity     ( pattern Many )
-import GHC.Core.DataCon ( dataConOrigResTy )
+import GHC.Core.TyCon
+import GHC.Core.Multiplicity     ( pattern Many, scaledThing )
 
 import GHC.Types.Id
 import GHC.Types.Id.Info
@@ -111,7 +106,7 @@ import GHC.Types.Unique       ( hasKey )
 import GHC.Types.Tickish      ( tickishIsCode )
 import GHC.Types.Unique.Supply
 import GHC.Types.Unique.DFM
-import GHC.Types.Basic  ( Arity, RecFlag(..), isRec )
+import GHC.Types.Basic  ( Arity, Boxity(..), RecFlag(..), isRec )
 
 import GHC.Builtin.Types
 import GHC.Builtin.Names      ( runRWKey )
@@ -714,6 +709,39 @@ lvlMFE env strict_ctxt ann_expr
        ; return (Let (NonRec (TB var (FloatMe dest_lvl)) float_rhs)
                      use_expr) }
 
+  | escapes_value_lam
+  , not expr_ok_for_spec -- Boxing/unboxing isn't worth it for cheap expressions
+                         -- See Note [Test cheapness with exprOkForSpeculation]
+  , Just (tc, tc_args) <- splitTyConApp_maybe expr_ty
+  , isUnboxedTupleTyCon tc
+  , let ubx_dc  = tyConSingleDataCon tc
+        arity   = dataConRepArity ubx_dc
+        arg_tys = scaledThing <$> dataConInstArgTys ubx_dc tc_args
+        bx_dc   = tupleDataCon Boxed arity
+  , not $ any mightBeUnliftedType arg_tys -- for now
+  = do { expr1 <- lvlExpr rhs_env ann_expr
+       ; let bx_expr = mkConApp2 bx_dc arg_tys fld_bndrs
+             bx_ty   = exprType bx_expr
+             (bx_bndr:ubx_bndr:fld_bndrs) = mkTemplateLocals (bx_ty:expr_ty:arg_tys)
+             l1r       = incMinorLvlFrom rhs_env
+
+             float_rhs = mkLams abs_vars_w_lvls $
+                         Case expr1 (stayPut l1r ubx_bndr) bx_ty
+                             [Alt (DataAlt ubx_dc)
+                                  (map (stayPut l1r) fld_bndrs)
+                                  (mkConApp2 bx_dc arg_tys fld_bndrs)]
+
+       ; var <- newLvlVar float_rhs Nothing is_mk_static
+       ; let l1u      = incMinorLvlFrom env
+             use_expr = Case (mkVarApps (Var var) abs_vars)
+                             (stayPut l1u bx_bndr) expr_ty
+                             [Alt (DataAlt bx_dc)
+                                  (map (stayPut l1u) fld_bndrs)
+                                  (mkUbxTup fld_bndrs)]
+
+       ; return (Let (NonRec (TB var (FloatMe dest_lvl)) float_rhs)
+                     use_expr) }
+
   | otherwise          -- e.g. do not float unboxed tuples
   = lvlExpr env ann_expr
 
-- 
cgit v1.2.1