Refactor GenStgRhs to include the Type in both constructors (#23280, #22576, #22364)

Carry the actual type of an expression through the PreStgRhs and into GenStgRhs
for use in later stages. Currently this is used in the JavaScript backend to fix
some tests from the above mentioned issues: EtaExpandLevPoly, RepPolyWrappedVar2,
T13822, T14749.

4 files changed, 44 insertions, 55 deletions

diff --git a/compiler/GHC/StgToJS/CodeGen.hs b/compiler/GHC/StgToJS/CodeGen.hs
index 55be51df9d..fdc431ef4c 100644
--- a/compiler/GHC/StgToJS/CodeGen.hs
+++ b/compiler/GHC/StgToJS/CodeGen.hs
@@ -290,10 +290,10 @@ genToplevelDecl i rhs = do
 
 genToplevelConEntry :: Id -> CgStgRhs -> G JStat
 genToplevelConEntry i rhs = case rhs of
-   StgRhsCon _cc con _mu _ts _args
+   StgRhsCon _cc con _mu _ts _args _typ
      | isDataConWorkId i
        -> genSetConInfo i con (stgRhsLive rhs) -- NoSRT
-   StgRhsClosure _ _cc _upd_flag _args _body
+   StgRhsClosure _ _cc _upd_flag _args _body _typ
      | Just dc <- isDataConWorkId_maybe i
        -> genSetConInfo i dc (stgRhsLive rhs) -- srt
    _ -> pure mempty
@@ -321,11 +321,11 @@ mkDataEntry = ValExpr $ JFunc [] returnStack
 genToplevelRhs :: Id -> CgStgRhs -> G JStat
 -- general cases:
 genToplevelRhs i rhs = case rhs of
-  StgRhsCon cc con _mu _tys args -> do
+  StgRhsCon cc con _mu _tys args _typ -> do
     ii <- identForId i
     allocConStatic ii cc con args
     return mempty
-  StgRhsClosure _ext cc _upd_flag {- srt -} args body -> do
+  StgRhsClosure _ext cc _upd_flag {- srt -} args body typ -> do
     {-
       algorithm:
        - collect all Id refs that are in the global id cache
@@ -335,7 +335,7 @@ genToplevelRhs i rhs = case rhs of
     -}
     eid@(TxtI eidt) <- identForEntryId i
     (TxtI idt)   <- identForId i
-    body <- genBody (initExprCtx i) i R2 args body
+    body <- genBody (initExprCtx i) R2 args body typ
     global_occs <- globalOccs (jsSaturate (Just "ghcjs_tmp_sat_") body)
     let lidents = map global_ident global_occs
     let lids    = map global_id    global_occs
diff --git a/compiler/GHC/StgToJS/Expr.hs b/compiler/GHC/StgToJS/Expr.hs
index 9f5a1f6d0a..0b8e34e14b 100644
--- a/compiler/GHC/StgToJS/Expr.hs
+++ b/compiler/GHC/StgToJS/Expr.hs
@@ -74,7 +74,6 @@ import GHC.Core.Type hiding (typeSize)
 import GHC.Utils.Misc
 import GHC.Utils.Monad
 import GHC.Utils.Panic
-import GHC.Utils.Panic.Plain
 import GHC.Utils.Outputable (ppr, renderWithContext, defaultSDocContext)
 import qualified Control.Monad.Trans.State.Strict as State
 import GHC.Data.FastString
@@ -148,7 +147,7 @@ genBind ctx bndr =
      ctx' = ctxClearLneFrame ctx
 
      assign :: Id -> CgStgRhs -> G (Maybe JStat)
-     assign b (StgRhsClosure _ _ccs {-[the_fv]-} _upd [] expr)
+     assign b (StgRhsClosure _ _ccs {-[the_fv]-} _upd [] expr _typ)
        | let strip = snd . stripStgTicksTop (not . tickishIsCode)
        , StgCase (StgApp scrutinee []) _ (AlgAlt _) [GenStgAlt (DataAlt _) params sel_expr] <- strip expr
        , StgApp selectee [] <- strip sel_expr
@@ -168,7 +167,7 @@ genBind ctx bndr =
              ([tgt], [the_fvj]) -> return $ Just
                (tgt ||= ApplExpr (var ("h$c_sel_" <> mkFastString sel_tag)) [the_fvj])
              _ -> panic "genBind.assign: invalid size"
-     assign b (StgRhsClosure _ext _ccs _upd [] expr)
+     assign b (StgRhsClosure _ext _ccs _upd [] expr _typ)
        | snd (isInlineExpr (ctxEvaluatedIds ctx) expr) = do
            d   <- declVarsForId b
            tgt <- varsForId b
@@ -180,9 +179,9 @@ genBind ctx bndr =
 
      addEvalRhs c [] = c
      addEvalRhs c ((b,r):xs)
-       | StgRhsCon{} <- r                       = addEvalRhs (ctxAssertEvaluated b c) xs
-       | (StgRhsClosure _ _ ReEntrant _ _) <- r = addEvalRhs (ctxAssertEvaluated b c) xs
-       | otherwise                              = addEvalRhs c xs
+       | StgRhsCon{} <- r                         = addEvalRhs (ctxAssertEvaluated b c) xs
+       | (StgRhsClosure _ _ ReEntrant _ _ _) <- r = addEvalRhs (ctxAssertEvaluated b c) xs
+       | otherwise                                = addEvalRhs c xs
 
 genBindLne :: HasDebugCallStack
            => ExprCtx
@@ -223,7 +222,7 @@ genBindLne ctx bndr = do
 -- is initially set to null, changed to h$blackhole when the thunk is being evaluated.
 --
 genEntryLne :: HasDebugCallStack => ExprCtx -> Id -> CgStgRhs -> G ()
-genEntryLne ctx i rhs@(StgRhsClosure _ext _cc update args body) =
+genEntryLne ctx i rhs@(StgRhsClosure _ext _cc update args body typ) =
   resetSlots $ do
   let payloadSize = ctxLneFrameSize ctx
       vars        = ctxLneFrameVars ctx
@@ -238,7 +237,7 @@ genEntryLne ctx i rhs@(StgRhsClosure _ext _cc update args body) =
               ])
          | otherwise = mempty
   lvs  <- popLneFrame True payloadSize ctx
-  body <- genBody ctx i R1 args body
+  body <- genBody ctx R1 args body typ
   ei@(TxtI eii) <- identForEntryId i
   sr   <- genStaticRefsRhs rhs
   let f = JFunc [] (bh <> lvs <> body)
@@ -251,7 +250,7 @@ genEntryLne ctx i rhs@(StgRhsClosure _ext _cc update args body) =
                 CIStackFrame
                 sr
   emitToplevel (ei ||= toJExpr f)
-genEntryLne ctx i (StgRhsCon cc con _mu _ticks args) = resetSlots $ do
+genEntryLne ctx i (StgRhsCon cc con _mu _ticks args _typ) = resetSlots $ do
   let payloadSize = ctxLneFrameSize ctx
   ei@(TxtI _eii) <- identForEntryId i
   -- di <- varForDataConWorker con
@@ -265,12 +264,12 @@ genEntryLne ctx i (StgRhsCon cc con _mu _ticks args) = resetSlots $ do
 -- | Generate the entry function for a local closure
 genEntry :: HasDebugCallStack => ExprCtx -> Id -> CgStgRhs -> G ()
 genEntry _ _i StgRhsCon {} = return ()
-genEntry ctx i rhs@(StgRhsClosure _ext cc {-_bi live-} upd_flag args body) = resetSlots $ do
+genEntry ctx i rhs@(StgRhsClosure _ext cc {-_bi live-} upd_flag args body typ) = resetSlots $ do
   let live = stgLneLiveExpr rhs -- error "fixme" -- probably find live vars in body
   ll    <- loadLiveFun live
   llv   <- verifyRuntimeReps live
   upd   <- genUpdFrame upd_flag i
-  body  <- genBody entryCtx i R2 args body
+  body  <- genBody entryCtx R2 args body typ
   ei@(TxtI eii) <- identForEntryId i
   et    <- genEntryType args
   setcc <- ifProfiling $
@@ -302,12 +301,12 @@ genEntryType args0 = do
 -- | Generate the body of an object
 genBody :: HasDebugCallStack
          => ExprCtx
-         -> Id
          -> StgReg
          -> [Id]
          -> CgStgExpr
+         -> Type
          -> G JStat
-genBody ctx i startReg args e = do
+genBody ctx startReg args e typ = do
   -- load arguments into local variables
   la <- do
     args' <- concatMapM genIdArgI args
@@ -318,7 +317,7 @@ genBody ctx i startReg args e = do
 
   -- compute PrimReps and their number of slots required to return the result of
   -- i applied to args.
-  let res_vars = resultSize args i
+  let res_vars = resultSize typ
 
   -- compute typed expressions for each slot and assign registers
   let go_var regs = \case
@@ -359,22 +358,12 @@ genBody ctx i startReg args e = do
 -- In case of failure to determine the type, we default to LiftedRep as it's
 -- probably what it is.
 --
-resultSize :: HasDebugCallStack => [Id] -> Id -> [(PrimRep, Int)]
-resultSize args i = result
+resultSize :: HasDebugCallStack => Type -> [(PrimRep, Int)]
+resultSize ty = result
   where
     result       = result_reps `zip` result_slots
     result_slots = fmap (slotCount . primRepSize) result_reps
-    result_reps  = trim_args (unwrapType (idType i)) (length args)
-
-    trim_args t 0 = typePrimRep t
-    trim_args t n
-      | Just (_af, _mult, arg, res) <- splitFunTy_maybe t
-      , nargs <- length (typePrimRepArgs arg)
-      , assert (n >= nargs) True
-      = trim_args (unwrapType res) (n - nargs)
-      | otherwise
-      = pprTrace "result_type: not a function type, assume LiftedRep" (ppr t)
-          [LiftedRep]
+    result_reps  = typePrimRep ty
 
 -- | Ensure that the set of identifiers has valid 'RuntimeRep's. This function
 -- returns a no-op when 'csRuntimeAssert' in 'StgToJSConfig' is False.
@@ -540,19 +529,19 @@ allocCls dynMiddle xs = do
       toCl (i, StgRhsCon cc con []) = do
       ii <- identForId i
       Left <$> (return (decl ii) <> allocCon ii con cc []) -}
-    toCl (i, StgRhsCon cc con _mui _ticjs [a]) | isUnboxableCon con = do
+    toCl (i, StgRhsCon cc con _mui _ticjs [a] _typ) | isUnboxableCon con = do
       ii <- identForId i
       ac <- allocCon ii con cc =<< genArg a
       pure (Left (decl ii <> ac))
 
     -- dynamics
-    toCl (i, StgRhsCon cc con _mu _ticks ar) =
+    toCl (i, StgRhsCon cc con _mu _ticks ar _typ) =
       -- fixme do we need to handle unboxed?
       Right <$> ((,,,) <$> identForId i
                        <*> varForDataConWorker con
                        <*> concatMapM genArg ar
                        <*> pure cc)
-    toCl (i, cl@(StgRhsClosure _ext cc _upd_flag _args _body)) =
+    toCl (i, cl@(StgRhsClosure _ext cc _upd_flag _args _body _typ)) =
       let live = stgLneLiveExpr cl
       in  Right <$> ((,,,) <$> identForId i
                        <*> varForEntryId i
diff --git a/compiler/GHC/StgToJS/Sinker.hs b/compiler/GHC/StgToJS/Sinker.hs
index 6df58d4fcf..f758a7ac94 100644
--- a/compiler/GHC/StgToJS/Sinker.hs
+++ b/compiler/GHC/StgToJS/Sinker.hs
@@ -64,11 +64,11 @@ sinkPgm' m pgm =
 alwaysSinkable :: CgStgBinding -> [(Id, CgStgExpr)]
 alwaysSinkable (StgRec {})       = []
 alwaysSinkable (StgNonRec b rhs) = case rhs of
-  StgRhsClosure _ _ _ _ e@(StgLit l)
+  StgRhsClosure _ _ _ _ e@(StgLit l) _
     | isSmallSinkableLit l
     , isLocal b
     -> [(b,e)]
-  StgRhsCon _ccs dc cnum _ticks as@[StgLitArg l]
+  StgRhsCon _ccs dc cnum _ticks as@[StgLitArg l] _typ
     | isSmallSinkableLit l
     , isLocal b
     , isUnboxableCon dc
@@ -88,9 +88,9 @@ onceSinkable _m (StgNonRec b rhs)
   , isLocal b = [(b,e)]
   where
     getSinkable = \case
-      StgRhsCon _ccs dc cnum _ticks args -> Just (StgConApp dc cnum args [])
-      StgRhsClosure _ _ _ _ e@(StgLit{}) -> Just e
-      _                                  -> Nothing
+      StgRhsCon _ccs dc cnum _ticks args _typ -> Just (StgConApp dc cnum args [])
+      StgRhsClosure _ _ _ _ e@(StgLit{}) _typ -> Just e
+      _                                       -> Nothing
 onceSinkable _ _ = []
 
 -- | collect all idents used only once in an argument at the top level
@@ -115,8 +115,8 @@ collectArgsTop = \case
 
 collectArgsTopRhs :: CgStgRhs -> [Id]
 collectArgsTopRhs = \case
-  StgRhsCon _ccs _dc _mu _ticks args -> concatMap collectArgsA args
-  StgRhsClosure {}                   -> []
+  StgRhsCon _ccs _dc _mu _ticks args _typ -> concatMap collectArgsA args
+  StgRhsClosure {}                        -> []
 
 -- | fold over all Id in StgArg in the AST
 collectArgs :: CgStgBinding -> [Id]
@@ -126,8 +126,8 @@ collectArgs = \case
 
 collectArgsR :: CgStgRhs -> [Id]
 collectArgsR = \case
-  StgRhsClosure _x0 _x1 _x2 _x3 e     -> collectArgsE e
-  StgRhsCon _ccs _con _mu _ticks args -> concatMap collectArgsA args
+  StgRhsClosure _x0 _x1 _x2 _x3 e _typ     -> collectArgsE e
+  StgRhsCon _ccs _con _mu _ticks args _typ -> concatMap collectArgsA args
 
 collectArgsAlt :: CgStgAlt -> [Id]
 collectArgsAlt alt = collectArgsE (alt_rhs alt)
@@ -171,7 +171,7 @@ topSortDecls _m binds = rest ++ nr'
     keys = mkUniqSet (map node_key vs)
     getV e@(StgNonRec b _) = DigraphNode e b []
     getV _                 = error "topSortDecls: getV, unexpected binding"
-    collectDeps (StgNonRec b (StgRhsCon _cc _dc _cnum _ticks args)) =
+    collectDeps (StgNonRec b (StgRhsCon _cc _dc _cnum _ticks args _typ)) =
       [ (i, b) | StgVarArg i <- args, i `elementOfUniqSet` keys ]
     collectDeps _ = []
     g = graphFromVerticesAndAdjacency vs (concatMap collectDeps nr)
diff --git a/compiler/GHC/StgToJS/StgUtils.hs b/compiler/GHC/StgToJS/StgUtils.hs
index 62c494c3a7..0632ce8fe6 100644
--- a/compiler/GHC/StgToJS/StgUtils.hs
+++ b/compiler/GHC/StgToJS/StgUtils.hs
@@ -67,8 +67,8 @@ bindingRefs u = \case
 
 rhsRefs :: UniqFM Id CgStgExpr -> CgStgRhs -> Set Id
 rhsRefs u = \case
-  StgRhsClosure _ _ _ _ body       -> exprRefs u body
-  StgRhsCon _ccs d _mu _ticks args -> l s [ i | AnId i <- dataConImplicitTyThings d] <> l (argRefs u) args
+  StgRhsClosure _ _ _ _ body _       -> exprRefs u body
+  StgRhsCon _ccs d _mu _ticks args _ -> l s [ i | AnId i <- dataConImplicitTyThings d] <> l (argRefs u) args
 
 exprRefs :: UniqFM Id CgStgExpr -> CgStgExpr -> Set Id
 exprRefs u = \case
@@ -97,7 +97,7 @@ hasExport bnd =
     StgNonRec b e -> isExportedBind b e
     StgRec bs     -> any (uncurry isExportedBind) bs
   where
-    isExportedBind _i (StgRhsCon _cc con _ _ _) =
+    isExportedBind _i (StgRhsCon _cc con _ _ _ _) =
       getUnique con == staticPtrDataConKey
     isExportedBind _ _ = False
 
@@ -152,8 +152,8 @@ stgBindRhsLive b =
 
 stgRhsLive :: CgStgRhs -> LiveVars
 stgRhsLive = \case
-  StgRhsClosure _ _ _ args e -> delDVarSetList (stgExprLive True e) args
-  StgRhsCon _ _ _ _ args     -> unionDVarSets (map stgArgLive args)
+  StgRhsClosure _ _ _ args e _ -> delDVarSetList (stgExprLive True e) args
+  StgRhsCon _ _ _ _ args _     -> unionDVarSets (map stgArgLive args)
 
 stgArgLive :: StgArg -> LiveVars
 stgArgLive = \case
@@ -189,8 +189,8 @@ bindees = \case
   StgRec bs      -> map fst bs
 
 isUpdatableRhs :: CgStgRhs -> Bool
-isUpdatableRhs (StgRhsClosure _ _ u _ _) = isUpdatable u
-isUpdatableRhs _                         = False
+isUpdatableRhs (StgRhsClosure _ _ u _ _ _) = isUpdatable u
+isUpdatableRhs _                           = False
 
 stgLneLive' :: CgStgBinding -> [Id]
 stgLneLive' b = filter (`notElem` bindees b) (stgLneLive b)
@@ -241,9 +241,9 @@ inspectInlineBinding v = \case
 
 inspectInlineRhs :: UniqSet Id -> Id -> CgStgRhs -> UniqSet Id
 inspectInlineRhs v i = \case
-  StgRhsCon{}                     -> addOneToUniqSet v i
-  StgRhsClosure _ _ ReEntrant _ _ -> addOneToUniqSet v i
-  _                               -> v
+  StgRhsCon{}                       -> addOneToUniqSet v i
+  StgRhsClosure _ _ ReEntrant _ _ _ -> addOneToUniqSet v i
+  _                                 -> v
 
 isInlineForeignCall :: ForeignCall -> Bool
 isInlineForeignCall (CCall (CCallSpec _ cconv safety)) =