{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase        #-}

module GHC.StgToJS.Utils
  ( assignCoerce1
  , assignToExprCtx
  , fixedLayout
  , assocIdExprs
  -- * Unboxable datacon
  , isUnboxableCon
  , isUnboxable
  , isBoolDataCon
  -- * JSRep
  , slotCount
  , varSize
  , typeSize
  , isVoid
  , isMultiVar
  , idJSRep
  , typeJSRep
  , unaryTypeJSRep
  , primRepToJSRep
  , primOrVoidRepToJSRep
  , stackSlotType
  , primRepSize
  , mkArityTag
  -- * References and Ids
  , exprRefs
  , hasExport
  , collectTopIds
  , collectIds
  -- * Live variables
  , LiveVars
  , liveStatic
  , liveVars
  , stgRhsLive
  , stgExprLive
  , isUpdatableRhs
  , stgLneLive'
  , stgLneLiveExpr
  , isInlineExpr
  )
where

import GHC.Prelude

import GHC.StgToJS.Types
import GHC.StgToJS.ExprCtx

import GHC.JS.JStg.Syntax
import GHC.JS.Make
import GHC.JS.Transform

import GHC.Core.DataCon
import GHC.Core.TyCo.Rep hiding (typeSize)
import GHC.Core.TyCon
import GHC.Core.Type hiding (typeSize)

import GHC.Stg.Syntax

import GHC.Tc.Utils.TcType

import GHC.Builtin.Names
import GHC.Builtin.PrimOps (PrimOp(SeqOp), primOpIsReallyInline)

import GHC.Types.RepType
import GHC.Types.Var
import GHC.Types.Var.Set
import GHC.Types.Id
import GHC.Types.Id.Info
import GHC.Types.Unique.FM
import GHC.Types.ForeignCall
import GHC.Types.TyThing
import GHC.Types.Name

import GHC.Utils.Misc
import GHC.Utils.Outputable hiding ((<>))
import GHC.Utils.Panic

import qualified Data.Bits as Bits
import qualified Data.Foldable as F
import qualified Data.Set      as S
import qualified Data.List     as L
import Data.Set (Set)
import Data.Monoid


assignToTypedExprs :: [TypedExpr] -> [JStgExpr] -> JStgStat
assignToTypedExprs tes es =
  assignAllEqual (concatMap typex_expr tes) es

assignTypedExprs :: [TypedExpr] -> [TypedExpr] -> JStgStat
assignTypedExprs tes es =
  let prim_tes = concatMap typex_expr tes
      prim_es  = concatMap typex_expr es
        -- extract the JExprs, effectively unarising a RuntimeRep thing to
        -- multiple VarType-repped things (e.g., AddrRep takes two VarType-regs)
  in assertPpr (equalLength prim_tes prim_es)
               (ppr (map typex_typ tes) $$ ppr (map typex_typ es))
               (assignAllEqual prim_tes prim_es)

assignToExprCtx :: ExprCtx -> [JStgExpr] -> JStgStat
assignToExprCtx ctx es = assignToTypedExprs (ctxTarget ctx) es

-- | Assign first expr only (if it exists), performing coercions between some
-- PrimReps (e.g. StablePtr# and Addr#).
assignCoerce1 :: [TypedExpr] -> [TypedExpr] -> JStgStat
assignCoerce1 [x] [y] = assignCoerce x y
assignCoerce1 []  []  = mempty
-- We silently ignore the case of an empty list on the first argument. It denotes
-- "assign nothing to n empty slots on the right". Usually this case shouldn't come
-- up, but rare cases where the earlier code can't correctly guess the size of type
-- classes causes slots to be allocated when they aren't needed.
assignCoerce1 []  _   = mempty
assignCoerce1 x y     = pprPanic "assignCoerce1"
                          (vcat [ text "lengths do not match"
                                -- FIXME: Outputable instance removed until JStg replaces JStat
                                , ppr x
                                , ppr y
                                ])

-- | Assign p2 to p1 with optional coercion
assignCoerce :: TypedExpr -> TypedExpr -> JStgStat
-- Coercion between StablePtr# and Addr#
assignCoerce (TypedExpr AddrRep [a_val, a_off]) (TypedExpr (BoxedRep (Just Unlifted)) [sptr]) = mconcat
    [ a_val |= var "h$stablePtrBuf"
    , a_off |= sptr
    ]
assignCoerce (TypedExpr (BoxedRep (Just Unlifted)) [sptr]) (TypedExpr AddrRep [_a_val, a_off]) =
  sptr |= a_off
assignCoerce p1 p2 = assignTypedExprs [p1] [p2]


--------------------------------------------------------------------------------
--                        Core Utils
--------------------------------------------------------------------------------

-- | can we unbox C x to x, only if x is represented as a Number
isUnboxableCon :: DataCon -> Bool
isUnboxableCon dc
  | [t] <- dataConRepArgTys dc
  , [t1] <- typeJSRep (scaledThing t)
  = isUnboxable t1 &&
    dataConTag dc == 1 &&
    length (tyConDataCons $ dataConTyCon dc) == 1
  | otherwise = False

-- | one-constructor types with one primitive field represented as a JS Number
-- can be unboxed
isUnboxable :: JSRep -> Bool
isUnboxable DoubleV = True
isUnboxable IntV    = True -- includes Char#
isUnboxable _       = False

-- | Number of slots occupied by a PrimRep
data SlotCount
  = NoSlot
  | OneSlot
  | TwoSlots
  deriving (Show,Eq,Ord)

instance Outputable SlotCount where
  ppr = text . show

-- | Return SlotCount as an Int
slotCount :: SlotCount -> Int
slotCount = \case
  NoSlot   -> 0
  OneSlot  -> 1
  TwoSlots -> 2


-- | Number of slots occupied by a value with the given JSRep
varSize :: JSRep -> Int
varSize = slotCount . jsRepSlots

jsRepSlots :: JSRep -> SlotCount
jsRepSlots VoidV = NoSlot
jsRepSlots LongV = TwoSlots -- hi, low
jsRepSlots AddrV = TwoSlots -- obj/array, offset
jsRepSlots _     = OneSlot

typeSize :: Type -> Int
typeSize t = sum . map varSize . typeJSRep $ t

isVoid :: JSRep -> Bool
isVoid VoidV = True
isVoid _     = False

isMultiVar :: JSRep -> Bool
isMultiVar v = case jsRepSlots v of
  NoSlot   -> False
  OneSlot  -> False
  TwoSlots -> True

idJSRep :: HasDebugCallStack => Id -> [JSRep]
idJSRep = typeJSRep . idType

typeJSRep :: HasDebugCallStack => Type -> [JSRep]
typeJSRep t = map primRepToJSRep (typePrimRep t)

-- only use if you know it's not an unboxed tuple
unaryTypeJSRep :: HasDebugCallStack => UnaryType -> JSRep
unaryTypeJSRep ut = primOrVoidRepToJSRep (typePrimRep1 ut)

primRepToJSRep :: HasDebugCallStack => PrimRep -> JSRep
primRepToJSRep (BoxedRep _) = PtrV
primRepToJSRep IntRep       = IntV
primRepToJSRep Int8Rep      = IntV
primRepToJSRep Int16Rep     = IntV
primRepToJSRep Int32Rep     = IntV
primRepToJSRep WordRep      = IntV
primRepToJSRep Word8Rep     = IntV
primRepToJSRep Word16Rep    = IntV
primRepToJSRep Word32Rep    = IntV
primRepToJSRep Int64Rep     = LongV
primRepToJSRep Word64Rep    = LongV
primRepToJSRep AddrRep      = AddrV
primRepToJSRep FloatRep     = DoubleV
primRepToJSRep DoubleRep    = DoubleV
primRepToJSRep (VecRep{})   = error "primRepToJSRep: vector types are unsupported"

primOrVoidRepToJSRep :: HasDebugCallStack => PrimOrVoidRep -> JSRep
primOrVoidRepToJSRep VoidRep = VoidV
primOrVoidRepToJSRep (NVRep rep) = primRepToJSRep rep

dataConType :: DataCon -> Type
dataConType dc = idType (dataConWrapId dc)

isBoolDataCon :: DataCon -> Bool
isBoolDataCon dc = isBoolTy (dataConType dc)

-- standard fixed layout: payload types
-- payload starts at .d1 for heap objects, entry closest to Sp for stack frames
fixedLayout :: [JSRep] -> CILayout
fixedLayout vts = CILayoutFixed (sum (map varSize vts)) vts

-- 2-var values might have been moved around separately, use DoubleV as substitute
-- ObjV is 1 var, so this is no problem for implicit metadata
stackSlotType :: Id -> JSRep
stackSlotType i
  | OneSlot <- jsRepSlots otype = otype
  | otherwise                   = DoubleV
  where otype = unaryTypeJSRep (idType i)

idPrimReps :: Id -> [PrimRep]
idPrimReps = typePrimReps . idType

typePrimReps :: Type -> [PrimRep]
typePrimReps = typePrimRep . unwrapType

primRepSize :: PrimRep -> SlotCount
primRepSize p = jsRepSlots (primRepToJSRep p)

-- | Associate the given values to each RrimRep in the given order, taking into
-- account the number of slots per PrimRep
assocPrimReps :: [PrimRep] -> [JStgExpr] -> [(PrimRep, [JStgExpr])]
assocPrimReps []     _  = []
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 $ map jStgExprToJS <$> err)

-- | Associate the given values to the Id's PrimReps, taking into account the
-- number of slots per PrimRep
assocIdPrimReps :: Id -> [JStgExpr] -> [(PrimRep, [JStgExpr])]
assocIdPrimReps i = assocPrimReps (idPrimReps i)

-- | Associate the given JExpr to the Id's PrimReps, taking into account the
-- number of slots per PrimRep
assocIdExprs :: Id -> [JStgExpr] -> [TypedExpr]
assocIdExprs i es = fmap (uncurry TypedExpr) (assocIdPrimReps i es)

mkArityTag :: Int -> Int -> Int
mkArityTag arity registers = arity Bits..|. (registers `Bits.shiftL` 8)

--------------------------------------------------------------------------------
--                        Stg Utils
--------------------------------------------------------------------------------

s :: a -> Set a
s = S.singleton

l :: (a -> Set Id) -> [a] -> Set Id
l = F.foldMap

-- | collect Ids that this binding refers to
--   (does not include the bindees themselves)
-- first argument is Id -> StgExpr map for unfloated arguments
bindingRefs :: UniqFM Id CgStgExpr -> CgStgBinding -> Set Id
bindingRefs u = \case
  StgNonRec _ rhs -> rhsRefs u rhs
  StgRec bs       -> l (rhsRefs u . snd) bs

rhsRefs :: UniqFM Id CgStgExpr -> CgStgRhs -> Set Id
rhsRefs u = \case
  StgRhsClosure _ _ _ _ body _       -> exprRefs u body
  StgRhsCon _ccs d _mu _ticks args _ -> l s [ i | AnId i <- dataConImplicitTyThings d] <> l (argRefs u) args

exprRefs :: UniqFM Id CgStgExpr -> CgStgExpr -> Set Id
exprRefs u = \case
  StgApp f args             -> s f <> l (argRefs u) args
  StgConApp d _n args _     -> l s [ i | AnId i <- dataConImplicitTyThings d] <> l (argRefs u) args
  StgOpApp _ args _         -> l (argRefs u) args
  StgLit {}                 -> mempty
  StgCase expr _ _ alts     -> exprRefs u expr <> mconcat (fmap (altRefs u) alts)
  StgLet _ bnd expr         -> bindingRefs u bnd <> exprRefs u expr
  StgLetNoEscape _ bnd expr -> bindingRefs u bnd <> exprRefs u expr
  StgTick _ expr            -> exprRefs u expr

altRefs :: UniqFM Id CgStgExpr -> CgStgAlt -> Set Id
altRefs u alt = exprRefs u (alt_rhs alt)

argRefs :: UniqFM Id CgStgExpr -> StgArg -> Set Id
argRefs u = \case
  StgVarArg id
    | Just e <- lookupUFM u id -> exprRefs u e
    | otherwise                -> s id
  _ -> mempty

hasExport :: CgStgBinding -> Bool
hasExport bnd =
  case bnd of
    StgNonRec b e -> isExportedBind b e
    StgRec bs     -> any (uncurry isExportedBind) bs
  where
    isExportedBind _i (StgRhsCon _cc con _ _ _ _) =
      getUnique con == staticPtrDataConKey
    isExportedBind _ _ = False

collectTopIds :: CgStgBinding -> [Id]
collectTopIds (StgNonRec b _) = [b]
collectTopIds (StgRec bs) = let xs = map (zapFragileIdInfo . fst) bs
                            in  seqList xs `seq` xs

collectIds :: UniqFM Id CgStgExpr -> CgStgBinding -> [Id]
collectIds unfloated b =
  let xs = map zapFragileIdInfo .
           filter acceptId $ S.toList (bindingRefs unfloated b)
  in  seqList xs `seq` xs
  where
    acceptId i = all ($ i) [not . isForbidden] -- fixme test this: [isExported[isGlobalId, not.isForbidden]
    isForbidden i
      -- the GHC.Prim module has no js source file
      | Just m <- nameModule_maybe (getName i)
      , m == gHC_PRIM
      = True
      -- unboxed tuples have no definition
      | isUnboxedTupleDataConLikeName (getName i)
      = True
      | otherwise
      = False

-----------------------------------------------------
-- Live vars
--
-- TODO: should probably be moved into GHC.Stg.LiveVars

type LiveVars = DVarSet

liveStatic :: LiveVars -> LiveVars
liveStatic = filterDVarSet isGlobalId

liveVars :: LiveVars -> LiveVars
liveVars = filterDVarSet (not . isGlobalId)

stgBindLive :: CgStgBinding -> [(Id, LiveVars)]
stgBindLive = \case
  StgNonRec b rhs -> [(b, stgRhsLive rhs)]
  StgRec bs       -> map (\(b,rhs) -> (b, stgRhsLive rhs)) bs

stgBindRhsLive :: CgStgBinding -> LiveVars
stgBindRhsLive b =
  let (bs, ls) = unzip (stgBindLive b)
  in  delDVarSetList (unionDVarSets ls) bs

stgRhsLive :: CgStgRhs -> LiveVars
stgRhsLive = \case
  StgRhsClosure _ _ _ args e _ -> delDVarSetList (stgExprLive True e) args
  StgRhsCon _ _ _ _ args _     -> unionDVarSets (map stgArgLive args)

stgArgLive :: StgArg -> LiveVars
stgArgLive = \case
  StgVarArg occ -> unitDVarSet occ
  StgLitArg {}  -> emptyDVarSet

stgExprLive :: Bool -> CgStgExpr -> LiveVars
stgExprLive includeLHS = \case
  StgApp occ args -> unionDVarSets (unitDVarSet occ : map stgArgLive args)
  StgLit {}       -> emptyDVarSet
  StgConApp _dc _n args _tys -> unionDVarSets (map stgArgLive args)
  StgOpApp _op args _ty      -> unionDVarSets (map stgArgLive args)
  StgCase e b _at alts
    | includeLHS -> el `unionDVarSet` delDVarSet al b
    | otherwise  -> delDVarSet al b
    where
      al = unionDVarSets (map stgAltLive alts)
      el = stgExprLive True e
  StgLet _ b e         -> delDVarSetList (stgBindRhsLive b `unionDVarSet` stgExprLive True e) (bindees b)
  StgLetNoEscape _ b e -> delDVarSetList (stgBindRhsLive b `unionDVarSet` stgExprLive True e) (bindees b)
  StgTick _ti e        -> stgExprLive True e

stgAltLive :: CgStgAlt -> LiveVars
stgAltLive alt =
  delDVarSetList (stgExprLive True (alt_rhs alt)) (alt_bndrs alt)

bindees :: CgStgBinding -> [Id]
bindees = \case
  StgNonRec b _e -> [b]
  StgRec bs      -> map fst bs

isUpdatableRhs :: CgStgRhs -> Bool
isUpdatableRhs (StgRhsClosure _ _ u _ _ _) = isUpdatable u
isUpdatableRhs _                           = False

stgLneLive' :: CgStgBinding -> [Id]
stgLneLive' b = filter (`notElem` bindees b) (stgLneLive b)

stgLneLive :: CgStgBinding -> [Id]
stgLneLive (StgNonRec _b e) = stgLneLiveExpr e
stgLneLive (StgRec bs)      = L.nub $ concatMap (stgLneLiveExpr . snd) bs

stgLneLiveExpr :: CgStgRhs -> [Id]
stgLneLiveExpr rhs = dVarSetElems (liveVars $ stgRhsLive rhs)
-- stgLneLiveExpr (StgRhsClosure _ _ _ _ e) = dVarSetElems (liveVars (stgExprLive e))
-- stgLneLiveExpr StgRhsCon {}              = []

-- | returns True if the expression is definitely inline
isInlineExpr :: CgStgExpr -> Bool
isInlineExpr = \case
  StgApp i args
    -> isInlineApp i args
  StgLit{}
    -> True
  StgConApp{}
    -> True
  StgOpApp (StgFCallOp f _) _ _
    -> isInlineForeignCall f
  StgOpApp (StgPrimOp SeqOp) [StgVarArg e] t
    -> ctxIsEvaluated e || isStrictType t
  StgOpApp (StgPrimOp op) _ _
    -> primOpIsReallyInline op
  StgOpApp (StgPrimCallOp _c) _ _
    -> True
  StgCase e _ _ alts
    ->let ie   = isInlineExpr e
          ias  = map isInlineExpr (fmap alt_rhs alts)
      in ie && and ias
  StgLet _ _ e
    -> isInlineExpr e
  StgLetNoEscape _ _ e
    -> isInlineExpr e
  StgTick _ e
    -> isInlineExpr e

isInlineForeignCall :: ForeignCall -> Bool
isInlineForeignCall (CCall (CCallSpec _ cconv safety)) =
  not (playInterruptible safety) &&
  not (cconv /= JavaScriptCallConv && playSafe safety)

isInlineApp :: Id -> [StgArg] -> Bool
isInlineApp i = \case
  _ | isJoinId i -> False
  [] -> isUnboxedTupleType (idType i) ||
                     isStrictType (idType i) ||
                     ctxIsEvaluated i

  [StgVarArg a]
    | DataConWrapId dc <- idDetails i
    , isNewTyCon (dataConTyCon dc)
    , isStrictType (idType a) || ctxIsEvaluated a || isStrictId a
    -> True
  _ -> False