1 files changed, 180 insertions, 0 deletions
diff --git a/compiler/GHC/StgToJS/Sinker.hs b/compiler/GHC/StgToJS/Sinker.hs
new file mode 100644
index 0000000000..6df58d4fcf
--- /dev/null
+++ b/compiler/GHC/StgToJS/Sinker.hs
@@ -0,0 +1,180 @@
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE LambdaCase #-}
+
+module GHC.StgToJS.Sinker (sinkPgm) where
+
+import GHC.Prelude
+import GHC.Types.Unique.Set
+import GHC.Types.Unique.FM
+import GHC.Types.Var.Set
+import GHC.Stg.Syntax
+import GHC.Types.Id
+import GHC.Types.Name
+import GHC.Unit.Module
+import GHC.Types.Literal
+import GHC.Data.Graph.Directed
+
+import GHC.StgToJS.CoreUtils
+
+import Data.Char
+import Data.Either
+import Data.List (partition)
+import Data.Maybe
+
+
+-- | Unfloat some top-level unexported things
+--
+-- GHC floats constants to the top level. This is fine in native code, but with JS
+-- they occupy some global variable name. We can unfloat some unexported things:
+--
+-- - global constructors, as long as they're referenced only once by another global
+--      constructor and are not in a recursive binding group
+-- - literals (small literals may also be sunk if they are used more than once)
+sinkPgm :: Module
+        -> [CgStgTopBinding]
+        -> (UniqFM Id CgStgExpr, [CgStgTopBinding])
+sinkPgm m pgm = (sunk, map StgTopLifted pgm'' ++ stringLits)
+  where
+    selectLifted (StgTopLifted b) = Left b
+    selectLifted x                = Right x
+    (pgm', stringLits) = partitionEithers (map selectLifted pgm)
+    (sunk, pgm'')      = sinkPgm' m pgm'
+
+sinkPgm'
+  :: Module
+       -- ^ the module, since we treat definitions from the current module
+       -- differently
+  -> [CgStgBinding]
+       -- ^ the bindings
+  -> (UniqFM Id CgStgExpr, [CgStgBinding])
+       -- ^ a map with sunken replacements for nodes, for where the replacement
+       -- does not fit in the 'StgBinding' AST and the new bindings
+sinkPgm' m pgm =
+  let usedOnce = collectUsedOnce pgm
+      sinkables = listToUFM $
+          concatMap alwaysSinkable pgm ++
+          filter ((`elementOfUniqSet` usedOnce) . fst) (concatMap (onceSinkable m) pgm)
+      isSunkBind (StgNonRec b _e) | elemUFM b sinkables = True
+      isSunkBind _                                      = False
+  in (sinkables, filter (not . isSunkBind) $ topSortDecls m pgm)
+
+-- | always sinkable, values that may be duplicated in the generated code (e.g.
+-- small literals)
+alwaysSinkable :: CgStgBinding -> [(Id, CgStgExpr)]
+alwaysSinkable (StgRec {})       = []
+alwaysSinkable (StgNonRec b rhs) = case rhs of
+  StgRhsClosure _ _ _ _ e@(StgLit l)
+    | isSmallSinkableLit l
+    , isLocal b
+    -> [(b,e)]
+  StgRhsCon _ccs dc cnum _ticks as@[StgLitArg l]
+    | isSmallSinkableLit l
+    , isLocal b
+    , isUnboxableCon dc
+    -> [(b,StgConApp dc cnum as [])]
+  _ -> []
+
+isSmallSinkableLit :: Literal -> Bool
+isSmallSinkableLit (LitChar c)     = ord c < 100000
+isSmallSinkableLit (LitNumber _ i) = abs i < 100000
+isSmallSinkableLit _               = False
+
+
+-- | once sinkable: may be sunk, but duplication is not ok
+onceSinkable :: Module -> CgStgBinding -> [(Id, CgStgExpr)]
+onceSinkable _m (StgNonRec b rhs)
+  | Just e <- getSinkable 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
+onceSinkable _ _ = []
+
+-- | collect all idents used only once in an argument at the top level
+--   and never anywhere else
+collectUsedOnce :: [CgStgBinding] -> IdSet
+collectUsedOnce binds = intersectUniqSets (usedOnce args) (usedOnce top_args)
+  where
+    top_args = concatMap collectArgsTop binds
+    args     = concatMap collectArgs    binds
+    usedOnce = fst . foldr g (emptyUniqSet, emptyUniqSet)
+    g i t@(once, mult)
+      | i `elementOfUniqSet` mult = t
+      | i `elementOfUniqSet` once
+        = (delOneFromUniqSet once i, addOneToUniqSet mult i)
+      | otherwise = (addOneToUniqSet once i, mult)
+
+-- | fold over all id in StgArg used at the top level in an StgRhsCon
+collectArgsTop :: CgStgBinding -> [Id]
+collectArgsTop = \case
+  StgNonRec _b r -> collectArgsTopRhs r
+  StgRec bs      -> concatMap (collectArgsTopRhs . snd) bs
+
+collectArgsTopRhs :: CgStgRhs -> [Id]
+collectArgsTopRhs = \case
+  StgRhsCon _ccs _dc _mu _ticks args -> concatMap collectArgsA args
+  StgRhsClosure {}                   -> []
+
+-- | fold over all Id in StgArg in the AST
+collectArgs :: CgStgBinding -> [Id]
+collectArgs = \case
+  StgNonRec _b r -> collectArgsR r
+  StgRec bs      -> concatMap (collectArgsR . snd) bs
+
+collectArgsR :: CgStgRhs -> [Id]
+collectArgsR = \case
+  StgRhsClosure _x0 _x1 _x2 _x3 e     -> collectArgsE e
+  StgRhsCon _ccs _con _mu _ticks args -> concatMap collectArgsA args
+
+collectArgsAlt :: CgStgAlt -> [Id]
+collectArgsAlt alt = collectArgsE (alt_rhs alt)
+
+collectArgsE :: CgStgExpr -> [Id]
+collectArgsE = \case
+  StgApp x args
+    -> x : concatMap collectArgsA args
+  StgConApp _con _mn args _ts
+    -> concatMap collectArgsA args
+  StgOpApp _x args _t
+    -> concatMap collectArgsA args
+  StgCase e _b _a alts
+    -> collectArgsE e ++ concatMap collectArgsAlt alts
+  StgLet _x b e
+    -> collectArgs b ++ collectArgsE e
+  StgLetNoEscape _x b e
+    -> collectArgs b ++ collectArgsE e
+  StgTick _i e
+    -> collectArgsE e
+  StgLit _
+    -> []
+
+collectArgsA :: StgArg -> [Id]
+collectArgsA = \case
+  StgVarArg i -> [i]
+  StgLitArg _ -> []
+
+isLocal :: Id -> Bool
+isLocal i = isNothing (nameModule_maybe . idName $ i) && not (isExportedId i)
+
+-- | since we have sequential initialization, topsort the non-recursive
+-- constructor bindings
+topSortDecls :: Module -> [CgStgBinding] -> [CgStgBinding]
+topSortDecls _m binds = rest ++ nr'
+  where
+    (nr, rest) = partition isNonRec binds
+    isNonRec StgNonRec{} = True
+    isNonRec _           = False
+    vs   = map getV 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)) =
+      [ (i, b) | StgVarArg i <- args, i `elementOfUniqSet` keys ]
+    collectDeps _ = []
+    g = graphFromVerticesAndAdjacency vs (concatMap collectDeps nr)
+    nr' | (not . null) [()| CyclicSCC _ <- stronglyConnCompG g]
+            = error "topSortDecls: unexpected cycle"
+        | otherwise = map node_payload (topologicalSortG g)