-- GRIN-like backend for Yhc Core.
-- State monad for the converter from Yhc Core.

module Yhc.Core.GRIN.GGMonad (
   GG (..)
  ,GGM
  ,getCnt
  ,pushFun
  ,popFun
  ,newVar
  ,newGVar
  ,mapVarName
  ,module Control.Monad.State) where

import Data.Maybe
import Yhc.Core.Extra
import qualified Data.Map as M
import qualified Data.Set as S
import Control.Monad
import Control.Monad.State
import Yhc.Core.GRIN.Type
import Yhc.Core.GRIN.SubstVars

data GG = GG {
   stateCnt :: Int                                -- counter to generate unique names
  ,coreRef :: Core                                -- reference to core
  ,coreAnno :: CoreAnnotations                    -- annotations for this core
  ,currFun :: CoreFunc                            -- current function being compiled
  ,funStack :: [CoreFunc]                         -- stack of functions (to make cf2gf recursive)
  ,varMap :: M.Map (CoreFuncName, CoreVarName) GName -- map to rename variables from the program
  ,funMap :: CoreFuncMap                          -- map of Core functions for faster lookup
  ,liftSet :: S.Set GName                         -- set of prepared lifting functions
  ,autoFuncs :: [GFunc]                           -- automatically created functions
  ,heapSeed :: HeapMap                            -- heap map seed obtained from Core annotations
}

getCnt :: GGM Int
                
getCnt = do     
  c <- gets stateCnt
  st <- get       
  put st {stateCnt = c + 1}
  return c

pushFun :: CoreFunc -> GGM ()

pushFun nf = do
  stk <- gets funStack
  cf <- gets currFun
  let stk' = cf : stk
  st <- get
  put st {currFun = nf, funStack = stk'}
  return ()

popFun :: GGM ()

popFun = do
  (cf:stk) <- gets funStack
  st <- get
  put st {currFun = cf, funStack = stk}
  return ()
  

-- Create a new variable name by using a unique numbers generator.

newVar = do
  c <- getCnt
  return $ "n" ++ show c

newGVar = newVar >>= return . gVar


-- Obtain or add a variable name mapping.

mapVarName vn = do
  mp <- gets varMap
  cf <- gets currFun >>= return . coreFuncName
  case M.lookup (cf, vn) mp of
    Just mn -> return mn
    Nothing -> do
      nv <- newVar
      let mp' = M.insert (cf, vn) nv mp
      st <- get
      put st {varMap = mp'}
      return nv


type GGM a = State GG a