-- {-# LANGUAGE MagicHash #-}
-- {-# OPTIONS_GHC -O2 #-}
-- {-# OPTIONS_GHC -O3 #-}

module UHC.Light.Compiler.CoreRun.Run.Val.RunExplStk
( cmodRun )
where
import UHC.Light.Compiler.Base.HsName.Builtin
import UHC.Light.Compiler.Base.Common
import UHC.Light.Compiler.Opts
import UHC.Light.Compiler.Ty
import UHC.Light.Compiler.Error
import UHC.Light.Compiler.Gam
import UHC.Light.Compiler.Gam.DataGam
import UHC.Light.Compiler.Base.Trace
import UHC.Light.Compiler.CoreRun
import UHC.Light.Compiler.CoreRun.Run
import UHC.Light.Compiler.CoreRun.Run.Val
import UHC.Light.Compiler.CoreRun.Run.Val.Prim
import UHC.Light.Compiler.CoreRun.Pretty
import UHC.Util.Pretty
import UHC.Util.Lens
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.List
import qualified Data.ByteString.Char8 as BSC8

{-# LINE 48 "src/ehc/CoreRun/Run/Val/RunExplStk.chs" #-}
-- | Arguments to a function, which may come from an RVal_App or from the stack
data ExplArgs = ExplArgs
  { eaVec		:: !RValV		-- ^ the accumulated part from RVal_App
  , eaStk		:: !Int			-- ^ the size of the part still on the stack
  }

emptyExplArgs = ExplArgs V.empty 0
-- {-# INLINE emptyExplArgs #-}

-- | The total nr of args
eaNrArgs :: ExplArgs -> Int
eaNrArgs (ExplArgs {eaVec=v, eaStk=na}) = V.length v + na
{-# INLINE eaNrArgs #-}

-- | Set total nr of args, taking into account what is in the vector part
eaSetNrArgs :: ExplArgs -> Int -> ExplArgs
eaSetNrArgs ea@(ExplArgs {eaVec=v}) n = ea {eaStk = n - V.length v}
{-# INLINE eaSetNrArgs #-}

-- | Pop from the ExplArgs partly embedded in the top frame and partly explicitly available
renvFrStkEaPopMV :: RunSem RValCxt RValEnv RVal m x => ExplArgs -> RValT m RValMV
renvFrStkEaPopMV ea@(ExplArgs {eaVec=v}) = (liftIO $ mvecAlloc eaLen) >>= \vs -> liftIO (mvecFillFromV 0 vs v) >> renvFrStkReversePopInMV vLen (eaLen-vLen) vs >> return vs
  where vLen  = V.length v
        eaLen = eaNrArgs ea
{-# INLINE renvFrStkEaPopMV #-}

{-# LINE 76 "src/ehc/CoreRun/Run/Val/RunExplStk.chs" #-}
-- | Allocate a new frame
explStkAllocFrameM
  :: (RunSem RValCxt RValEnv RVal m x)
  => Ref2Nm					-- ^ ref <-> name mapping
  -> RCxt					-- ^ context
  -> Int					-- ^ size
  -> ExplArgs				-- ^ arguments
  -> RValT m HpPtr
explStkAllocFrameM r2n cx sz as@(ExplArgs {eaVec=vsArgs, eaStk=nrArgs}) = do
  a <- liftIO $ mvecAllocInit sz -- (sz+3)		-- TBD: stack overflow somewhere...
  let vsLen = V.length vsArgs
  when (vsLen  > 0) $ liftIO $ mvecFillFromV 0 a vsArgs
  when (nrArgs > 0) $ renvFrStkReversePopInMV vsLen nrArgs a
  spref <- liftIO $ newIORef (eaNrArgs as)
  p <- heapAllocM $ RVal_Frame r2n cx a spref
  return p

-- | Push a new stack frame
explStkPushFrameM :: (RunSem RValCxt RValEnv RVal m x) => HpPtr -> RValT m ()
explStkPushFrameM frptr = do
  (RValEnv {renvStack=st, renvTopFrame=tfref}) <- get
  liftIO $ do
    tf <- readIORef tfref
    unless (isNullPtr tf) $ modifyIORef st (tf:)
    writeIORef tfref frptr
{-# INLINE explStkPushFrameM #-}

-- | Allocate and push a new stack frame
explStkPushAllocFrameM
  :: (RunSem RValCxt RValEnv RVal m x)
  => Ref2Nm					-- ^ ref <-> name mapping
  -> RCxt					-- ^ context
  -> Int					-- ^ size
  -> ExplArgs				-- ^ arguments
  -> RValT m HpPtr
explStkPushAllocFrameM r2n cx sz as = do
  p <- explStkAllocFrameM r2n cx sz as
  explStkPushFrameM p
  return p
{-# INLINE explStkPushAllocFrameM #-}

-- | Allocate and replace top stack frame
explStkReplaceAllocFrameM
  :: (RunSem RValCxt RValEnv RVal m x)
  => Ref2Nm					-- ^ ref <-> name mapping
  -> RCxt					-- ^ context
  -> Int					-- ^ size
  -> ExplArgs				-- ^ arguments
  -> RValT m ()
explStkReplaceAllocFrameM r2n cx sz as = do
  p <- explStkAllocFrameM r2n cx sz as
  (RValEnv {renvTopFrame=tf}) <- get
  liftIO $ writeIORef tf p
{-# INLINE explStkReplaceAllocFrameM #-}

-- | Pop a stack frame, copying the top of the stack embedded in the frame
explStkPopFrameM :: (RunSem RValCxt RValEnv RVal m x) => RValT m HpPtr
explStkPopFrameM = do
  (RValEnv {renvStack=stref, renvTopFrame=tfref}) <- get
  liftIO $ do
    tf  <- readIORef tfref
    stk <- readIORef stref
    case stk of
      [] -> writeIORef tfref nullPtr
      (h:t) -> do
        writeIORef tfref h
        writeIORef stref t
    return tf
{-# INLINE explStkPopFrameM #-}

{-# LINE 152 "src/ehc/CoreRun/Run/Val/RunExplStk.chs" #-}
cmodRun :: (RunSem RValCxt RValEnv RVal m ()) => EHCOpts -> Mod -> RValT m ()
cmodRun opts (Mod_Mod {mbbody_Mod_Mod = Just e}) = do
  -- dumpEnvM True
  mustReturn $ rsemExp e
  -- v <- renvFrStkPop1
  -- return v
cmodRun opts _ = return ()

{-# LINE 170 "src/ehc/CoreRun/Run/Val/RunExplStk.chs" #-}
-- | Apply Lam in context of static link with exact right amount of params, otherwise the continuation is used
rvalExplStkAppLam :: RunSem RValCxt RValEnv RVal m () => RCxt -> Exp -> ExplArgs -> (Int -> RValT m ()) -> RValT m ()
rvalExplStkAppLam cx f as failcont = do
  let nrActualArgs = eaNrArgs as
  case f of
    Exp_Lam {{- lev_Exp_Lam=l, -} mbNm_Exp_Lam=mn, nrArgs_Exp_Lam=nrRequiredArgs, stkDepth_Exp_Lam=sz, ref2nm_Exp_Lam=r2n, body_Exp_Lam=b}
      | nrActualArgs == nrRequiredArgs -> do
           -- rsemTr $ ">V (" ++ show mn ++ ") app lam ==, na=" ++ show nrRequiredArgs ++ ", sz=" ++ show sz
           needRet <- asks rcxtInRet
           rvalTrEnterLam mn $
             if needRet
               then do
                 explStkPushAllocFrameM r2n cx sz as
                 rsemExp b
                 v <- renvFrStkPop1
                 explStkPopFrameM
                 renvFrStkPush1 v
               else do
                 explStkReplaceAllocFrameM r2n cx sz as
                 mustReturn $ rsemExp b
           -- rsemTr $ "<V (" ++ show mn ++ ")"
      | otherwise -> failcont nrRequiredArgs
    _   -> err $ "CoreRun.Run.Val.rvalExplStkAppLam:" >#< f
-- {-# SPECIALIZE rvalExplStkAppLam :: HpPtr -> Exp -> RValMV -> (Int -> RValT IO RVal) -> RValT IO RVal #-}
-- {-# INLINE rvalExplStkAppLam #-}

{-# LINE 198 "src/ehc/CoreRun/Run/Val/RunExplStk.chs" #-}
-- | Apply. Assume: function 'f' is already evaluated (responsibility lies outside)
rvalExplStkApp :: RunSem RValCxt RValEnv RVal m () => RVal -> ExplArgs -> RValT m ()
rvalExplStkApp f as = do
  -- rsemTr $ "V app f(" ++ show (MV.length as) ++ "): " ++ show (pp f)
  let nrActualArgs = eaNrArgs as
  case f of
    RVal_Lam {rvalCx=rcx, rvalBody=b} -> do
      -- sl <- liftIO $ readIORef (rcxtSlRef rcx)
      rvalExplStkAppLam rcx b as $ \narg -> do
        if nrActualArgs < narg
          then do
            renvFrStkEaPopMV as >>= \as -> heapAllocAsPtrM (RVal_App f as) >>= renvFrStkPush1
          else do
            ap <- mustReturn $ rvalExplStkApp f (eaSetNrArgs as narg) >>= rsemPop >>= rsemDeref >>= rsemPop
            rvalExplStkApp ap (eaSetNrArgs emptyExplArgs (nrActualArgs - narg))
    RVal_App appf appas
      | nrActualArgs > 0 -> do
           appas' <- liftIO $ V.freeze appas
           rvalExplStkApp appf (as {eaVec=appas' V.++ eaVec as})
    _   -> err $ "CoreRun.Run.Val.rvalExplStkApp:" >#< f
-- {-# SPECIALIZE rvalExplStkApp :: RunSem RValCxt RValEnv RVal IO RVal => RVal -> RValMV -> RValT IO RVal #-}
-- {-# INLINE rvalExplStkApp #-}

{-# LINE 223 "src/ehc/CoreRun/Run/Val/RunExplStk.chs" #-}
-- | rsemExp for RVal, without explicit use of expr stack, i.e. implicit stack via Haskell thereby preventing correct GC
rvalExplStkExp :: RunSem RValCxt RValEnv RVal m () => Exp -> RValT m ()
{-# SPECIALIZE rvalExplStkExp :: RunSem RValCxt RValEnv RVal IO () => Exp -> RValT IO () #-}
-- {-# INLINE rvalExplStkExp #-}
rvalExplStkExp e = do
  rsemTr'' TraceOn_RunEval $ ">E:" >#< e
  -- e' <- case e of
  case e of
    -- app, call
    Exp_App f as -> do
        vecReverseForM_ as rsemSExp
        f' <- mustReturn $ rsemExp f
        rsemPop f' >>= ptr2valM >>= \f' -> rvalExplStkApp f' (emptyExplArgs {eaStk=V.length as})

    -- heap node
    Exp_Tup t as -> do
        V.forM_ as rsemSExp
        renvFrStkPopMV (V.length as) >>= rsemNode t >>= rsemPush

    -- lam as is, being a heap allocated thunk when 0 args are required
    Exp_Lam {nrArgs_Exp_Lam=na, mbNm_Exp_Lam=mn}
      | na == 0   -> mk heapAllocAsPtrM RVal_Thunk
      | otherwise -> mk return          RVal_Lam
      where mk alloc rv = do
             (sl,fr) <- renvTopFramePtrAndFrameM
             cx <- liftIO $ rcxtCloneWithNewFrame sl (rvalCx fr)
             alloc (rv mn e cx) >>= rsemPush

    -- let
    Exp_Let {firstOff_Exp_Let=fillFrom, ref2nm_Exp_Let=r2n, binds_Exp_Let=bs, body_Exp_Let=b} -> do
        mustReturn $ V.forM_ bs rsemExp
        rsemExp b

    -- case, scrutinee already evaluated
    Exp_Case e as -> do
      v <- ptr2valM =<< rsemPop =<< rsemSExp e
      case v of
        -- RVal_NodeMV {rvalTag=tg} -> rsemAlt $ as V.! tg
        RVal_Int  tg           -> rsemAlt $ as V.! tg
        _ -> err $ "CoreRun.Run.Val.RunExplStk.rvalExplStkExp.Case: scrutinee:" >#< v

    -- force evaluation immediately
    Exp_Force e -> rsemExp e >>= rsemPop >>= rsemEvl

    -- setup for context requiring a return (TBD: should be done via CPS style, but is other issue)
    -- Exp_Ret e -> mustReturn $ rsemExp e

    -- setup for context requiring a return from case alternative
    -- Exp_RetCase _ e -> rsemExp e

    -- setup for context not requiring a return
    Exp_Tail e -> needNotReturn $ rsemExp e

    -- simple expressions
    Exp_SExp se -> rsemSExp se

    -- FFI
    Exp_FFI pr as -> V.mapM_ rsemSExp as >> renvFrStkPopMV (V.length as) >>= (liftIO . V.freeze) >>= rsemPrim pr

    -- necessary only when case is non-saturated w.r.t. alternatives of datatype Exp
    -- e -> err $ "CoreRun.Run.Val.RunExplStk.rvalExplStkExp:" >#< e

  rsemTr'' TraceOn_RunEval $ "<E:" >#< (e) -- >-< e')
  -- return e'

{-# LINE 302 "src/ehc/CoreRun/Run/Val/RunExplStk.chs" #-}
-- | Add module
rvalExplAddModule :: RunSem RValCxt RValEnv RVal m () => Mod -> RValT m HpPtr
rvalExplAddModule mod@(Mod_Mod {moduleNm_Mod_Mod=nm, ref2nm_Mod_Mod=r2n, binds_Mod_Mod=bs, stkDepth_Mod_Mod=sz, imports_Mod_Mod=imports}) = do
    rsemTr'' TraceOn_RunMod $ ">rvalExplAddModule:" >#< nm
    -- add new entry
    env@(RValEnv {renvModulesMV=mods}) <- get
    let nr = MV.length mods
    -- add new entry
    mods' <- liftIO $ MV.grow mods 1
    -- frame IORef with dummy ptr
    fr <- liftIO $ newIORef nullPtr
    m  <- heapAllocM $ RVal_Module nm (crarrayFromList []) fr
    -- write module entry which requires later patching
    liftIO $ MV.write mods' nr m
    -- set as new module array
    put $ env {renvModulesMV = mods'}
    -- get updated modules
    env@(RValEnv {renvModulesMV=mods, renvGlobalsMV=globs}) <- get
    -- get import indirection table
    imptbl <- renvResolveModNames (nr-1) [ nm | Import_Import {nm_Import_Import=nm} <- imports ]
    -- construct context (module is patched in later)
    cx <- liftIO $ mkRCxt nullPtr nullPtr
    -- construct frame
    f <- explStkPushAllocFrameM r2n cx sz emptyExplArgs
    -- patch frame ref with frame
    liftIO $ writeIORef fr f
    -- patch module with imp table
    heapUpdM m (\m -> return $ m {rvalModImpsV=crarrayFromList imptbl})
    -- patch module ref in context
    liftIO $ writeIORef (rcxtMdRef cx) m
    -- compute module bindings into current frame
    V.forM_ bs rsemExp
    rsemTr'' TraceOn_RunMod $ "<rvalExplAddModule:" >#< nm >#< "-> modhpptr=" >|< m
    -- remove the frame
    f <- explStkPopFrameM
    -- and return it
    return f
  where

{-# LINE 343 "src/ehc/CoreRun/Run/Val/RunExplStk.chs" #-}
instance
    ( Monad m, MonadIO m, Functor m
    ) => RunSem RValCxt RValEnv RVal m ()
  where
    {-# SPECIALIZE instance RunSem RValCxt RValEnv RVal IO () #-}
    rsemInitial = do
      s <- liftIO $ newRValEnv 100000 -- 100000 --
      return (emptyRValCxt, s, undefined)

    rsemSetup opts modImpL mod {- @(Mod_Mod {moduleNr_Mod_Mod=mainModNr}) -} = do
{-
-}
        rsemSetupTracing opts
        let modAllL = modImpL ++ [mod]
            updTr   = rcxtTraceOnS ^= ehcOptTraceOn opts
        modFrames <- local updTr $ forM modAllL $ \mod -> do
          rvalExplAddModule mod
        explStkPushFrameM (last modFrames)
        fmap updTr $ rcxtUpdDatatypes modAllL
{-
        -- rsemSetTrace True
        rsemGcEnterRootLevel
        let modAllL = modImpL ++ [mod]
        ms <- liftIO $ MV.new (maximum (map moduleNr_Mod_Mod modAllL) + 1)
        modify $ \env -> env {renvGlobalsMV = ms}
        forM_ modAllL $ \(Mod_Mod {ref2nm_Mod_Mod=r2n, moduleNr_Mod_Mod=nr, binds_Mod_Mod=bs, stkDepth_Mod_Mod=sz}) -> do
          -- context (ignoring module stuff, TBD)
          cx <- liftIO $ mkRCxtSl nullPtr
          -- construct frame for each module
          p <- explStkPushAllocFrameM r2n cx sz emptyExplArgs
          -- and store the frame into the array holding module frames
          (liftIO $ MV.write ms nr p >> newIORef p) >>= \r -> rsemGcPushRoot (RVal_Ptr r)
          -- holding all local defs
          V.forM_ bs rsemExp
          -- p <-
          explStkPopFrameM
        -- use the main module's stackframe for evaluating 'main'
        liftIO (MV.read ms mainModNr) >>= explStkPushFrameM
        rsemGcLeaveRootLevel
        rsemSetupTracing opts
        rcxtUpdDatatypes modAllL
-}

    rsemSetTrace doTrace doExtensive = modify $ \env ->
      env {renvDoTrace = doTrace, renvDoTraceExt = doExtensive}

    rsemTraceOnS = asks _rcxtTraceOnS

    rsemExp = rvalExplStkExp

    rsemSExp se = do
      case se of
        SExp_Int 	v -> rsemPush $ RVal_Int v
        SExp_Char 	v -> rsemPush $ RVal_Char v
        SExp_Var    r -> do v <- ref2valM r
                            -- rsemTr $ "R->V:" >#< v
                            rsemPush v
        SExp_String v -> rsemPush $ RVal_PackedString $ BSC8.pack v
        _ -> rsemPush (RVal_Lit se)
    {-# INLINE rsemSExp #-}

    rsemEvl v = do
        case v of
          RVal_Ptr {rvalPtrRef=pref} -> do
            rsemGcEnterRootLevel
            rsemGcPushRoot v
            liftIO (readIORef pref) >>= evlPtr pref
            rsemGcLeaveRootLevel
          RVal_BlackHole             -> err $ "CoreRun.Run.Val.rsemEvl.RVal_BlackHole:" >#< "Black hole"
          _                          -> return () -- rsemPush v
        rsemPush v
      where
        evlPtr pref p = do
          hp <- gets renvHeap
          v <- heapGetM' hp p
          case v of
            RVal_Thunk {rvalMbNm=mn, rvalCx=rcx, rvalBody=e} -> do
              -- rsemGcPushRoot v
              -- sl <- liftIO $ readIORef (rcxtSlRef rcx)
              heapSetM' hp p RVal_BlackHole
              v' <- rvalExplStkAppLam rcx e (emptyExplArgs {eaStk=0}) $ \_ -> err $ "CoreRun.Run.Val.rsemEvl.RVal_Thunk:" >#< e
              hp <- gets renvHeap
              p <- liftIO (readIORef pref)
              v'' <- rsemPop v'
              heapSetM' hp p v''
              return v''
            RVal_Ptr {rvalPtrRef=pref} -> do
              v' <- evlPtr pref =<< liftIO (readIORef pref)
              hp <- gets renvHeap
              p <- liftIO (readIORef pref)
              heapSetM' hp p v'
              return v'
            v -> do
              return v

    rsemDeref v = do
      v' <- ptr2valM v
      -- rsemTr $ "Deref:" >#< (v >-< v')
      rsemPush v'
    {-# INLINE rsemDeref #-}

    -- apply a known primitive
    rsemPrim = rvalPrim
    {-# INLINE rsemPrim #-}

    rsemPush = renvFrStkPush1
    {-# INLINE rsemPush #-}
    rsemPop  = \_ -> renvFrStkPop1
    {-# INLINE rsemPop #-}
    rsemNode t vs = {- heapAllocAsPtrM -} return $ RVal_NodeMV t vs
    {-# INLINE rsemNode #-}


    rsemGcEnterRootLevel = gets renvGcRootStack >>= \r -> liftIO $ modifyIORef r $ ([]:)
    {-# INLINE rsemGcEnterRootLevel #-}

    rsemGcPushRoot v = gets renvGcRootStack >>= \r -> liftIO $ modifyIORef r $ \(h:t) -> (v:h) : t
    {-# INLINE rsemGcPushRoot #-}

    rsemGcLeaveRootLevel = gets renvGcRootStack >>= \r -> liftIO $ modifyIORef r tail
    {-# INLINE rsemGcLeaveRootLevel #-}