-- | Module containing code that is shared by Interpreter and
-- TypeChecker

module Common where

import Control.Monad
import Data.Map (Map)
import Data.Monoid ((<>))
import qualified Data.Map as M
import CPP.Abs
import CPP.ErrM

type Env s a = (Sig s, [Context a])
type Sig s     = Map Id s
type Context a = Map Id a

-- | Looks up `Type` of `Id` in the environment
lookupVar :: Monad m => Env s a -> Id -> m a
lookupVar (_, (c:cs)) id' =
  case M.lookup id' c of
    Nothing -> lookupVar (M.empty, cs) id'
    Just v  -> return v

lookupVar (_,[]) id' = fail $ "use of undeclared variable " <> show id'

lookupFun :: Monad m => Env s a -> Id -> m s
lookupFun (sig,_) id' =
  case M.lookup id' sig of
    Nothing -> fail $ "use of undeclared function " <> show id'
    Just f  -> return f

containsVar :: Env s a -> Id -> Bool
containsVar (sig, cs) id' = foldr (\c b -> M.member id' c || b) False cs

-- | Extends the variable context with a variable and a value
extendVar :: Monad m => Env sig val -> Id -> val -> m (Env sig val)
extendVar (_, []) _ _       = fail $ "cannot extend a non-existing context"
extendVar (sig, c:cs) id' a =
  case M.lookup id' c of
    Just _  -> fail $ "declaration of already declared variable " <> show id'
    Nothing -> return (sig, M.insert id' a c : cs)

-- | Extends the function context with a variable and a function signature
extendFun :: Monad m => Env sig val -> Id -> sig -> m (Env sig val)
extendFun (sigs, cs) id' def =
  case M.lookup id' sigs of
    Just f  -> fail $ "redeclaration of function " <> show id'
    Nothing -> return (M.insert id' def sigs, cs)

-- | Updates a variable context with a new value
updateVar :: Monad m => Env sig val -> Id -> val -> m (Env sig val)
updateVar (sigs, cs) id' a = let
  -- Does the traversing stuff
  update :: Monad m => [Context val] -> Id -> val -> m ([Context val])
  update [] _ _ = fail $ "assignment of undeclared variable " <> show id'
  update (c:cs) id' a = case M.lookup id' c of
    Nothing -> return ((:) c =<< update cs id' a)
    Just _  -> return $ (M.insert id' a c) : cs
  in do
    cs' <- update cs id' a
    return (sigs, cs')

-- | newBlock inserts a fresh variable context
newBlock :: Env s a -> Env s a
newBlock (s, cs) = (s, M.empty:cs)

-- | exitBlock pops the current context
exitBlock :: Env s a -> Env s a
exitBlock (s, cs) = (s, drop 1 cs)

-- | Applies function `f` on a temporary scope in the environment
withBlock :: Env s a -> (Env s a -> Env s a) -> Env s a
withBlock env f = exitBlock $ f $ newBlock env

-- | Applies function `f` on a temporary scope (inside a monad)
withBlockM :: Monad m => Env s a -> (Env s a -> m (Env s a)) -> m (Env s a)
withBlockM env f = f (newBlock env) >>= (return . exitBlock)