{-# LANGUAGE InstanceSigs #-}
module Primitives
  ( PrologIO(..)
  , Globals
  , newGlobals
  , newVarPrim, unnewVarPrim
  , getVarPrim, setVarPrim
  , errorPrim
  )
where

import Data.IORef
import qualified Data.IntMap.Strict as M
import Terms
import Classes (LiftIO(..))

newtype PrologIO a = PrologIO { runPrologIO :: Globals -> IO a }

instance Functor PrologIO where
  fmap :: (a -> b) -> PrologIO a -> PrologIO b
  fmap f (PrologIO a) = PrologIO (\globals -> fmap f (a globals))

instance Applicative PrologIO where
  pure :: a -> PrologIO a
  pure x = PrologIO (\_ -> pure x)

  (PrologIO f) <*> (PrologIO x) = PrologIO (
    \globals -> do
      f' <- f globals
      x' <- x globals
      pure (f' x')
    )

instance Monad PrologIO where
  return :: a -> PrologIO a
  return = pure

  (PrologIO a) >>= f = PrologIO (
    \globals -> do
      a' <- a globals
      let PrologIO b' = f a'
      b' globals
    )

data Globals = Globals
  { values :: IORef (M.IntMap Term)
  , names :: IORef (M.IntMap Name)
  , nextVar :: IORef Int
  }

newGlobals :: IO (Globals)
newGlobals = do
  vs <- newIORef (M.empty)
  ns <- newIORef (M.empty)
  next <- newIORef 0
  return $ Globals { values = vs, names = ns, nextVar = next }

newVarPrim :: Name -> PrologIO Var
newVarPrim name = PrologIO $ \globals -> do
  var <- readIORef (nextVar globals) 
  modifyIORef (names globals) (M.insert var name)
  modifyIORef (nextVar globals) succ
  return (VInt var)

unnewVarPrim :: PrologIO ()
unnewVarPrim = PrologIO $ \globals -> do
  modifyIORef (nextVar globals) pred
  var <- readIORef (nextVar globals) 
  modifyIORef (names globals) (M.delete var)
  return ()

getVarPrim :: Var -> PrologIO (Maybe Term) 
getVarPrim (VInt var) = PrologIO $ \globals ->
  fmap (M.lookup var) (readIORef (values globals))

setVarPrim :: Var -> Maybe Term -> PrologIO ()
setVarPrim (VInt var) term = PrologIO $ \globals -> do
  modifyIORef (values globals) (M.alter (\_ -> term) var)
  return ()

errorPrim :: String -> PrologIO a
errorPrim message = error message

instance LiftIO PrologIO where 
  liftIO x = PrologIO (\_ -> x)