{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE LambdaCase #-}

module GHC.StgToJS.Sinker.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.Sinker.Collect
import GHC.StgToJS.Sinker.StringsUnfloat

import GHC.Utils.Misc (partitionWith)
import GHC.StgToJS.Utils

import Data.Char
import Data.List (partition)
import Data.Maybe
import Data.ByteString (ByteString)

-- | 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 :: CgStgTopBinding -> Either CgStgBinding (Id, ByteString)
    selectLifted (StgTopLifted b)      = Left b
    selectLifted (StgTopStringLit i b) = Right (i, b)

    (pgm', allStringLits) = partitionWith selectLifted pgm
    usedOnceIds = selectUsedOnce $ concatMap collectArgs pgm'

    stringLitsUFM = listToUFM $ (\(i, b) -> (idName i, (i, b))) <$> allStringLits
    (pgm'', _actuallyUnfloatedStringLitNames) =
      unfloatStringLits
        (idName `mapUniqSet` usedOnceIds)
        (snd `mapUFM` stringLitsUFM)
        pgm'

    stringLits = uncurry StgTopStringLit <$> allStringLits

    (sunk, pgm''') = sinkPgm' m usedOnceIds pgm''

sinkPgm'
  :: Module
       -- ^ the module, since we treat definitions from the current module
       -- differently
  -> IdSet
       -- ^ the set of used once ids
  -> [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 usedOnceIds pgm =
  let usedOnce = collectTopLevelUsedOnce usedOnceIds pgm
      sinkables = listToUFM $
          concatMap alwaysSinkable pgm ++
          concatMap (filter ((`elementOfUniqSet` usedOnce) . fst) . 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] _typ
    | 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 _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
--   and never anywhere else
collectTopLevelUsedOnce :: IdSet -> [CgStgBinding] -> IdSet
collectTopLevelUsedOnce usedOnceIds binds = intersectUniqSets usedOnceIds (selectUsedOnce top_args)
  where
    top_args = concatMap collectArgsTop binds

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 _typ)) =
      [ (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)