-- | The Vectorisation monad.
module Vectorise.Monad.Base (
	-- * The Vectorisation Monad
	VResult(..),
	VM(..),

	-- * Lifting
	liftDs,

	-- * Error Handling
	cantVectorise,
	maybeCantVectorise,
	maybeCantVectoriseM,
	
	-- * Debugging
	traceVt, dumpOptVt, dumpVt,
	
	-- * Control
	noV,	 traceNoV,
	ensureV, traceEnsureV,
	onlyIfV,
	tryV,
	maybeV,  traceMaybeV,
	orElseV,
	fixV,
) where

import Vectorise.Builtins
import Vectorise.Env

import DsMonad
import TcRnMonad
import ErrUtils
import Outputable
import DynFlags
import StaticFlags

import Control.Monad
import System.IO (stderr)


-- The Vectorisation Monad ----------------------------------------------------

-- | Vectorisation can either succeed with new envionment and a value,
--   or return with failure.
data VResult a 
	= Yes GlobalEnv LocalEnv a | No

newtype VM a 
	= VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }

instance Monad VM where
  return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
  VM p >>= f = VM $ \bi genv lenv -> do
                                      r <- p bi genv lenv
                                      case r of
                                        Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
                                        No                -> return No

instance Functor VM where
  fmap = liftM
  
instance MonadIO VM where
  liftIO = liftDs . liftIO


-- Lifting --------------------------------------------------------------------
-- | Lift a desugaring computation into the vectorisation monad.
liftDs :: DsM a -> VM a
liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }


-- Error Handling -------------------------------------------------------------
-- | Throw a `pgmError` saying we can't vectorise something.
cantVectorise :: String -> SDoc -> a
cantVectorise s d = pgmError
                  . showSDocDump
                  $ vcat [text "*** Vectorisation error ***",
                          nest 4 $ sep [text s, nest 4 d]]


-- | Like `fromJust`, but `pgmError` on Nothing.
maybeCantVectorise :: String -> SDoc -> Maybe a -> a
maybeCantVectorise s d Nothing  = cantVectorise s d
maybeCantVectorise _ _ (Just x) = x


-- | Like `maybeCantVectorise` but in a `Monad`.
maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
maybeCantVectoriseM s d p
  = do
      r <- p
      case r of
        Just x  -> return x
        Nothing -> cantVectorise s d


-- Debugging ------------------------------------------------------------------

-- |Output a trace message if -ddump-vt-trace is active.
--
traceVt :: String -> SDoc -> VM () 
traceVt herald doc
  | 1 <= opt_TraceLevel = liftDs $
                            traceOptIf Opt_D_dump_vt_trace $
                              hang (text herald) 2 doc
  | otherwise           = return ()

-- |Dump the given program conditionally.
--
dumpOptVt :: DynFlag -> String -> SDoc -> VM ()
dumpOptVt flag header doc 
  = do { b <- liftDs $ doptM flag
       ; if b 
         then dumpVt header doc 
         else return () 
       }

-- |Dump the given program unconditionally.
--
dumpVt :: String -> SDoc -> VM ()
dumpVt header doc 
  = do { unqual <- liftDs mkPrintUnqualifiedDs
       ; liftIO $ printForUser stderr unqual (mkDumpDoc header doc)
       }

-- Control --------------------------------------------------------------------
-- | Return some result saying we've failed.
noV :: VM a
noV	= VM $ \_ _ _ -> return No


-- | Like `traceNoV` but also emit some trace message to stderr.
traceNoV :: String -> SDoc -> VM a
traceNoV s d  = pprTrace s d noV


-- | If `True` then carry on, otherwise fail.
ensureV :: Bool -> VM ()
ensureV False = noV
ensureV True  = return ()


-- | Like `ensureV` but if we fail then emit some trace message to stderr.
traceEnsureV :: String -> SDoc -> Bool -> VM ()
traceEnsureV s d False = traceNoV s d
traceEnsureV _ _ True  = return ()


-- | If `True` then return the first argument, otherwise fail.
onlyIfV :: Bool -> VM a -> VM a
onlyIfV b p = ensureV b >> p


-- | Try some vectorisation computaton.
--	If it succeeds then return `Just` the result,
--	otherwise return `Nothing`.
tryV :: VM a -> VM (Maybe a)
tryV (VM p) = VM $ \bi genv lenv ->
  do
    r <- p bi genv lenv
    case r of
      Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
      No                -> return (Yes genv  lenv  Nothing)


-- | If `Just` then return the value, otherwise fail.
maybeV :: VM (Maybe a) -> VM a
maybeV p = maybe noV return =<< p


-- | Like `maybeV` but emit a message to stderr if we fail.
traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
traceMaybeV s d p = maybe (traceNoV s d) return =<< p


-- | Try the first computation,
--	if it succeeds then take the returned value,
--	if it fails then run the second computation instead.
orElseV :: VM a -> VM a -> VM a
orElseV p q = maybe q return =<< tryV p


-- | Fixpoint in the vectorisation monad.
fixV :: (a -> VM a) -> VM a
fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
  where
    -- NOTE: It is essential that we are lazy in r above so do not replace
    --       calls to this function by an explicit case.
    unYes (Yes _ _ x) = x
    unYes No          = panic "Vectorise.Monad.Base.fixV: no result"