{-# LANGUAGE CPP #-}

module GHC.Stg.DepAnal (depSortStgPgm) where

import GHC.Prelude

import GHC.Stg.Syntax
import GHC.Types.Id
import GHC.Types.Name (Name, nameIsLocalOrFrom)
import GHC.Types.Name.Env
import GHC.Utils.Outputable
import GHC.Types.Unique.Set (nonDetEltsUniqSet)
import GHC.Types.Var.Set
import GHC.Unit.Module (Module)

import Data.Graph (SCC (..))
import Data.Bifunctor (first)

--------------------------------------------------------------------------------
-- * Dependency analysis

-- | Set of bound variables
type BVs = VarSet

-- | Set of free variables
type FVs = VarSet

-- | Dependency analysis on STG terms.
--
-- Dependencies of a binding are just free variables in the binding. This
-- includes imported ids and ids in the current module. For recursive groups we
-- just return one set of free variables which is just the union of dependencies
-- of all bindings in the group.
--
-- Implementation: pass bound variables (BVs) to recursive calls, get free
-- variables (FVs) back. We ignore imported FVs as they do not change the
-- ordering but it improves performance.
--
annTopBindingsDeps :: Module -> [StgTopBinding] -> [(StgTopBinding, FVs)]
annTopBindingsDeps this_mod bs = zip bs (map top_bind bs)
  where
    top_bind :: StgTopBinding -> FVs
    top_bind StgTopStringLit{} =
      emptyVarSet

    top_bind (StgTopLifted bs) =
      binding emptyVarSet bs

    binding :: BVs -> StgBinding -> FVs
    binding bounds (StgNonRec _ r) =
      rhs bounds r
    binding bounds (StgRec bndrs) =
      unionVarSets $
        map (bind_non_rec (extendVarSetList bounds (map fst bndrs))) bndrs

    bind_non_rec :: BVs -> (Id, StgRhs) -> FVs
    bind_non_rec bounds (_, r) =
        rhs bounds r

    rhs :: BVs -> StgRhs -> FVs
    rhs bounds (StgRhsClosure _ _ _ as e) =
      expr (extendVarSetList bounds as) e

    rhs bounds (StgRhsCon _ _ as) =
      args bounds as

    var :: BVs -> Var -> FVs
    var bounds v
      | not (elemVarSet v bounds)
      , nameIsLocalOrFrom this_mod (idName v)
      = unitVarSet v
      | otherwise
      = emptyVarSet

    arg :: BVs -> StgArg -> FVs
    arg bounds (StgVarArg v) = var bounds v
    arg _ StgLitArg{} = emptyVarSet

    args :: BVs -> [StgArg] -> FVs
    args bounds as = unionVarSets (map (arg bounds) as)

    expr :: BVs -> StgExpr -> FVs
    expr bounds (StgApp f as) =
      var bounds f `unionVarSet` args bounds as

    expr _ StgLit{} =
      emptyVarSet

    expr bounds (StgConApp _ as _) =
      args bounds as
    expr bounds (StgOpApp _ as _) =
      args bounds as
    expr _ lam@StgLam{} =
      pprPanic "annTopBindingsDeps" (text "Found lambda:" $$ pprStgExpr panicStgPprOpts lam)
    expr bounds (StgCase scrut scrut_bndr _ as) =
      expr bounds scrut `unionVarSet`
        alts (extendVarSet bounds scrut_bndr) as
    expr bounds (StgLet _ bs e) =
      binding bounds bs `unionVarSet`
        expr (extendVarSetList bounds (bindersOf bs)) e
    expr bounds (StgLetNoEscape _ bs e) =
      binding bounds bs `unionVarSet`
        expr (extendVarSetList bounds (bindersOf bs)) e

    expr bounds (StgTick _ e) =
      expr bounds e

    alts :: BVs -> [StgAlt] -> FVs
    alts bounds = unionVarSets . map (alt bounds)

    alt :: BVs -> StgAlt -> FVs
    alt bounds (_, bndrs, e) =
      expr (extendVarSetList bounds bndrs) e

--------------------------------------------------------------------------------
-- * Dependency sorting

-- | Dependency sort a STG program so that dependencies come before uses.
depSortStgPgm :: Module -> [StgTopBinding] -> [StgTopBinding]
depSortStgPgm this_mod =
    {-# SCC "STG.depSort" #-}
    map fst . depSort . annTopBindingsDeps this_mod

-- | Sort free-variable-annotated STG bindings so that dependencies come before
-- uses.
depSort :: [(StgTopBinding, FVs)] -> [(StgTopBinding, FVs)]
depSort = concatMap get_binds . depAnal defs uses
  where
    uses, defs :: (StgTopBinding, FVs) -> [Name]

    -- TODO (osa): I'm unhappy about two things in this code:
    --
    --     * Why do we need Name instead of Id for uses and dependencies?
    --     * Why do we need a [Name] instead of `Set Name`? Surely depAnal
    --       doesn't need any ordering.

    uses (StgTopStringLit{}, _) = []
    uses (StgTopLifted{}, fvs)  = map idName (nonDetEltsUniqSet fvs)

    defs (bind, _) = map idName (bindersOfTop bind)

    get_binds (AcyclicSCC bind) =
      [bind]
    get_binds (CyclicSCC binds) =
      pprPanic "depSortStgBinds"
               (text "Found cyclic SCC:"
               $$ ppr (map (first (pprStgTopBinding panicStgPprOpts)) binds))