JS: refactor jsSaturate to return a saturated JStat (#23328)

9 files changed, 85 insertions, 103 deletions

diff --git a/compiler/GHC/JS/Transform.hs b/compiler/GHC/JS/Transform.hs
index 9c45ebdd57..3415b04b86 100644
--- a/compiler/GHC/JS/Transform.hs
+++ b/compiler/GHC/JS/Transform.hs
@@ -6,13 +6,15 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE TupleSections #-}
 
 module GHC.JS.Transform
   ( identsS
   , identsV
   , identsE
   -- * Saturation
-  , jsSaturate
+  , satJStat
+  , satJExpr
   -- * Generic traversal (via compos)
   , JMacro(..)
   , JMGadt(..)
@@ -21,8 +23,6 @@ module GHC.JS.Transform
   , composOpM
   , composOpM_
   , composOpFold
-  , satJExpr
-  , satJStat
   )
 where
 
@@ -33,11 +33,12 @@ import GHC.JS.Unsat.Syntax
 
 import Data.Functor.Identity
 import Control.Monad
-import Control.Arrow ((***))
+import Data.List (sortBy)
 
 import GHC.Data.FastString
 import GHC.Utils.Monad.State.Strict
 import GHC.Types.Unique.Map
+import GHC.Types.Unique.FM
 
 
 {-# INLINE identsS #-}
@@ -205,69 +206,62 @@ jmcompos ret app f' v =
 
 -- | Given an optional prefix, fills in all free variable names with a supply
 -- of names generated by the prefix.
-jsSaturate :: (JMacro a) => Maybe FastString -> a -> a
-jsSaturate str x = evalState (runIdentSupply $ jsSaturate_ x) (newIdentSupply str)
-
-jsSaturate_ :: (JMacro a) => a -> IdentSupply a
-jsSaturate_ e = IS $ jfromGADT <$> go (jtoGADT e)
+satJStat :: Maybe FastString -> JStat -> Sat.JStat
+satJStat str x = evalState (jsSaturateS x) (newIdentSupply str)
+
+satJExpr :: Maybe FastString -> JExpr -> Sat.JExpr
+satJExpr str x = evalState (jsSaturateE x) (newIdentSupply str)
+
+jsSaturateS :: JStat -> State [Ident] Sat.JStat
+jsSaturateS  = \case
+  DeclStat i rhs        -> Sat.DeclStat i <$> mapM jsSaturateE rhs
+  ReturnStat e          -> Sat.ReturnStat <$> jsSaturateE e
+  IfStat c t e          -> Sat.IfStat <$> jsSaturateE c <*> jsSaturateS t <*> jsSaturateS e
+  WhileStat is_do c e   -> Sat.WhileStat is_do <$> jsSaturateE c <*> jsSaturateS e
+  ForStat init p step body -> Sat.ForStat <$> jsSaturateS init <*> jsSaturateE p
+                                          <*> jsSaturateS step <*> jsSaturateS body
+  ForInStat is_each i iter body -> Sat.ForInStat is_each i <$> jsSaturateE iter <*> jsSaturateS body
+  SwitchStat struct ps def -> Sat.SwitchStat <$> jsSaturateE struct
+                                             <*> mapM (\(p1, p2) -> (,) <$> jsSaturateE p1 <*> jsSaturateS p2) ps
+                                             <*> jsSaturateS def
+  TryStat t i c f       -> Sat.TryStat <$> jsSaturateS t <*> pure i <*> jsSaturateS c <*> jsSaturateS f
+  BlockStat bs          -> fmap Sat.BlockStat $! mapM jsSaturateS bs
+  ApplStat rator rand   -> Sat.ApplStat <$> jsSaturateE rator <*> mapM jsSaturateE rand
+  UOpStat  rator rand   -> Sat.UOpStat (satJUOp rator) <$> jsSaturateE rand
+  AssignStat lhs rhs    -> Sat.AssignStat <$> jsSaturateE lhs <*> pure Sat.AssignOp <*> jsSaturateE rhs
+  LabelStat lbl stmt    -> Sat.LabelStat lbl <$> jsSaturateS stmt
+  BreakStat m_l         -> return $ Sat.BreakStat $! m_l
+  ContinueStat m_l      -> return $ Sat.ContinueStat $! m_l
+  FuncStat i args body  -> Sat.FuncStat i args <$> jsSaturateS body
+  UnsatBlock us         -> jsSaturateS =<< runIdentSupply us
+
+jsSaturateE :: JExpr -> State [Ident] Sat.JExpr
+jsSaturateE = \case
+  ValExpr v            -> Sat.ValExpr <$> jsSaturateV v
+  SelExpr obj i        -> Sat.SelExpr <$> jsSaturateE obj <*> pure i
+  IdxExpr o i          -> Sat.IdxExpr <$> jsSaturateE o <*> jsSaturateE i
+  InfixExpr op l r     -> Sat.InfixExpr (satJOp op) <$> jsSaturateE l <*> jsSaturateE r
+  UOpExpr op r         -> Sat.UOpExpr (satJUOp op) <$> jsSaturateE r
+  IfExpr c t e         -> Sat.IfExpr <$> jsSaturateE c <*> jsSaturateE t <*> jsSaturateE e
+  ApplExpr rator rands -> Sat.ApplExpr <$> jsSaturateE rator <*> mapM jsSaturateE rands
+  UnsatExpr us         -> jsSaturateE =<< runIdentSupply us
+
+jsSaturateV :: JVal -> State [Ident] Sat.JVal
+jsSaturateV = \case
+  JVar i   -> return $ Sat.JVar i
+  JList xs -> Sat.JList <$> mapM jsSaturateE xs
+  JDouble d -> return $ Sat.JDouble (Sat.SaneDouble (unSaneDouble d))
+  JInt i    -> return $ Sat.JInt   i
+  JStr s    -> return $ Sat.JStr   s
+  JRegEx f  -> return $ Sat.JRegEx f
+  JHash m   -> Sat.JHash <$> mapUniqMapM satHash m
     where
-      go :: forall a. JMGadt a -> State [Ident] (JMGadt a)
-      go v = case v of
-               JMGStat (UnsatBlock us) -> go =<< (JMGStat <$> runIdentSupply us)
-               JMGExpr (UnsatExpr  us) -> go =<< (JMGExpr <$> runIdentSupply us)
-               JMGVal  (UnsatVal   us) -> go =<< (JMGVal  <$> runIdentSupply us)
-               _ -> composOpM go v
-
-
---------------------------------------------------------------------------------
---                            Translation
---
---------------------------------------------------------------------------------
-satJStat :: JStat -> Sat.JStat
-satJStat = witness . proof
-  where proof = jsSaturate Nothing
-
-        -- This is an Applicative but we can't use it because no type variables :(
-        witness :: JStat -> Sat.JStat
-        witness (DeclStat i rhs)      = Sat.DeclStat i (fmap satJExpr rhs)
-        witness (ReturnStat e)        = Sat.ReturnStat (satJExpr e)
-        witness (IfStat c t e)        = Sat.IfStat (satJExpr c) (witness t) (witness e)
-        witness (WhileStat is_do c e) = Sat.WhileStat is_do (satJExpr c) (witness e)
-        witness (ForStat init p step body) = Sat.ForStat
-                                             (witness init) (satJExpr p)
-                                             (witness step) (witness body)
-        witness (ForInStat is_each i iter body) = Sat.ForInStat is_each i
-                                                  (satJExpr iter)
-                                                  (witness body)
-        witness (SwitchStat struct ps def) = Sat.SwitchStat
-                                             (satJExpr struct)
-                                             (map (satJExpr *** witness) ps)
-                                             (witness def)
-        witness (TryStat t i c f)     = Sat.TryStat (witness t) i (witness c) (witness f)
-        witness (BlockStat bs)        = Sat.BlockStat $! fmap witness bs
-        witness (ApplStat rator rand) = Sat.ApplStat (satJExpr rator) (satJExpr <$> rand)
-        witness (UOpStat rator rand)  = Sat.UOpStat  (satJUOp rator) (satJExpr rand)
-        witness (AssignStat lhs rhs)  = Sat.AssignStat (satJExpr lhs) Sat.AssignOp (satJExpr rhs)
-        witness (LabelStat lbl stmt)  = Sat.LabelStat lbl (witness stmt)
-        witness (BreakStat Nothing)   = Sat.BreakStat Nothing
-        witness (BreakStat (Just l))  = Sat.BreakStat $! Just l
-        witness (ContinueStat Nothing)  = Sat.ContinueStat Nothing
-        witness (ContinueStat (Just l)) = Sat.ContinueStat $! Just l
-        witness (FuncStat i args body)  = Sat.FuncStat i args (witness body)
-        witness UnsatBlock{}            = error "satJStat: discovered an Unsat...impossibly"
-
-
-satJExpr :: JExpr -> Sat.JExpr
-satJExpr = go
-  where
-    go (ValExpr v)        = Sat.ValExpr (satJVal v)
-    go (SelExpr obj i)    = Sat.SelExpr (satJExpr obj) i
-    go (IdxExpr o i)      = Sat.IdxExpr (satJExpr o) (satJExpr i)
-    go (InfixExpr op l r) = Sat.InfixExpr (satJOp op) (satJExpr l) (satJExpr r)
-    go (UOpExpr op r)     = Sat.UOpExpr (satJUOp op) (satJExpr r)
-    go (IfExpr c t e)     = Sat.IfExpr (satJExpr c) (satJExpr t) (satJExpr e)
-    go (ApplExpr rator rands) = Sat.ApplExpr (satJExpr rator) (satJExpr <$> rands)
-    go UnsatExpr{}        = error "satJExpr: discovered an Unsat...impossibly"
+      satHash (i, x) = (i,) . (i,) <$> jsSaturateE x
+      compareHash (i,_) (j,_) = lexicalCompareFS i j
+      -- By lexically sorting the elements, the non-determinism introduced by nonDetEltsUFM is avoided
+      mapUniqMapM f (UniqMap m) = UniqMap . listToUFM <$> (mapM f . sortBy compareHash $ nonDetEltsUFM m)
+  JFunc args body   -> Sat.JFunc args <$> jsSaturateS body
+  UnsatVal us       -> jsSaturateV =<< runIdentSupply us
 
 satJOp :: JOp -> Sat.Op
 satJOp = go
@@ -313,15 +307,3 @@ satJUOp = go
     go PreDecOp  = Sat.PreDecOp
     go PostDecOp = Sat.PostDecOp
 
-satJVal :: JVal -> Sat.JVal
-satJVal = go
-  where
-    go (JVar i)    = Sat.JVar i
-    go (JList xs)  = Sat.JList (satJExpr <$> xs)
-    go (JDouble d) = Sat.JDouble (Sat.SaneDouble (unSaneDouble d))
-    go (JInt i)    = Sat.JInt   i
-    go (JStr f)    = Sat.JStr   f
-    go (JRegEx f)  = Sat.JRegEx f
-    go (JHash m)   = Sat.JHash (satJExpr <$> m)
-    go (JFunc args body) = Sat.JFunc args (satJStat body)
-    go UnsatVal{} = error "jvalToSatVar: discovered an Sat...impossibly"
diff --git a/compiler/GHC/StgToJS/CodeGen.hs b/compiler/GHC/StgToJS/CodeGen.hs
index 6110135afb..4d557526bd 100644
--- a/compiler/GHC/StgToJS/CodeGen.hs
+++ b/compiler/GHC/StgToJS/CodeGen.hs
@@ -134,10 +134,9 @@ genUnits m ss spt_entries foreign_stubs = do
         glbl <- State.gets gsGlobal
         staticInit <-
           initStaticPtrs spt_entries
-        let stat = ( jsOptimize
-                     . satJStat
-                     . jsSaturate (Just $ modulePrefix m 1)
-                     $ mconcat (reverse glbl) <> staticInit)
+        let stat = ( jsOptimize .
+                     satJStat (Just $ modulePrefix m 1)
+                   $ mconcat (reverse glbl) <> staticInit)
         let syms = [moduleGlobalSymbol m]
         let oi = ObjUnit
                   { oiSymbols  = syms
@@ -210,8 +209,7 @@ genUnits m ss spt_entries foreign_stubs = do
               si        <- State.gets (ggsStatic . gsGroup)
               let body = mempty -- mconcat (reverse extraTl) <> b1 ||= e1 <> b2 ||= e2
               let stat =  jsOptimize
-                          . satJStat
-                          $ jsSaturate (Just $ modulePrefix m n) body
+                          $ satJStat (Just $ modulePrefix m n) body
               let ids = [bnd]
               syms <- (\(TxtI i) -> [i]) <$> identForId bnd
               let oi = ObjUnit
@@ -249,8 +247,7 @@ genUnits m ss spt_entries foreign_stubs = do
               topDeps  = collectTopIds decl
               required = hasExport decl
               stat     = jsOptimize
-                         . satJStat
-                         . jsSaturate (Just $ modulePrefix m n)
+                         . satJStat (Just $ modulePrefix m n)
                          $ mconcat (reverse extraTl) <> tl
           syms <- mapM (fmap (\(TxtI i) -> i) . identForId) topDeps
           let oi = ObjUnit
@@ -339,7 +336,7 @@ genToplevelRhs i rhs = case rhs of
     eid@(TxtI eidt) <- identForEntryId i
     (TxtI idt)   <- identForId i
     body <- genBody (initExprCtx i) R2 args body typ
-    global_occs <- globalOccs (jsSaturate (Just "ghcjs_tmp_sat_") body)
+    global_occs <- globalOccs (satJStat (Just "ghcjs_tmp_sat_") body)
     let lidents = map global_ident global_occs
     let lids    = map global_id    global_occs
     let lidents' = map identFS lidents
diff --git a/compiler/GHC/StgToJS/CoreUtils.hs b/compiler/GHC/StgToJS/CoreUtils.hs
index 751661b11b..fbb89a1339 100644
--- a/compiler/GHC/StgToJS/CoreUtils.hs
+++ b/compiler/GHC/StgToJS/CoreUtils.hs
@@ -253,7 +253,7 @@ assocPrimReps (r:rs) vs = case (primRepSize r,vs) of
   (NoSlot,   xs)     -> (r,[])    : assocPrimReps rs xs
   (OneSlot,  x:xs)   -> (r,[x])   : assocPrimReps rs xs
   (TwoSlots, x:y:xs) -> (r,[x,y]) : assocPrimReps rs xs
-  err                -> pprPanic "assocPrimReps" (ppr $ fmap (map satJExpr) $ err)
+  err                -> pprPanic "assocPrimReps" (ppr $ map (satJExpr Nothing) <$> err)
 
 -- | Associate the given values to the Id's PrimReps, taking into account the
 -- number of slots per PrimRep
diff --git a/compiler/GHC/StgToJS/DataCon.hs b/compiler/GHC/StgToJS/DataCon.hs
index 675fd6d583..5f8444092b 100644
--- a/compiler/GHC/StgToJS/DataCon.hs
+++ b/compiler/GHC/StgToJS/DataCon.hs
@@ -60,8 +60,8 @@ genCon ctx con args
 
   | xs <- concatMap typex_expr (ctxTarget ctx)
   = pprPanic "genCon: unhandled DataCon" (ppr (con
-                                              , fmap satJExpr args
-                                              , fmap satJExpr xs
+                                              , satJExpr Nothing <$> args
+                                              , satJExpr Nothing <$> xs
                                               ))
 
 -- | Allocate a data constructor. Allocate in this context means bind the data
@@ -90,7 +90,7 @@ allocUnboxedCon con = \case
     | isBoolDataCon con && dataConTag con == 2 -> true_
   [x]
     | isUnboxableCon con -> x
-  xs -> pprPanic "allocUnboxedCon: not an unboxed constructor" (ppr (con, fmap satJExpr xs))
+  xs -> pprPanic "allocUnboxedCon: not an unboxed constructor" (ppr (con, satJExpr Nothing <$> xs))
 
 -- | Allocate an entry function. See 'GHC.StgToJS.hs' for the object layout.
 allocDynamicE :: Bool          -- ^ csInlineAlloc from StgToJSConfig
diff --git a/compiler/GHC/StgToJS/Expr.hs b/compiler/GHC/StgToJS/Expr.hs
index 835140791d..09213bcf12 100644
--- a/compiler/GHC/StgToJS/Expr.hs
+++ b/compiler/GHC/StgToJS/Expr.hs
@@ -899,7 +899,7 @@ caseCond = \case
   DataAlt da -> return $ Just (toJExpr $ dataConTag da)
   LitAlt l   -> genLit l >>= \case
     [e] -> pure (Just e)
-    es  -> pprPanic "caseCond: expected single-variable literal" (ppr $ fmap satJExpr es)
+    es  -> pprPanic "caseCond: expected single-variable literal" (ppr $ satJExpr Nothing <$> es)
 
 -- fixme use single tmp var for all branches
 -- | Load parameters from constructor
diff --git a/compiler/GHC/StgToJS/FFI.hs b/compiler/GHC/StgToJS/FFI.hs
index effaa1f122..79409d1719 100644
--- a/compiler/GHC/StgToJS/FFI.hs
+++ b/compiler/GHC/StgToJS/FFI.hs
@@ -14,6 +14,7 @@ import GHC.Prelude
 import GHC.JS.Unsat.Syntax
 import GHC.JS.Make
 import GHC.JS.Transform
+import qualified GHC.JS.Syntax as Sat
 
 import GHC.StgToJS.Arg
 import GHC.StgToJS.ExprCtx
@@ -176,8 +177,8 @@ genFFIArg isJavaScriptCc a@(StgVarArg i)
      arg_ty = stgArgType a
      r      = uTypeVt arg_ty
 
-saturateFFI :: JMacro a => Int -> a -> a
-saturateFFI u = jsSaturate (Just . mkFastString $ "ghcjs_ffi_sat_" ++ show u)
+saturateFFI :: Int -> JStat -> Sat.JStat
+saturateFFI u = satJStat (Just . mkFastString $ "ghcjs_ffi_sat_" ++ show u)
 
 genForeignCall :: HasDebugCallStack
                => ExprCtx
diff --git a/compiler/GHC/StgToJS/Linker/Linker.hs b/compiler/GHC/StgToJS/Linker/Linker.hs
index 067616b1cb..5c6659092a 100644
--- a/compiler/GHC/StgToJS/Linker/Linker.hs
+++ b/compiler/GHC/StgToJS/Linker/Linker.hs
@@ -333,7 +333,7 @@ renderLinker h mods jsFiles = do
     pure (mod_mod, mod_size)
 
   -- commoned up metadata
-  !meta_length <- fromIntegral <$> putJS (jsOptimize $ satJStat meta)
+  !meta_length <- fromIntegral <$> putJS (jsOptimize $ satJStat Nothing meta)
 
   -- module exports
   mapM_ (putBS . cmc_exports) compacted_mods
diff --git a/compiler/GHC/StgToJS/Monad.hs b/compiler/GHC/StgToJS/Monad.hs
index f64275c399..64a5375061 100644
--- a/compiler/GHC/StgToJS/Monad.hs
+++ b/compiler/GHC/StgToJS/Monad.hs
@@ -25,6 +25,7 @@ where
 import GHC.Prelude
 
 import GHC.JS.Unsat.Syntax
+import qualified GHC.JS.Syntax as Sat
 import GHC.JS.Transform
 
 import GHC.StgToJS.Types
@@ -160,7 +161,7 @@ data GlobalOcc = GlobalOcc
 
 -- | Return number of occurrences of every global id used in the given JStat.
 -- Sort by increasing occurrence count.
-globalOccs :: JStat -> G [GlobalOcc]
+globalOccs :: Sat.JStat -> G [GlobalOcc]
 globalOccs jst = do
   GlobalIdCache gidc <- getGlobalIdCache
   -- build a map form Ident Unique to (Ident, Id, Count)
@@ -180,4 +181,4 @@ globalOccs jst = do
               let g = GlobalOcc i gid 1
               in go (addToUFM_C inc gids i g) is
 
-  pure $ go emptyUFM (identsS $ satJStat jst)
+  pure $ go emptyUFM (identsS jst)
diff --git a/compiler/GHC/StgToJS/Rts/Rts.hs b/compiler/GHC/StgToJS/Rts/Rts.hs
index a25cd6c4e7..0d21430602 100644
--- a/compiler/GHC/StgToJS/Rts/Rts.hs
+++ b/compiler/GHC/StgToJS/Rts/Rts.hs
@@ -31,6 +31,7 @@ import GHC.JS.Unsat.Syntax
 import GHC.JS.Make
 import GHC.JS.Transform
 import GHC.JS.Optimizer
+import qualified GHC.JS.Syntax as Sat
 
 import GHC.StgToJS.Apply
 import GHC.StgToJS.Closure
@@ -298,8 +299,8 @@ closureTypes = mconcat (map mkClosureType (enumFromTo minBound maxBound)) <> clo
     ifCT arg ct = jwhenS (arg .===. toJExpr ct) (returnS (toJExpr (show ct)))
 
 -- | JS payload declaring the RTS functions.
-rtsDecls :: JStat
-rtsDecls = jsSaturate (Just "h$RTSD") $
+rtsDecls :: Sat.JStat
+rtsDecls = satJStat (Just "h$RTSD") $
   mconcat [ TxtI "h$currentThread"   ||= null_                   -- thread state object for current thread
           , TxtI "h$stack"           ||= null_                   -- stack for the current thread
           , TxtI "h$sp"              ||= 0                       -- stack pointer for the current thread
@@ -314,15 +315,15 @@ rtsDecls = jsSaturate (Just "h$RTSD") $
 
 -- | print the embedded RTS to a String
 rtsText :: StgToJSConfig -> String
-rtsText = show . pretty . jsOptimize . satJStat . rts
+rtsText = show . pretty . jsOptimize . rts
 
 -- | print the RTS declarations to a String.
 rtsDeclsText :: String
-rtsDeclsText = show . pretty . jsOptimize . satJStat $ rtsDecls
+rtsDeclsText = show . pretty . jsOptimize $ rtsDecls
 
 -- | Wrapper over the RTS to guarentee saturation, see 'GHC.JS.Transform'
-rts :: StgToJSConfig -> JStat
-rts = jsSaturate (Just "h$RTS") . rts'
+rts :: StgToJSConfig -> Sat.JStat
+rts = satJStat (Just "h$RTS") . rts'
 
 -- | JS Payload which defines the embedded RTS.
 rts' :: StgToJSConfig -> JStat