Add new bot element to latticewip/andreask/infer_closures

4 files changed, 129 insertions, 71 deletions

diff --git a/compiler/GHC/Stg/InferTags.hs b/compiler/GHC/Stg/InferTags.hs
index 3c342d93bc..a3a659af23 100644
--- a/compiler/GHC/Stg/InferTags.hs
+++ b/compiler/GHC/Stg/InferTags.hs
@@ -13,6 +13,8 @@ module GHC.Stg.InferTags ( inferTags ) where
 import GHC.Prelude
 
 import GHC.Core.DataCon
+import GHC.Core.TyCon
+import GHC.Core.Type
 import GHC.Types.Id
 import GHC.Stg.Syntax
 import GHC.Types.Basic ( Arity, TopLevelFlag(..), RecFlag(..) )
@@ -148,6 +150,11 @@ seemed more efficient.
 *                                                                      *
 ********************************************************************* -}
 
+type OutputableInferPass p = (Outputable (TagEnv p)
+                              , Outputable (GenStgExpr p)
+                              , Outputable (BinderP p)
+                              , Outputable (GenStgRhs p))
+
 inferTags :: [GenStgTopBinding 'Vanilla] -> [GenStgTopBinding 'InferTaggedBinders]
 inferTags binds =
   -- pprTrace "Binds" (pprGenStgTopBindings shortStgPprOpts $ binds) $
@@ -165,7 +172,7 @@ inferTagTopBind env (StgTopLifted bind)
 
 
 -----------------------
-inferTagExpr :: Outputable (TagEnv p)
+inferTagExpr :: OutputableInferPass p
   => TagEnv p -> GenStgExpr p -> (TagInfo, GenStgExpr 'InferTaggedBinders)
 inferTagExpr env (StgApp _ext fun args)
   = (info, StgApp noEnterInfo fun args)
@@ -174,7 +181,8 @@ inferTagExpr env (StgApp _ext fun args)
          , arity == length args  -- Saturated
          = res_info
          | otherwise
-         = TagDunno
+         = --pprTrace "inferAppUnknown" (ppr fun) $
+           TagDunno
 
 inferTagExpr env (StgConApp con cn args tys)
   = (info, StgConApp con cn args tys)
@@ -214,27 +222,47 @@ inferTagExpr env (StgLetNoEscape ext bind body)
 inferTagExpr env (StgCase scrut bndr ty alts)
   | [(DataAlt con, bndrs, rhs)] <- alts
   , isUnboxedTupleDataCon con
-  , TagTuple infos <- scrut_info
+  , Just infos <- scrut_infos bndrs
+  -- , pprTrace "scrut info:" (ppr infos $$ ppr scrut $$ ppr bndrs) True
   , let bndrs' = zipWithEqual "inferTagExpr" mk_bndr bndrs infos
-        mk_bndr bndr info = (getBinderId env bndr, TagSig 0 info)
+        mk_bndr bndr info =
+            --  pprTrace "mk_ubx_bndr_info" ( ppr bndr <+> ppr info ) $
+            (getBinderId env bndr, TagSig 0 info)
+        -- no case binder in alt_env here, unboxed tuple binders are dead after unarise
         alt_env = extendSigEnv env bndrs'
         (info, rhs') = inferTagExpr alt_env rhs
-  = (info, StgCase scrut' (noSig env bndr) ty [(DataAlt con, bndrs', rhs')])
+  = -- pprTrace "inferCase1" (ppr scrut $$ ppr bndr $$ ppr infos $$ ppr bndrs') $
+    (info, StgCase scrut' (noSig env bndr) ty [(DataAlt con, bndrs', rhs')])
+
   | null alts -- Empty case, but I might just be paranoid.
-  = (TagDunno, StgCase scrut' bndr' ty [])
+  = -- pprTrace "inferCase2" empty $
+    (TagDunno, StgCase scrut' bndr' ty [])
+  -- More than one alternative OR non-tuple single alternative.
   | otherwise
-  = ( foldr combineAltInfo TagProper infos
+  = -- pprTrace "inferCase3" empty $
+    let
+        alt_env = extendSigEnv env [bndr']
+        (infos, alts')
+          = unzip [ (info, (con, bndrs', rhs'))
+                  | (con, bndrs, rhs) <- alts
+                  , let (info, rhs') = inferTagExpr alt_env rhs
+                        bndrs' = addAltBndrInfo env con bndrs ]
+        alt_info = foldr combineAltInfo TagTagged infos
+    in -- pprTrace "combine alts:" (ppr alt_info $$ ppr infos)
+    ( foldr combineAltInfo TagTagged infos
     , StgCase scrut' bndr' ty alts')
   where
+    -- Single unboxed tuple alternative
+    scrut_infos bndrs = case scrut_info of
+      TagTagged -> Just $ replicate (length bndrs) TagProper
+      TagTuple infos -> Just infos
+      _ -> Nothing
     (scrut_info, scrut') = inferTagExpr env scrut
     bndr' = (getBinderId env bndr, TagSig 0 TagProper)
-    alt_env = extendSigEnv env [bndr']
-    (infos, alts')
-       = unzip [ (info, (con, bndrs', rhs'))
-               | (con, bndrs, rhs) <- alts
-               , let (info, rhs') = inferTagExpr alt_env rhs
-                     bndrs' = addAltBndrInfo env con bndrs ]
 
+
+
+-- Not used if we have tuple info about the scrutinee
 addAltBndrInfo :: TagEnv p -> AltCon -> [BinderP p] -> [BinderP 'InferTaggedBinders]
 addAltBndrInfo env (DataAlt con) bndrs
   = zipWithEqual "inferTagAlt" mk_bndr bndrs marks
@@ -249,14 +277,14 @@ addAltBndrInfo env (DataAlt con) bndrs
 addAltBndrInfo env _ bndrs = map (noSig env) bndrs
 
 -----------------------------
-inferTagBind :: Outputable (TagEnv p)
+inferTagBind :: OutputableInferPass p
   => TopLevelFlag -> TagEnv p -> GenStgBinding p -> (TagEnv p, GenStgBinding 'InferTaggedBinders)
 inferTagBind top env (StgNonRec bndr rhs)
   = (env', StgNonRec (id, sig) rhs')
   where
     id   = getBinderId env bndr
     env' = extendSigEnv env [(id, sig)]
-    (sig,rhs') = inferTagRhs top [] env rhs
+    (sig,rhs') = inferTagRhs top [id] env rhs
 
 inferTagBind top env (StgRec pairs)
   = (env { te_env = sig_env }, StgRec pairs')
@@ -266,7 +294,7 @@ inferTagBind top env (StgRec pairs)
     init_sigs         = map initSig rhss
     (sig_env, pairs') = go env init_sigs rhss
 
-    go :: forall q. TagEnv q -> [TagSig] -> [GenStgRhs q]
+    go :: forall q. OutputableInferPass q => TagEnv q -> [TagSig] -> [GenStgRhs q]
                  -> (TagSigEnv, [((Id,TagSig), GenStgRhs 'InferTaggedBinders)])
     go env sigs rhss
       --  | pprTrace "go" (ppr ids $$ ppr sigs $$ ppr sigs') False
@@ -281,11 +309,11 @@ inferTagBind top env (StgRec pairs)
 
 initSig :: GenStgRhs p -> TagSig
 -- Initial signature for the fixpoint loop
-initSig StgRhsCon {}                = TagSig 0              TagProper
-initSig (StgRhsClosure _ _ _ bndrs _) = TagSig (length bndrs) TagProper
+initSig (StgRhsCon {})                = TagSig 0              TagProper
+initSig (StgRhsClosure _ _ _ bndrs _) = TagSig (length bndrs) TagTagged
 
 -----------------------------
-inferTagRhs :: Outputable (TagEnv p)
+inferTagRhs :: OutputableInferPass p
   => TopLevelFlag -- ^
   -> [Id] -- ^ List of ids in the recursive group, or [] otherwise
   -> TagEnv p -- ^
@@ -307,29 +335,27 @@ inferTagRhs _top _grp_ids env (StgRhsClosure ext cc upd bndrs body)
       | otherwise  = info
     bndrs' = map (noSig env) bndrs
 
-inferTagRhs top grp_ids env (StgRhsCon cc con cn ticks args)
+inferTagRhs _top _grp_ids env rhs@(StgRhsCon cc con cn ticks args)
 -- Top level constructors, which have untagged arguments to strict fields
 -- become thunks. Same goes for rhs which are part of a recursive group.
 -- We encode this by giving changing RhsCon nodes the info TagDunno
   = --pprTrace "inferTagRhsCon" (ppr grp_ids) $
     let
         strictArgs = zipEqual "inferTagRhs" args (dataConRuntimeRepStrictness con)
+        -- argInfo = [(lookupInfo env (StgVarArg v)) | StgVarArg v <- args ]
         strictUntaggedIds = [v | (StgVarArg v, MarkedStrict) <- strictArgs
-                            , lookupInfo env (StgVarArg v) /= TagProper] :: [Id]
-
-        mkResult x = (TagSig 0 x, StgRhsCon cc con cn ticks args)
+                            , not (isTaggedInfo (lookupInfo env (StgVarArg v))) ] :: [Id]
+
+        mkResult x =
+          -- pprTrace "inferTagRhsCon"
+          --   ( ppr _grp_ids <+> ppr x <+> ppr rhs $$
+          --     ppr strictArgs $$
+          --     ppr strictUntaggedIds $$
+          --     ppr argInfo $$
+          --     text "con:" <> ppr con
+          --     ) $
+            (TagSig 0 x, StgRhsCon cc con cn ticks args)
     in case () of
           -- All fields tagged or non-strict
         _ | null strictUntaggedIds -> mkResult TagProper
-          -- -- Non-recursive local let
-          -- | null grp_ids
-          -- , NotTopLevel <- top
-          -- -> mkResult TagProper
-          -- Recursive local let, no bindings from grp in args
-          -- | NotTopLevel <- top
-          -- , mkVarSet grp_ids `disjointVarSet` mkVarSet strictUntaggedIds
-          -- -> mkResult TagProper
-          -- Otherwise we have a top level let with untagged args,
-          -- or a recursive group where a bindings of the group is
-          -- passed into a strict field
           | otherwise -> mkResult TagDunno
diff --git a/compiler/GHC/Stg/InferTags/Rewrite.hs b/compiler/GHC/Stg/InferTags/Rewrite.hs
index 127e75f848..4f605d4064 100644
--- a/compiler/GHC/Stg/InferTags/Rewrite.hs
+++ b/compiler/GHC/Stg/InferTags/Rewrite.hs
@@ -158,6 +158,7 @@ isTagged v = do
                         case info of
                             TagDunno -> False
                             TagProper -> True
+                            TagTagged -> True
                             TagTuple _ -> True -- Consider unboxed tuples tagged
         False -- Imported
             | Just con <- (isDataConWorkId_maybe v)
@@ -214,26 +215,27 @@ rewriteTop (StgTopStringLit v s) = return $! (StgTopStringLit v s)
 rewriteTop (StgTopLifted bind)   = do
     -- Top level bindings can, and must remain in scope
     addBind bind
-    (StgTopLifted . fst) <$!> (rewriteBinds bind)
+    (StgTopLifted) <$!> (rewriteBinds bind)
 
 -- For top level binds, the wrapper is guaranteed to be `id`
-rewriteBinds :: InferStgBinding -> RM (TgStgBinding, TgStgExpr -> TgStgExpr)
+rewriteBinds :: InferStgBinding -> RM (TgStgBinding)
 rewriteBinds (StgNonRec v rhs) = do
-        (!rhs, wrapper) <-  rewriteRhs v rhs
-        return $! (StgNonRec (fst v) rhs, wrapper)
+        (!rhs) <-  rewriteRhs v rhs
+        return $! (StgNonRec (fst v) rhs)
 rewriteBinds b@(StgRec binds) =
     -- Bring sigs of binds into scope for all rhss
     withBind b $ do
-        (rhss, wrappers) <- unzip <$> mapM (uncurry rewriteRhs) binds
-        let wrapper = foldl1 (.) wrappers
-        return $! (mkRec rhss, wrapper)
+        (rhss) <- mapM (uncurry rewriteRhs) binds
+        return $! (mkRec rhss)
         where
             mkRec :: [TgStgRhs] -> TgStgBinding
             mkRec rhss = StgRec (zip (map (fst . fst) binds) rhss)
 
 -- Rewrite a RHS, the rewriteFlag tells us weither or not the RHS is in a context in which
 -- we can avoid turning the RhsCon into a closure. (e.g. for top level bindings)
-rewriteRhs :: (Id,TagSig) -> InferStgRhs -> RM (TgStgRhs, TgStgExpr -> TgStgExpr)
+rewriteRhs :: (Id,TagSig) -> InferStgRhs
+           -> RM (-- Bool, -- Should we turn it into an updateable closure
+                TgStgRhs)
 rewriteRhs (_id, tagSig) (StgRhsCon ccs con cn ticks args) = {-# SCC rewriteRhs_ #-} do
     -- pprTraceM "rewriteRhs" (ppr _id)
 
@@ -250,35 +252,24 @@ rewriteRhs (_id, tagSig) (StgRhsCon ccs con cn ticks args) = {-# SCC rewriteRhs_
     let evalArgs = [v | StgVarArg v <- needsEval] :: [Id]
 
     if (null evalArgs)
-        then return $! (StgRhsCon ccs con cn ticks args, id)
+        then return $! (StgRhsCon ccs con cn ticks args)
         else do
+            --assert not (isTaggedSig tagSig)
             -- pprTraceM "CreatingSeqs for " $ ppr _id <+> ppr node_id
 
-            evaldArgs <- mapM mkLocalArgId evalArgs -- Create case binders
-            let varMap = zip evalArgs evaldArgs -- Match them up with original ids
-            let updateArg (StgLitArg lit) = (StgLitArg lit)
-                updateArg (StgVarArg v)
-                    | Just v' <- lookup v varMap
-                    = StgVarArg v'
-                    | otherwise = StgVarArg v
-            let evaldConArgs = map updateArg args
             -- At this point iff we have:
             --  * possibly untagged arguments to strict fields
-            --  * and Dunno as tag signature
-            -- Then we return a RhsClosure, otherwise we return a wrapper
+            --  * then inference marked the binder as tag Dunno
+            -- So we convert it into a RhsClosure.
             -- which will evaluate the arguments first when applied to an expression.
-            if not (isTaggedSig tagSig) || True --rewriteFlag == MaybeClosure
-                then do -- Turn the rhs into a closure that evaluates the arguments to the strict fields
-                    conExpr <- mkSeqs evalArgs con cn args (panic "mkSeqs should not need to provide types")
-                    return $! (StgRhsClosure noExtFieldSilent ccs ReEntrant [] $! conExpr, id)
-                else do -- Return a case expression that will evaluate the arguments.
-                    let evalExpr expr = foldr (\(v, vEvald) e -> mkSeq v vEvald e) expr varMap
-                    return $! ((StgRhsCon ccs con cn ticks evaldConArgs), evalExpr)
+            -- Turn the rhs into a closure that evaluates the arguments to the strict fields
+            conExpr <- mkSeqs evalArgs con cn args (panic "mkSeqs should not need to provide types")
+            return $! (StgRhsClosure noExtFieldSilent ccs Updatable [] $! conExpr)
 rewriteRhs _binding (StgRhsClosure ext ccs flag args body) = do
     -- mapM_ addBinder  args
     withBinders args $ do
         closure <- StgRhsClosure ext ccs flag (map fst args) <$> rewriteExpr False body
-        return (closure, id)
+        return (closure)
 
 type IsScrut = Bool
 
@@ -312,18 +303,18 @@ rewriteAlt (altCon, bndrs, rhs) = do
 
 rewriteLet :: InferStgExpr -> RM TgStgExpr
 rewriteLet (StgLet xt bind expr) = do
-    (!bind', !wrapper) <- rewriteBinds bind
+    (!bind') <- rewriteBinds bind
     withBind bind $ do
         !expr' <- rewriteExpr False expr
-        return $! wrapper (StgLet xt bind' expr')
+        return $! (StgLet xt bind' expr')
 rewriteLet _ = panic "Impossible"
 
 rewriteLetNoEscape :: InferStgExpr -> RM TgStgExpr
 rewriteLetNoEscape (StgLetNoEscape xt bind expr) = do
-    (!bind', wrapper) <- rewriteBinds bind
+    (!bind') <- rewriteBinds bind
     withBind bind $ do
         !expr' <- rewriteExpr False expr
-        return $! wrapper (StgLetNoEscape xt bind' expr')
+        return $! (StgLetNoEscape xt bind' expr')
 rewriteLetNoEscape _ = panic "Impossible"
 
 rewriteConApp :: InferStgExpr -> RM TgStgExpr
@@ -373,7 +364,7 @@ mkSeqs :: [Id] -> DataCon -> ConstructorNumber -> [StgArg] -> [Type] -> RM TgStg
 mkSeqs untaggedIds con cn args tys = do
     argMap <- mapM (\arg -> (arg,) <$> mkLocalArgId arg ) untaggedIds :: RM [(InId, OutId)]
     -- mapM_ (pprTraceM "Forcing strict args before allocation:" . ppr) argMap
-    let taggedArgs
+    let taggedArgs :: [StgArg]
             = map   (\v -> case v of
                         StgVarArg v' -> StgVarArg $ fromMaybe v' $ lookup v' argMap
                         lit -> lit)
diff --git a/compiler/GHC/Stg/InferTags/Types.hs b/compiler/GHC/Stg/InferTags/Types.hs
index 2beb1523fe..9d94054d0d 100644
--- a/compiler/GHC/Stg/InferTags/Types.hs
+++ b/compiler/GHC/Stg/InferTags/Types.hs
@@ -22,6 +22,9 @@ import GHC.Types.Var.Env
 import GHC.Utils.Outputable
 import GHC.Utils.Misc( zipWithEqual )
 
+import GHC.StgToCmm.Types ( LambdaFormInfo(..) )
+import GHC.Utils.Panic
+
 {- *********************************************************************
 *                                                                      *
                          Supporting data types
@@ -48,20 +51,24 @@ data TagInfo
   = TagDunno
   | TagTuple [TagInfo]  -- Unboxed tuple
   | TagProper           -- Heap pointer to properly-tagged value
-                        -- Bottom of the domain
+  | TagTagged           -- Bottom of the domain.
   deriving( Eq )
 
 instance Outputable TagInfo where
+  ppr TagTagged      = text "TagTagged"
   ppr TagDunno       = text "TagDunno"
   ppr TagProper      = text "TagProper"
   ppr (TagTuple tis) = text "TagTuple" <> brackets (pprWithCommas ppr tis)
 
 combineAltInfo :: TagInfo -> TagInfo -> TagInfo
 combineAltInfo TagDunno         _              = TagDunno
-combineAltInfo TagProper        ti             = ti
-combineAltInfo (TagTuple {})    TagDunno       = TagDunno
-combineAltInfo ti@(TagTuple {}) TagProper      = ti
-combineAltInfo (TagTuple is1)   (TagTuple is2) = TagTuple (zipWithEqual "combineAltInfo" combineAltInfo is1 is2)
+combineAltInfo _                TagDunno       = TagDunno
+combineAltInfo (TagTuple {})    TagProper      = panic "Combining unboxed tuple with non-tuple result"
+combineAltInfo TagProper       (TagTuple {})   = panic "Combining unboxed tuple with non-tuple result"
+combineAltInfo TagProper        TagProper      = TagProper
+combineAltInfo (TagTuple is1)  (TagTuple is2)  = TagTuple (zipWithEqual "combineAltInfo" combineAltInfo is1 is2)
+combineAltInfo (TagTagged)      ti             = ti
+combineAltInfo ti               TagTagged      = ti
 
 type TagSigEnv = IdEnv TagSig
 data TagEnv p = TE { te_env :: TagSigEnv
@@ -120,6 +127,24 @@ lookupInfo env (StgVarArg var)
   | Just (TagSig 0 info) <- lookupVarEnv (te_env env) var
   = info
 
+  -- | Just lf_info <- idLFInfo_maybe var
+  -- =   case lf_info of
+  --         -- Function, tagged (with arity)
+  --         LFReEntrant {}
+  --             -> TagProper
+  --         -- Thunks need to be entered.
+  --         LFThunk {}
+  --             -> TagDunno
+  --         -- Constructors, already tagged.
+  --         LFCon {}
+  --             -> TagProper
+  --         LFUnknown {}
+  --             -> TagDunno
+  --         LFUnlifted {}
+  --             -> TagProper
+  --         -- Shouldn't be possible. I don't think we can export letNoEscapes
+  --         LFLetNoEscape {} -> panic "LFLetNoEscape exported"
+
   | otherwise
   = TagDunno
 
@@ -130,11 +155,18 @@ isDunnoSig :: TagSig -> Bool
 isDunnoSig (TagSig _ TagDunno) = True
 isDunnoSig (TagSig _ TagProper) = False
 isDunnoSig (TagSig _ TagTuple{}) = False
+isDunnoSig (TagSig _ TagTagged{}) = False
 
 isTaggedSig :: TagSig -> Bool
 isTaggedSig (TagSig _ TagProper) = True
+isTaggedSig (TagSig _ TagTagged) = True
 isTaggedSig _ = False
 
+isTaggedInfo :: TagInfo -> Bool
+isTaggedInfo TagProper = True
+isTaggedInfo TagTagged = True
+isTaggedInfo _         = False
+
 extendSigEnv :: TagEnv p -> [(Id,TagSig)] -> TagEnv p
 extendSigEnv env@(TE { te_env = sig_env }) bndrs
   = env { te_env = extendVarEnvList sig_env bndrs }
diff --git a/compiler/GHC/Stg/Syntax.hs b/compiler/GHC/Stg/Syntax.hs
index 7bc18dc07b..c9f9a6cd5f 100644
--- a/compiler/GHC/Stg/Syntax.hs
+++ b/compiler/GHC/Stg/Syntax.hs
@@ -755,6 +755,9 @@ pprGenStgBinding opts b = case b of
                              = hang (hsep [pprBndr LetBind bndr, equals])
                                     4 (pprStgRhs opts expr <> semi)
 
+instance OutputablePass pass => Outputable  (GenStgBinding pass) where
+  ppr = pprGenStgBinding panicStgPprOpts
+
 pprGenStgTopBindings :: (OutputablePass pass) => StgPprOpts -> [GenStgTopBinding pass] -> SDoc
 pprGenStgTopBindings opts binds
   = vcat $ intersperse blankLine (map (pprGenStgTopBinding opts) binds)
@@ -775,6 +778,9 @@ pprStgArg :: StgArg -> SDoc
 pprStgArg (StgVarArg var) = ppr var
 pprStgArg (StgLitArg con) = ppr con
 
+instance OutputablePass pass => Outputable  (GenStgExpr pass) where
+  ppr = pprStgExpr panicStgPprOpts
+
 pprStgExpr :: OutputablePass pass => StgPprOpts -> GenStgExpr pass -> SDoc
 pprStgExpr opts e = case e of
                            -- special case
@@ -891,3 +897,6 @@ pprStgRhs opts rhs = case rhs of
                   NoNumber -> empty
                   Numbered n -> hcat [ppr n, space]
               , ppr con, text "! ", brackets (sep (map pprStgArg args))]
+
+instance OutputablePass pass => Outputable  (GenStgRhs pass) where
+  ppr = pprStgRhs panicStgPprOpts