Perf tweaks

.gitmodules

@@ -116,5 +116,5 @@
 	ignore = untracked
 [submodule "libraries/haskell-transients"]
 	path = libraries/haskell-transients
-	url = https://github.com/PascEll/haskell-transients.git
+	url = git@github.com:PascEll/haskell-transients.git
 	ignore = untracked

compiler/GHC/Builtin/Names.hs

@@ -497,6 +497,12 @@ basicKnownKeyNames
         , unsafeEqualityTyConName
         , unsafeReflDataConName
         , unsafeCoercePrimName
+
+        -- unsafePersistNode/shareNode rewrite rule
+        , unsafePersistNodeName
+        , shareNodeName
+        , wordMapTyConName
+        , tWordMapTyConName
     ]
 
 genericTyConNames :: [Name]
@@ -644,6 +650,9 @@ gHC_RECORDS = mkBaseModule (fsLit "GHC.Records")
 rOOT_MAIN :: Module
 rOOT_MAIN       = mkMainModule (fsLit ":Main") -- Root module for initialisation
 
+wORD_MAP :: Module
+wORD_MAP = mkModule transientsUnit (mkModuleNameFS (fsLit "WordMap"))
+
 mkInteractiveModule :: Int -> Module
 -- (mkInteractiveMoudule 9) makes module 'interactive:Ghci9'
 mkInteractiveModule n = mkModule interactiveUnit (mkModuleName ("Ghci" ++ show n))
@@ -1102,16 +1111,6 @@ joinMName, alternativeClassName :: Name
 joinMName            = varQual gHC_BASE (fsLit "join")        joinMIdKey
 alternativeClassName = clsQual mONAD (fsLit "Alternative") alternativeClassKey
 
---
-joinMIdKey, apAClassOpKey, pureAClassOpKey, thenAClassOpKey,
-    alternativeClassKey :: Unique
-joinMIdKey          = mkPreludeMiscIdUnique 750
-apAClassOpKey       = mkPreludeMiscIdUnique 751 -- <*>
-pureAClassOpKey     = mkPreludeMiscIdUnique 752
-thenAClassOpKey     = mkPreludeMiscIdUnique 753
-alternativeClassKey = mkPreludeMiscIdUnique 754
-
-
 -- Functions for GHC extensions
 considerAccessibleName :: Name
 considerAccessibleName = varQual gHC_EXTS (fsLit "considerAccessible") considerAccessibleIdKey
@@ -1654,6 +1653,12 @@ fingerprintDataConName :: Name
 fingerprintDataConName =
     dcQual gHC_FINGERPRINT_TYPE (fsLit "Fingerprint") fingerprintDataConKey
 
+shareNodeName, unsafePersistNodeName, wordMapTyConName, tWordMapTyConName :: Name
+shareNodeName = varQual  wORD_MAP (fsLit "shareNode") shareNodeIdKey
+unsafePersistNodeName = varQual  wORD_MAP (fsLit "unsafePersistNode") unsafePersistNodeIdKey
+wordMapTyConName = tcQual  wORD_MAP (fsLit "WordMap") wordMapTyConKey
+tWordMapTyConName = tcQual  wORD_MAP (fsLit "TWordMap") tWordMapTyConKey
+
 {-
 ************************************************************************
 *                                                                      *
@@ -2674,6 +2679,29 @@ mkRationalBase2IdKey, mkRationalBase10IdKey :: Unique
 mkRationalBase2IdKey  = mkPreludeMiscIdUnique 700
 mkRationalBase10IdKey = mkPreludeMiscIdUnique 701 :: Unique
 
+------------------------------------------------------
+-- Applicative/Alternative/Monad combinators 750-754 uniques
+------------------------------------------------------
+
+joinMIdKey, apAClassOpKey, pureAClassOpKey, thenAClassOpKey,
+    alternativeClassKey :: Unique
+joinMIdKey          = mkPreludeMiscIdUnique 750
+apAClassOpKey       = mkPreludeMiscIdUnique 751 -- <*>
+pureAClassOpKey     = mkPreludeMiscIdUnique 752
+thenAClassOpKey     = mkPreludeMiscIdUnique 753
+alternativeClassKey = mkPreludeMiscIdUnique 754
+
+
+------------------------------------------------------
+-- haskell-transients 755-759 uniques
+------------------------------------------------------
+
+shareNodeIdKey, unsafePersistNodeIdKey, wordMapTyConKey, tWordMapTyConKey :: Unique
+shareNodeIdKey  = mkPreludeMiscIdUnique 755
+unsafePersistNodeIdKey  = mkPreludeMiscIdUnique 756
+wordMapTyConKey  = mkPreludeMiscIdUnique 757
+tWordMapTyConKey  = mkPreludeMiscIdUnique 758
+
 {-
 ************************************************************************
 *                                                                      *

compiler/GHC/Core/Opt/Simplify.hs

@@ -53,7 +53,7 @@ import GHC.Types.SourceText
 import GHC.Types.Id
 import GHC.Types.Id.Make   ( seqId )
 import GHC.Types.Id.Info
-import GHC.Types.Name   ( mkSystemVarName, isExternalName, getOccFS )
+import GHC.Types.Name   ( mkSystemVarName, isExternalName, getOccFS, nameModule, nameUnique )
 import GHC.Types.Demand
 import GHC.Types.Cpr    ( mkCprSig, botCpr )
 import GHC.Types.Unique ( hasKey )
@@ -62,7 +62,7 @@ import GHC.Types.Tickish
 import GHC.Types.Var    ( isTyCoVar )
 import GHC.Builtin.PrimOps ( PrimOp (SeqOp) )
 import GHC.Builtin.Types.Prim( realWorldStatePrimTy )
-import GHC.Builtin.Names( runRWKey )
+import GHC.Builtin.Names( runRWKey, unsafePersistNodeIdKey, shareNodeIdKey, unsafePersistNodeName )
 
 import GHC.Data.Maybe   ( isNothing, orElse )
 import GHC.Data.FastString
@@ -77,6 +77,11 @@ import GHC.Utils.Logger
 import GHC.Utils.Misc
 
 import Control.Monad
+import GHC.Unit.Types (moduleUnit)
+import GHC.Unit.State (unitPackageName)
+import GHC.Unit (unitFS)
+import GHC.Builtin.Utils (knownKeyNames)
+import GHC.Core.TyCo.Rep (UnivCoProvenance(CorePrepProv))
 
 {-
 The guts of the simplifier is in this module, but the driver loop for
@@ -2889,13 +2894,82 @@ rebuildCase env scrut case_bndr alts@[Alt _ bndrs rhs] cont
        ; case mb_rule of
            Just (env', rule_rhs, cont') -> simplExprF env' rule_rhs cont'
            Nothing                      -> reallyRebuildCase env scrut case_bndr alts cont }
+
+  -- 2d. Try the builtin unsafePersistNode/shareNode rule. Recall that
+  --
+  -- > shareNode :: forall s a. Node a -> TNode s a
+  -- > unsafePersistNode :: forall s a. TNode s a -> State# s -> (# State# s, Node a #)
+  --
+  -- Now if
+  --     a) the scrutinee is `unsafePersistNode @RealWorld @ty_a tn s`
+  --     b) The case alt looks like `(# s', n #) -> rhs[s', shareNode @RealWorld @ty_a n]`
+  --     c) n's OccInfo tells us that its only OneOcc is in that call to shareNode
+  -- Then rewrite the expr to `rhs[s, tn]`.
+  | Just (env', rhs') <- rewriteTransientPersistent env scrut case_bndr bndrs rhs
+  = simplExprF env' rhs' cont
   where
     all_dead_bndrs = all isDeadBinder bndrs       -- bndrs are [InId]
     is_plain_seq   = all_dead_bndrs && isDeadBinder case_bndr -- Evaluation *only* for effect
 
 rebuildCase env scrut case_bndr alts cont
   = reallyRebuildCase env scrut case_bndr alts cont
 
+-- shareNode :: forall s a. Node a -> TNode s a
+-- unsafePersistNode :: forall s a. TNode s a -> State# s -> (# State# s, Node a #)
+-- case unsafePersistNode @ty_s @ty_a tm s of
+--   (# s', m #) -> E[s', shareNode @ty_s @ty_a m]
+-- ==> { if `shareNode m` is the only occurrence of m and not under a multishot lambda }
+-- E[s, tm]
+rewriteTransientPersistent :: SimplEnv -> OutExpr -> InId -> [InId] -> InExpr -> Maybe (SimplEnv, InExpr)
+rewriteTransientPersistent env scrut case_bndr alt_bndrs rhs
+  | isDeadBinder case_bndr -- Otherwise we might dup the map
+  , [s', m] <- alt_bndrs -- (# s', m #) ->
+--  , pprTrace "rewrite? 1" (ppr scrut) True
+  , OneOcc { occ_in_lam = NotInsideLam } <- idOccInfo m
+--  , pprTrace "rewrite? occ OK 2" empty True
+  , (Var fun_id, [_ty_s, _ty_a, tm, Var s]) <- collectArgs scrut -- unsafePersistNode @ty_s @ty_a (tm :: TWordMap ty_s ty_a) RealWorld
+--  , pprTrace "rewrite? App recogn 3" (ppr fun_id) True
+  , fun_id `hasKey` unsafePersistNodeIdKey
+--  , pprTrace "rewrite? Id matches 4" empty True
+  , let tm_in_id = mkLocalId (idName m) (idMult m) (exprType tm) -- we simply take m's InId Unique as the InId Unique for tm
+  , Just rhs' <- rw m tm_in_id rhs
+--  , pprTrace "rewrite? Yes!" (ppr rhs $$ ppr rhs') True
+  , let env' = extendIdSubst (extendIdSubst env tm_in_id (DoneEx tm Nothing)) s' (DoneId s)
+  = Just (env', rhs') -- substitute s for s'
+  | otherwise
+  = Nothing
+  where
+    rw :: InId -> InId -> InExpr -> Maybe InExpr
+    rw m tm e = go e
+      where
+        go e@App{} = case collectArgs e of
+          (Var fun_id, [_ty_s, _ty_a, Var m'])
+            | fun_id `hasKey` shareNodeIdKey
+            , m == m'
+            , let co = mkUnivCo (CorePrepProv False) Representational (idType tm) (exprType e)
+            , pprTrace "rw: found app" (ppr e $$ ppr co) True
+            -> Just (mkCast (Var tm) co)
+          (f, args)
+            -> mkApps <$> go f <*> traverse go args
+        go (Var v) | v == m = Nothing -- at least one occurrence outside of a unsafePersistNode call ==> abort
+        go (Tick t e) = Tick t <$> go e
+        go (Cast e c) = flip Cast c <$> go e
+        go (Lam b e) = Lam b <$> shadow [b] go e
+        go (Case s b ty alts) = Case <$> go s <*> pure b <*> pure ty <*> traverse (shadow [b] go_alt) alts
+        go (Let b e) = Let <$> go_bind b <*> shadow (bindersOf b) go e
+        go e = Just e
+
+        shadow :: [Var] -> (a -> Maybe a) -> a -> Maybe a
+        shadow bs f a | m `elem` bs = Just a -- NB: Don't need to check tm because it has the same unique as m
+                      | otherwise   = f a
+
+        go_bind (NonRec b e) = NonRec b <$> go e
+        go_bind (Rec pairs)  = (Rec . zip bs) <$> traverse (shadow bs go) rhss
+          where
+            (bs, rhss) = unzip pairs
+
+        go_alt (Alt con bs rhs) = Alt con bs <$> shadow bs go rhs
+
 
 doCaseToLet :: OutExpr          -- Scrutinee
             -> InId             -- Case binder

compiler/GHC/Core/Opt/Simplify/Env.hs

@@ -1141,8 +1141,9 @@ substTyVarBndr env tv
 
 substTyVarBndrT :: TSimplEnv s -> TyVar -> ST s (TSimplEnv s, TyVar)
 substTyVarBndrT env tv = do
-  (TTCvSubst in_scope' tv_env' cv_env', tv') <- Type.substTyVarBndrT (getTTCvSubst env) tv
-  return (env { tseInScope = in_scope', tseTvSubst = tv_env', tseCvSubst = cv_env' }, tv')
+  (TTCvSubst in_scope' tv_env' cv_env', !tv') <- Type.substTyVarBndrT (getTTCvSubst env) tv
+  let !env' = env { tseInScope = in_scope', tseTvSubst = tv_env', tseCvSubst = cv_env' }
+  return (env', tv')
 
 substCoVar :: SimplEnv -> CoVar -> Coercion
 substCoVar env tv = Coercion.substCoVar (getTCvSubst env) tv
@@ -1155,8 +1156,9 @@ substCoVarBndr env cv
 
 substCoVarBndrT :: TSimplEnv s -> CoVar -> ST s (TSimplEnv s, CoVar)
 substCoVarBndrT env cv = do
-  (TTCvSubst in_scope' tv_env' cv_env', cv') <- Coercion.substCoVarBndrT (getTTCvSubst env) cv
-  return (env { tseInScope = in_scope', tseTvSubst = tv_env', tseCvSubst = cv_env' }, cv')
+  (TTCvSubst in_scope' tv_env' cv_env', !cv') <- Coercion.substCoVarBndrT (getTTCvSubst env) cv
+  let !env' = env { tseInScope = in_scope', tseTvSubst = tv_env', tseCvSubst = cv_env' }
+  return (env', cv')
 
 substCo :: SimplEnv -> Coercion -> Coercion
 substCo env co = Coercion.substCo (getTCvSubst env) co

compiler/GHC/Core/TyCo/Subst.hs

@@ -580,19 +580,19 @@ substitution.
 substTyWith :: HasDebugCallStack => [TyVar] -> [Type] -> Type -> Type
 -- Works only if the domain of the substitution is a
 -- superset of the type being substituted into
-substTyWith tvs tys = {-#SCC "substTyWith" #-}
-                      assert (tvs `equalLength` tys )
-                      substTy (zipTvSubst tvs tys)
+substTyWith tvs tys ty = {-#SCC "substTyWith" #-}
+                         assert (tvs `equalLength` tys )
+                         substTy (zipTvSubst tvs tys) ty
 
 -- | Type substitution, see 'zipTvSubst'. Disables sanity checks.
 -- The problems that the sanity checks in substTy catch are described in
 -- Note [The substitution invariant].
 -- The goal of #11371 is to migrate all the calls of substTyUnchecked to
 -- substTy and remove this function. Please don't use in new code.
 substTyWithUnchecked :: [TyVar] -> [Type] -> Type -> Type
-substTyWithUnchecked tvs tys
+substTyWithUnchecked tvs tys ty
   = assert (tvs `equalLength` tys )
-    substTyUnchecked (zipTvSubst tvs tys)
+    substTyUnchecked (zipTvSubst tvs tys) ty
 
 -- | Substitute tyvars within a type using a known 'InScopeSet'.
 -- Pre-condition: the 'in_scope' set should satisfy Note [The substitution
@@ -606,34 +606,34 @@ substTyWithInScope in_scope tvs tys ty =
 
 -- | Coercion substitution, see 'zipTvSubst'
 substCoWith :: HasDebugCallStack => [TyVar] -> [Type] -> Coercion -> Coercion
-substCoWith tvs tys = assert (tvs `equalLength` tys )
-                      substCo (zipTvSubst tvs tys)
+substCoWith tvs tys co = assert (tvs `equalLength` tys )
+                         substCo (zipTvSubst tvs tys) co
 
 -- | Coercion substitution, see 'zipTvSubst'. Disables sanity checks.
 -- The problems that the sanity checks in substCo catch are described in
 -- Note [The substitution invariant].
 -- The goal of #11371 is to migrate all the calls of substCoUnchecked to
 -- substCo and remove this function. Please don't use in new code.
 substCoWithUnchecked :: [TyVar] -> [Type] -> Coercion -> Coercion
-substCoWithUnchecked tvs tys
+substCoWithUnchecked tvs tys co
   = assert (tvs `equalLength` tys )
-    substCoUnchecked (zipTvSubst tvs tys)
+    substCoUnchecked (zipTvSubst tvs tys) co
 
 
 
 -- | Substitute covars within a type
 substTyWithCoVars :: [CoVar] -> [Coercion] -> Type -> Type
-substTyWithCoVars cvs cos = substTy (zipCvSubst cvs cos)
+substTyWithCoVars cvs cos ty = substTy (zipCvSubst cvs cos) ty
 
 -- | Type substitution, see 'zipTvSubst'
 substTysWith :: [TyVar] -> [Type] -> [Type] -> [Type]
-substTysWith tvs tys = assert (tvs `equalLength` tys )
-                       substTys (zipTvSubst tvs tys)
+substTysWith tvs tys ty = assert (tvs `equalLength` tys )
+                          substTys (zipTvSubst tvs tys) ty
 
 -- | Type substitution, see 'zipTvSubst'
 substTysWithCoVars :: [CoVar] -> [Coercion] -> [Type] -> [Type]
-substTysWithCoVars cvs cos = assert (cvs `equalLength` cos )
-                             substTys (zipCvSubst cvs cos)
+substTysWithCoVars cvs cos tys = assert (cvs `equalLength` cos )
+                                 substTys (zipCvSubst cvs cos) tys
 
 -- | Substitute within a 'Type' after adding the free variables of the type
 -- to the in-scope set. This is useful for the case when the free variables
@@ -774,7 +774,7 @@ subst_ty :: TCvSubst -> Type -> Type
 --
 -- Note that the in_scope set is poked only if we hit a forall
 -- so it may often never be fully computed
-subst_ty subst ty
+subst_ty !subst ty
    = go ty
   where
     go (TyVarTy tv)      = substTyVar subst tv
@@ -1306,4 +1306,3 @@ substTyCoBndr subst (Anon af ty)          = (subst, Anon af (substScaledTy subst
 substTyCoBndr subst (Named (Bndr tv vis)) = (subst', Named (Bndr tv' vis))
                                           where
                                             (subst', tv') = substVarBndr subst tv
-

compiler/GHC/Types/Var/Env.hs

@@ -246,7 +246,7 @@ mkInScopeSet :: VarSet -> InScopeSet
 mkInScopeSet in_scope
   = InScope (FastVarSet (WordMap.fromAscList $ map intKeyToWordKey $ varSetToAscList in_scope))
   where
-    intKeyToWordKey (k, v) = (fromIntegral k, v)
+    intKeyToWordKey (k, v) = ((,) $! fromIntegral k) v
 
 extendInScopeSet :: InScopeSet -> Var -> InScopeSet
 extendInScopeSet (InScope in_scope) v
@@ -293,7 +293,7 @@ mkTInScopeSet in_scope = do
   tmap <- WordMap.fromAscListT $ map intKeyToWordKey $ varSetToAscList in_scope
   return $ TInScope (TFastVarSet tmap)
   where
-    intKeyToWordKey (k, v) = (fromIntegral k, v)
+    intKeyToWordKey (k, v) = ((,) $! fromIntegral k) v
 
 extendTInScopeSet :: TInScopeSet s -> Var -> ST s (TInScopeSet s)
 extendTInScopeSet (TInScope in_scope) v
@@ -401,7 +401,7 @@ unsafeGetFreshLocalUniqueT (TInScope (TFastVarSet set)) = do
   case mb_uniq of
     Just (uniq,_)
       | let uniq' = mkLocalUnique (fromIntegral uniq)
-      , not $ uniq' `ltUnique` minLocalUnique 
+      , not $ uniq' `ltUnique` minLocalUnique
       -> return $! incrUnique uniq'
     _ -> return minLocalUnique
 

compiler/GHC/Unit/Types.hs

@@ -67,6 +67,7 @@ module GHC.Unit.Types
    , mainUnitId
    , thisGhcUnitId
    , interactiveUnitId
+   , transientsUnitId
 
    , primUnit
    , bignumUnit
@@ -76,6 +77,7 @@ module GHC.Unit.Types
    , mainUnit
    , thisGhcUnit
    , interactiveUnit
+   , transientsUnit
 
    , isInteractiveModule
    , wiredInUnitIds
@@ -581,10 +583,12 @@ Make sure you change 'GHC.Unit.State.findWiredInUnits' if you add an entry here.
 -}
 
 bignumUnitId, primUnitId, baseUnitId, rtsUnitId,
-  thUnitId, mainUnitId, thisGhcUnitId, interactiveUnitId  :: UnitId
+  thUnitId, mainUnitId, thisGhcUnitId, interactiveUnitId,
+  transientsUnitId  :: UnitId
 
 bignumUnit, primUnit, baseUnit, rtsUnit,
-  thUnit, mainUnit, thisGhcUnit, interactiveUnit  :: Unit
+  thUnit, mainUnit, thisGhcUnit, interactiveUnit,
+  transientsUnit  :: Unit
 
 primUnitId        = UnitId (fsLit "ghc-prim")
 bignumUnitId      = UnitId (fsLit "ghc-bignum")
@@ -593,6 +597,7 @@ rtsUnitId         = UnitId (fsLit "rts")
 thisGhcUnitId     = UnitId (fsLit "ghc")
 interactiveUnitId = UnitId (fsLit "interactive")
 thUnitId          = UnitId (fsLit "template-haskell")
+transientsUnitId  = UnitId (fsLit "haskell-transients")
 
 thUnit            = RealUnit (Definite thUnitId)
 primUnit          = RealUnit (Definite primUnitId)
@@ -601,6 +606,7 @@ baseUnit          = RealUnit (Definite baseUnitId)
 rtsUnit           = RealUnit (Definite rtsUnitId)
 thisGhcUnit       = RealUnit (Definite thisGhcUnitId)
 interactiveUnit   = RealUnit (Definite interactiveUnitId)
+transientsUnit    = RealUnit (Definite transientsUnitId)
 
 -- | This is the package Id for the current program.  It is the default
 -- package Id if you don't specify a package name.  We don't add this prefix
@@ -619,6 +625,7 @@ wiredInUnitIds =
    , rtsUnitId
    , thUnitId
    , thisGhcUnitId
+   , transientsUnitId
    ]
 
 ---------------------------------------------------------------------