{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module GHC.StgToJS.Closure
  ( closureInfoStat
  , closure
  , conClosure
  , Closure (..)
  , newClosure
  , assignClosure
  , CopyCC (..)
  , copyClosure
  , mkClosure
  -- $names
  , allocData
  , allocClsA
  , dataName
  , clsName
  , dataFieldName
  , varName
  , jsClosureCount
  )
where

import GHC.Prelude
import GHC.Data.FastString

import GHC.StgToJS.Heap
import GHC.StgToJS.Types
import GHC.StgToJS.Utils
import GHC.StgToJS.Regs (stack,sp)

import GHC.JS.Make
import GHC.JS.JStg.Syntax
import GHC.JS.JStg.Monad
import GHC.JS.Ident

import GHC.Types.Unique.Map

import Data.Array
import Data.Monoid
import qualified Data.Bits as Bits

closureInfoStat :: Bool -> ClosureInfo -> JStgStat
closureInfoStat debug (ClosureInfo obj rs name layout ctype srefs)
  = setObjInfoL debug obj rs layout ty name tag srefs
      where
        !ty = case ctype of
          CIThunk      -> Thunk
          CIFun {}     -> Fun
          CICon {}     -> Con
          CIBlackhole  -> Blackhole
          CIPap        -> Pap
          CIStackFrame -> StackFrame
        !tag = case ctype of
          CIThunk           -> 0
          CIFun arity nregs -> mkArityTag arity nregs
          CICon con         -> con
          CIBlackhole       -> 0
          CIPap             -> 0
          CIStackFrame      -> 0


setObjInfoL :: Bool        -- ^ debug: output symbol names
            -> Ident       -- ^ the object name
            -> CIRegs      -- ^ things in registers
            -> CILayout    -- ^ layout of the object
            -> ClosureType -- ^ closure type
            -> FastString  -- ^ object name, for printing
            -> Int         -- ^ `a' argument, depends on type (arity, conid)
            -> CIStatic    -- ^ static refs
            -> JStgStat
setObjInfoL debug obj rs layout t n a
  = setObjInfo debug obj t n field_types a size rs
      where
        size = case layout of
          CILayoutVariable   -> (-1)
          CILayoutUnknown sz -> sz
          CILayoutFixed sz _ -> sz
        field_types = case layout of
          CILayoutVariable     -> []
          CILayoutUnknown size -> to_type_list (replicate size ObjV)
          CILayoutFixed _ fs   -> to_type_list fs
        to_type_list = concatMap (\x -> replicate (varSize x) (fromEnum x))

setObjInfo :: Bool        -- ^ debug: output all symbol names
           -> Ident       -- ^ the thing to modify
           -> ClosureType -- ^ closure type
           -> FastString  -- ^ object name, for printing
           -> [Int]       -- ^ list of item types in the object, if known (free variables, datacon fields)
           -> Int         -- ^ extra 'a' parameter, for constructor tag or arity
           -> Int         -- ^ object size, -1 (number of vars) for unknown
           -> CIRegs      -- ^ things in registers
           -> CIStatic    -- ^ static refs
           -> JStgStat
setObjInfo debug obj t name fields a size regs static
   | debug     = appS "h$setObjInfo" [ toJExpr obj
                                     , toJExpr t
                                     , toJExpr name
                                     , toJExpr fields
                                     , toJExpr a
                                     , toJExpr size
                                     , toJExpr (regTag regs)
                                     , toJExpr static
                                     ]
   | otherwise = appS "h$o" [ toJExpr obj
                            , toJExpr t
                            , toJExpr a
                            , toJExpr size
                            , toJExpr (regTag regs)
                            , toJExpr static
                            ]
  where
    regTag CIRegsUnknown       = -1
    regTag (CIRegs skip types) =
      let nregs = sum $ fmap varSize types
      in  skip + (nregs `Bits.shiftL` 8)

-- | Special case of closures that do not need to generate any @fresh@ names
closure :: ClosureInfo    -- ^ object being info'd see @ciVar@
         -> (JSM JStgStat) -- ^ rhs
         -> JSM JStgStat
closure ci body = do f <- (jFunction' (ciVar ci) body)
                     return $ f `mappend` closureInfoStat False ci

conClosure :: Ident -> FastString -> CILayout -> Int -> JSM JStgStat
conClosure symbol name layout constr = closure ci body
  where
    ci = (ClosureInfo symbol (CIRegs 0 [PtrV]) name layout (CICon constr) mempty)
    body   = pure . returnS $ stack .! sp

-- | Used to pass arguments to newClosure with some safety
data Closure = Closure
  { clEntry  :: JStgExpr
  , clField1 :: JStgExpr
  , clField2 :: JStgExpr
  , clMeta   :: JStgExpr
  , clCC     :: Maybe JStgExpr
  }

newClosure :: Closure -> JStgExpr
newClosure Closure{..} =
  let xs = [ (closureEntry_ , clEntry)
           , (closureField1_, clField1)
           , (closureField2_, clField2)
           , (closureMeta_  , clMeta)
           ]
  in case clCC of
    -- CC field is optional (probably to minimize code size as we could assign
    -- null_, but we get the same effect implicitly)
    Nothing -> ValExpr (jhFromList xs)
    Just cc -> ValExpr (jhFromList $ (closureCC_,cc) : xs)

assignClosure :: JStgExpr -> Closure -> JStgStat
assignClosure t Closure{..} = BlockStat
  [ closureEntry  t |= clEntry
  , closureField1 t |= clField1
  , closureField2 t |= clField2
  , closureMeta   t |= clMeta
  ] <> case clCC of
      Nothing -> mempty
      Just cc -> closureCC t |= cc

data CopyCC = CopyCC | DontCopyCC

copyClosure :: CopyCC -> JStgExpr -> JStgExpr -> JStgStat
copyClosure copy_cc t s = BlockStat
  [ closureEntry  t |= closureEntry  s
  , closureField1 t |= closureField1 s
  , closureField2 t |= closureField2 s
  , closureMeta   t |= closureMeta   s
  ] <> case copy_cc of
      DontCopyCC -> mempty
      CopyCC     -> closureCC t |= closureCC s

mkClosure :: JStgExpr -> [JStgExpr] -> JStgExpr -> Maybe JStgExpr -> Closure
mkClosure entry fields meta cc = Closure
  { clEntry  = entry
  , clField1 = x1
  , clField2 = x2
  , clMeta   = meta
  , clCC     = cc
  }
  where
    x1 = case fields of
           []  -> null_
           x:_ -> x
    x2 = case fields of
           []     -> null_
           [_]    -> null_
           [_,x]  -> x
           _:x:xs -> ValExpr . JHash . listToUniqMap $ zip (fmap dataFieldName [1..]) (x:xs)


-------------------------------------------------------------------------------
--                             Name Caches
-------------------------------------------------------------------------------
-- $names

-- | Cache "dXXX" field names
dataFieldCache :: Array Int FastString
dataFieldCache = listArray (0,nFieldCache) (fmap (mkFastString . ('d':) . show) [(0::Int)..nFieldCache])

-- | Data names are used in the AST, and logging has determined that 255 is the maximum number we see.
nFieldCache :: Int
nFieldCache  = 255

-- | We use this in the RTS to determine the number of generated closures. These closures use the names
-- cached here, so we bind them to the same number.
jsClosureCount :: Int
jsClosureCount  = 24

dataFieldName :: Int -> FastString
dataFieldName i
  | i < 0 || i > nFieldCache = mkFastString ('d' : show i)
  | otherwise                = dataFieldCache ! i

-- | Cache "h$dXXX" names
dataCache :: Array Int FastString
dataCache = listArray (0,jsClosureCount) (fmap (mkFastString . ("h$d"++) . show) [(0::Int)..jsClosureCount])

dataName :: Int -> FastString
dataName i
  | i < 0 || i > nFieldCache = mkFastString ("h$d" ++ show i)
  | otherwise                = dataCache ! i

allocData :: Int -> JStgExpr
allocData i = toJExpr (global (dataName i))

-- | Cache "h$cXXX" names
clsCache :: Array Int FastString
clsCache = listArray (0,jsClosureCount) (fmap (mkFastString . ("h$c"++) . show) [(0::Int)..jsClosureCount])

clsName :: Int -> FastString
clsName i
  | i < 0 || i > jsClosureCount = mkFastString ("h$c" ++ show i)
  | otherwise                   = clsCache ! i

allocClsA :: Int -> JStgExpr
allocClsA i = toJExpr (global (clsName i))

-- | Cache "xXXX" names
varCache :: Array Int Ident
varCache = listArray (0,jsClosureCount) (fmap (global . mkFastString . ('x':) . show) [(0::Int)..jsClosureCount])

varName :: Int -> Ident
varName i
  | i < 0 || i > jsClosureCount = global $ mkFastString ('x' : show i)
  | otherwise                   = varCache ! i