{-# LANGUAGE CPP, ScopedTypeVariables, TypeFamilies #-}
{-# LANGUAGE DataKinds #-}

module GHC.Stg.Utils
    ( mkStgAltTypeFromStgAlts
    , bindersOf, bindersOfX, bindersOfTop, bindersOfTopBinds

    , stripStgTicksTop, stripStgTicksTopE
    , idArgs

    , mkUnarisedId, mkUnarisedIds
    ) where

import GHC.Prelude

import GHC.Types.Id
import GHC.Core.Type
import GHC.Core.TyCon
import GHC.Core.DataCon
import GHC.Core ( AltCon(..) )
import GHC.Types.Tickish
import GHC.Types.Unique.Supply

import GHC.Types.RepType
import GHC.Stg.Syntax

import GHC.Utils.Outputable

import GHC.Utils.Panic.Plain
import GHC.Utils.Panic

import GHC.Data.FastString

mkUnarisedIds :: MonadUnique m => FastString -> [UnaryType] -> m [Id]
mkUnarisedIds :: forall (m :: * -> *).
MonadUnique m =>
FastString -> [UnaryType] -> m [Id]
mkUnarisedIds FastString
fs [UnaryType]
tys = (UnaryType -> m Id) -> [UnaryType] -> m [Id]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (FastString -> UnaryType -> m Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> UnaryType -> m Id
mkUnarisedId FastString
fs) [UnaryType]
tys

mkUnarisedId :: MonadUnique m => FastString -> UnaryType -> m Id
mkUnarisedId :: forall (m :: * -> *).
MonadUnique m =>
FastString -> UnaryType -> m Id
mkUnarisedId FastString
s UnaryType
t = FastString -> UnaryType -> UnaryType -> m Id
forall (m :: * -> *).
MonadUnique m =>
FastString -> UnaryType -> UnaryType -> m Id
mkSysLocalM FastString
s UnaryType
Many UnaryType
t

-- Checks if id is a top level error application.
-- isErrorAp_maybe :: Id ->

-- | Extract the default case alternative
-- findDefaultStg :: [Alt b] -> ([Alt b], Maybe (Expr b))
findDefaultStg
  :: [GenStgAlt p]
  -> ([GenStgAlt p], Maybe (GenStgExpr p))
findDefaultStg :: forall (p :: StgPass).
[GenStgAlt p] -> ([GenStgAlt p], Maybe (GenStgExpr p))
findDefaultStg (GenStgAlt{ alt_con :: forall (pass :: StgPass). GenStgAlt pass -> AltCon
alt_con    = AltCon
DEFAULT
                         , alt_bndrs :: forall (pass :: StgPass). GenStgAlt pass -> [BinderP pass]
alt_bndrs  = [BinderP p]
args
                         , alt_rhs :: forall (pass :: StgPass). GenStgAlt pass -> GenStgExpr pass
alt_rhs    = GenStgExpr p
rhs} : [GenStgAlt p]
alts) = Bool
-> ([GenStgAlt p], Maybe (GenStgExpr p))
-> ([GenStgAlt p], Maybe (GenStgExpr p))
forall a. HasCallStack => Bool -> a -> a
assert( [BinderP p] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [BinderP p]
args ) ([GenStgAlt p]
alts, GenStgExpr p -> Maybe (GenStgExpr p)
forall a. a -> Maybe a
Just GenStgExpr p
rhs)
findDefaultStg [GenStgAlt p]
alts                                  = ([GenStgAlt p]
alts, Maybe (GenStgExpr p)
forall a. Maybe a
Nothing)

mkStgAltTypeFromStgAlts :: forall p. Id -> [GenStgAlt p] -> AltType
mkStgAltTypeFromStgAlts :: forall (p :: StgPass). Id -> [GenStgAlt p] -> AltType
mkStgAltTypeFromStgAlts Id
bndr [GenStgAlt p]
alts
  | UnaryType -> Bool
isUnboxedTupleType UnaryType
bndr_ty Bool -> Bool -> Bool
|| UnaryType -> Bool
isUnboxedSumType UnaryType
bndr_ty
  = Int -> AltType
MultiValAlt ([PrimRep] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimRep]
prim_reps)  -- always use MultiValAlt for unboxed tuples

  | Bool
otherwise
  = case [PrimRep]
prim_reps of
      [PrimRep
rep] | PrimRep -> Bool
isGcPtrRep PrimRep
rep ->
        case UnaryType -> Maybe TyCon
tyConAppTyCon_maybe (UnaryType -> UnaryType
unwrapType UnaryType
bndr_ty) of
          Just TyCon
tc
            | TyCon -> Bool
isAbstractTyCon TyCon
tc -> AltType
look_for_better_tycon
            | TyCon -> Bool
isAlgTyCon TyCon
tc      -> TyCon -> AltType
AlgAlt TyCon
tc
            | Bool
otherwise          -> Bool -> SDoc -> AltType -> AltType
forall a. HasCallStack => Bool -> SDoc -> a -> a
assertPpr ( TyCon -> Bool
_is_poly_alt_tycon TyCon
tc) (TyCon -> SDoc
forall a. Outputable a => a -> SDoc
ppr TyCon
tc)
                                    AltType
PolyAlt
          Maybe TyCon
Nothing                -> AltType
PolyAlt
      [PrimRep
non_gcd] -> PrimRep -> AltType
PrimAlt PrimRep
non_gcd
      [PrimRep]
not_unary -> Int -> AltType
MultiValAlt ([PrimRep] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PrimRep]
not_unary)
  where
   bndr_ty :: UnaryType
bndr_ty   = Id -> UnaryType
idType Id
bndr
   prim_reps :: [PrimRep]
prim_reps = (() :: Constraint) => UnaryType -> [PrimRep]
UnaryType -> [PrimRep]
typePrimRep UnaryType
bndr_ty

   _is_poly_alt_tycon :: TyCon -> Bool
_is_poly_alt_tycon TyCon
tc
        =  TyCon -> Bool
isFunTyCon TyCon
tc
        Bool -> Bool -> Bool
|| TyCon -> Bool
isPrimTyCon TyCon
tc   -- "Any" is lifted but primitive
        Bool -> Bool -> Bool
|| TyCon -> Bool
isFamilyTyCon TyCon
tc -- Type family; e.g. Any, or arising from strict
                            -- function application where argument has a
                            -- type-family type

   -- Sometimes, the TyCon is a AbstractTyCon which may not have any
   -- constructors inside it.  Then we may get a better TyCon by
   -- grabbing the one from a constructor alternative
   -- if one exists.
   look_for_better_tycon :: AltType
look_for_better_tycon
        | (DataAlt DataCon
con : [AltCon]
_) <- GenStgAlt p -> AltCon
forall (pass :: StgPass). GenStgAlt pass -> AltCon
alt_con (GenStgAlt p -> AltCon) -> [GenStgAlt p] -> [AltCon]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [GenStgAlt p]
data_alts =
                TyCon -> AltType
AlgAlt (DataCon -> TyCon
dataConTyCon DataCon
con)
        | Bool
otherwise =
                Bool -> AltType -> AltType
forall a. HasCallStack => Bool -> a -> a
assert([GenStgAlt p] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [GenStgAlt p]
data_alts)
                AltType
PolyAlt
        where
                ([GenStgAlt p]
data_alts, Maybe (GenStgExpr p)
_deflt) = [GenStgAlt p] -> ([GenStgAlt p], Maybe (GenStgExpr p))
forall (p :: StgPass).
[GenStgAlt p] -> ([GenStgAlt p], Maybe (GenStgExpr p))
findDefaultStg [GenStgAlt p]
alts

bindersOf :: BinderP a ~ Id => GenStgBinding a -> [Id]
bindersOf :: forall (a :: StgPass). (BinderP a ~ Id) => GenStgBinding a -> [Id]
bindersOf (StgNonRec BinderP a
binder GenStgRhs a
_) = [Id
BinderP a
binder]
bindersOf (StgRec [(BinderP a, GenStgRhs a)]
pairs)       = [Id
binder | (Id
binder, GenStgRhs a
_) <- [(Id, GenStgRhs a)]
[(BinderP a, GenStgRhs a)]
pairs]

bindersOfX :: GenStgBinding a -> [BinderP a]
bindersOfX :: forall (a :: StgPass). GenStgBinding a -> [BinderP a]
bindersOfX (StgNonRec BinderP a
binder GenStgRhs a
_) = [BinderP a
binder]
bindersOfX (StgRec [(BinderP a, GenStgRhs a)]
pairs)       = [BinderP a
binder | (BinderP a
binder, GenStgRhs a
_) <- [(BinderP a, GenStgRhs a)]
pairs]

bindersOfTop :: BinderP a ~ Id => GenStgTopBinding a -> [Id]
bindersOfTop :: forall (a :: StgPass).
(BinderP a ~ Id) =>
GenStgTopBinding a -> [Id]
bindersOfTop (StgTopLifted GenStgBinding a
bind) = GenStgBinding a -> [Id]
forall (a :: StgPass). (BinderP a ~ Id) => GenStgBinding a -> [Id]
bindersOf GenStgBinding a
bind
bindersOfTop (StgTopStringLit Id
binder ByteString
_) = [Id
binder]

-- All ids we bind something to on the top level.
bindersOfTopBinds :: BinderP a ~ Id => [GenStgTopBinding a] -> [Id]
-- bindersOfTopBinds binds = mapUnionVarSet (mkVarSet . bindersOfTop) binds
bindersOfTopBinds :: forall (a :: StgPass).
(BinderP a ~ Id) =>
[GenStgTopBinding a] -> [Id]
bindersOfTopBinds [GenStgTopBinding a]
binds = (GenStgTopBinding a -> [Id] -> [Id])
-> [Id] -> [GenStgTopBinding a] -> [Id]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ([Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
(++) ([Id] -> [Id] -> [Id])
-> (GenStgTopBinding a -> [Id])
-> GenStgTopBinding a
-> [Id]
-> [Id]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GenStgTopBinding a -> [Id]
forall (a :: StgPass).
(BinderP a ~ Id) =>
GenStgTopBinding a -> [Id]
bindersOfTop) [] [GenStgTopBinding a]
binds

idArgs :: [StgArg] -> [Id]
idArgs :: [StgArg] -> [Id]
idArgs [StgArg]
args = [Id
v | StgVarArg Id
v <- [StgArg]
args]

-- | Strip ticks of a given type from an STG expression.
stripStgTicksTop :: (StgTickish -> Bool) -> GenStgExpr p -> ([StgTickish], GenStgExpr p)
stripStgTicksTop :: forall (p :: StgPass).
(StgTickish -> Bool)
-> GenStgExpr p -> ([StgTickish], GenStgExpr p)
stripStgTicksTop StgTickish -> Bool
p = [StgTickish] -> GenStgExpr p -> ([StgTickish], GenStgExpr p)
go []
   where go :: [StgTickish] -> GenStgExpr p -> ([StgTickish], GenStgExpr p)
go [StgTickish]
ts (StgTick StgTickish
t GenStgExpr p
e) | StgTickish -> Bool
p StgTickish
t = [StgTickish] -> GenStgExpr p -> ([StgTickish], GenStgExpr p)
go (StgTickish
tStgTickish -> [StgTickish] -> [StgTickish]
forall a. a -> [a] -> [a]
:[StgTickish]
ts) GenStgExpr p
e
         -- This special case avoid building a thunk for "reverse ts" when there are no ticks
         go [] GenStgExpr p
other               = ([], GenStgExpr p
other)
         go [StgTickish]
ts GenStgExpr p
other               = ([StgTickish] -> [StgTickish]
forall a. [a] -> [a]
reverse [StgTickish]
ts, GenStgExpr p
other)

-- | Strip ticks of a given type from an STG expression returning only the expression.
stripStgTicksTopE :: (StgTickish -> Bool) -> GenStgExpr p -> GenStgExpr p
stripStgTicksTopE :: forall (p :: StgPass).
(StgTickish -> Bool) -> GenStgExpr p -> GenStgExpr p
stripStgTicksTopE StgTickish -> Bool
p = GenStgExpr p -> GenStgExpr p
go
   where go :: GenStgExpr p -> GenStgExpr p
go (StgTick StgTickish
t GenStgExpr p
e) | StgTickish -> Bool
p StgTickish
t = GenStgExpr p -> GenStgExpr p
go GenStgExpr p
e
         go GenStgExpr p
other               = GenStgExpr p
other