EventT.hs

{-
 -  ``Control/Monad/Event/EventT''
 -  (c) 2008 Cook, J. MR  SSD, Inc.
 -}
{-# LANGUAGE
    GeneralizedNewtypeDeriving,
    MultiParamTypeClasses,
    FlexibleContexts,
    FlexibleInstances,
    UndecidableInstances,
    TypeSynonymInstances,
    KindSignatures
  #-}

module Control.Monad.EventT
    ( 
      EventT
    , EventM
    , runEventT
    , runEventGraph
    , runEventGraphWithState
    
    , newEventT_RState
    , newEventT_RWState
    
    , onClockChanged
    , onEventDispatch
    , onEventComplete
    , onEventSchedule
    , onEventCancel
    
    , addHandler
    , removeHandler
    
    , SetDebugHandlers(..)
    , DebugHandlerOptions(..)
    , defaultHandlerOptions
    ) where

import Control.Monad.Event.Classes
import Control.Monad.Event.Internal.Types

import Data.Handler
import Data.PriorityQueue
import Data.StateRef
import Control.Monad.RWS
import Control.Monad.Loops

import Text.PrettyPrint.Leijen

import IO
import Control.Monad.Trans
import Control.Concurrent (threadDelay)

-- imports just for classes to implement
import Control.Monad.Error (MonadError(..))
import Control.Monad (MonadPlus(..))
import Control.Monad.Cont (MonadCont(..))

import qualified Data.Map as M

{- Time management -}
instance Monad m => MonadTime (EventT t m) t where
    getCurrentTime = EventT (gets currentTime)

{- Time Management (internal) -}
setCurrentTime :: Monad m => t -> EventT t m ()
setCurrentTime t1 = do
    rwState <- EventT get
    
    let t0 = currentTime rwState
    EventT (put (rwState {currentTime = t1}))
    
    invokeHandler clockChangedHandler (t0,t1)

{- simulation control -}
instance Monad m => MonadSimControl (EventT t m) where
    resumeSimulation = do
        rwState <- EventT get
        EventT (put (rwState {simRunning = True}))
    pauseSimulation = do
        rwState <- EventT get
        EventT (put (rwState {simRunning = False}))
    isSimulationRunning = EventT (gets simRunning)

{- MonadEvent instance for unadorned (EventT t m) actions -}
instance Monad m => MonadEvent (EventT t m) (EventT t m a) where
    describeEvent e = return (text "Undocumented Event")
    runEvent e = e >> return ()

{- MonadEvent instance for Event Descriptors - this does most of 
    the real grunt work of running an event  -}
instance (Monad m, Pretty t) => MonadEvent (EventT t m) (EventDescriptor (EventT t m) t) where
    describeEvent (EventDescriptor {eventId = eid, eventTime = t, event = e}) = do
        eventDescription <- describeEvent e
        return (
            brackets (
                fill 5 (pretty eid)
                <> text "|"
                <> fill 10 (pretty t)
            ) <> colon
            <+> align eventDescription
         )
    runEvent event@(EventDescriptor {event = e}) = do
        setCurrentTime (eventTime event)
        
        invokeHandler eventDispatchHandler event
        
        EventT (local (\env -> env {currentEvent = Just event}) (unWrapEventT (runEvent e)))
        
        invokeHandler eventCompleteHandler event

{- event scheduling, canceling, etc. -}
instance (Monad m, Ord t, Num t, MonadEvent (EventT t m) e) => ScheduleEvent (EventT t m) t e where
    scheduleEventIn dt e = do
        now <- getCurrentTime
        let t = (now + dt)
        
        eid <- getNextEventId
        let event = EventDescriptor 
                {eventId = eid, eventTime = t, event = e}
        
        enqueue eventQueueRef event
        invokeHandler eventScheduleHandler event
        
        return eid
    
    doNext e = scheduleEventIn 0 e >> return ()

instance (Monad m, Ord t) => CancelEvent (EventT t m) t where
    cancelEvent eid = do
        maybeEvent <- dequeueWhere eventQueueRef (\e -> eventId e == eid)
        
        invokeHandler eventCancelHandler (maybe (Left eid) Right maybeEvent)
        return maybeEvent

{- services for use in implementing events -}
instance Monad m => GetCurrentEvent (EventT t m) t where
    getCurrentEvent = EventT (asks currentEvent)

instance Monad m => RetryEvent (EventT t m) t where
    retryEventIn dt = do
        maybeCe <- getCurrentEvent
        case maybeCe of
            Nothing -> fail "retry called outside an event"
            Just (EventDescriptor {event = e}) ->
                scheduleEventIn dt e

instance (Monad m, Ord t) => MonadEventQueueInfo (EventT t m) t where
    eventQueueSize = withEventQueue queueSize
    eventQueueContents = withEventQueue peekQueue


{- The EventT type and related administrative bits (state management stuff) -}
-- |A monad transformer which adds an event queue and related operations to
-- an underlying monad.  The \"t\" parameter specifies the type of the
-- simulation time.
--
-- Several hooks are provided to allow special handling of various events,
-- such as the progression of time, the scheduling or canceling or dispatch
-- of an event, etc.
newtype EventT t m a = EventT { unWrapEventT :: RWST (EventT_RState t m) () (EventT_RWState t m) m a }
    deriving (Functor, Monad, MonadIO)

instance MonadTrans (EventT t) where
    lift = EventT . lift

{- Monad type-class pass-through implementations -}
instance MonadReader r m => MonadReader r (EventT t m) where
    ask = EventT (lift ask)
    local f (EventT x) = EventT $ do
        rState  <- ask
        rwState <- get
        (a, rwState, w) <- lift (local f (runRWST x rState rwState))
        put rwState
        tell w
        return a

instance MonadWriter w m => MonadWriter w (EventT t m) where
    tell w = EventT (lift (tell w))
    listen (EventT x) = EventT $ do
        rState  <- ask
        rwState <- get
        ((a, rwState, w), w2) <- lift (listen (runRWST x rState rwState))
        put rwState
        tell w
        return (a, w2)
    pass (EventT x) = EventT $ do
        rState  <- ask
        rwState <- get
        let knot = do
            ((a, f), rwState, w) <- runRWST x rState rwState 
            return ((a, rwState, w), f)
        (a, rwState, w) <- lift (pass knot)
        put rwState
        tell w
        return a

instance MonadState s m => MonadState s (EventT t m) where
    get = EventT (lift get)
    put s = EventT (lift (put s))

instance MonadError e m => MonadError e (EventT t m) where
    throwError e = EventT (lift (throwError e))
    catchError (EventT x) onErr = EventT $ do
        rState  <- ask
        rwState <- get
        let onErr' err = runRWST (unWrapEventT (onErr err)) rState rwState
        (a, rwState, w) <- lift (catchError (runRWST x rState rwState) onErr')
        put rwState
        tell w
        return a

-- is this valid?
instance MonadPlus m => MonadPlus (EventT t m) where
    mzero = EventT (lift mzero)
    mplus (EventT a) (EventT b) = EventT $ do
        rState  <- ask
        rwState <- get
        (c, rwState, w) <- lift (runRWST a rState rwState `mplus` runRWST b rState rwState)
        put rwState
        tell w
        return c

-- what semantics does this inherit?  is it sensible?  is it unique?
instance MonadCont m => MonadCont (EventT t m) where
    callCC f = EventT $
        callCC $ \cont ->
            unWrapEventT (f (EventT . cont))

-- |'EventM' is a shorthand for an event graph monad in IO with clock of type Double
type EventM = EventT Double IO

{- Running (EventT t m) actions and whole event graphs -}
-- |Run an 'EventT' wrapped action.  This is a \"raw\" action - there is no 
-- interaction with the state (including event graph) except whatever the
-- action itself does.
runEventT :: Monad m => EventT t m a -> EventT_RState t m -> EventT_RWState t m -> m (a, EventT_RWState t m)
runEventT (EventT x) rState rwState = do
    (a, rwState, _) <- runRWST x rState rwState
    return (a, rwState)

-- |Repeatedly pull and run the next event in the queue until it's 
-- empty or until the simulation is paused using 'pauseSimulation'
-- or something equivalent.
runEventGraphWithState :: 
    ( Monad m, Ord t, Pretty t
    ) => EventT_RState t m -> EventT_RWState t m -> m (EventT_RWState t m)
runEventGraphWithState rState rwState = do
    (_, rwState) <- runEventT (whileJust_ dequeueNextEvent runEvent) rState rwState
    return rwState

-- |Initialize the event queue and other stuff, enqueue the provided \"start 
-- event\", and run the queue until it's empty or until the simulation is
-- paused.
runEventGraph ::
    ( Monad m, MonadEvent (EventT t m) e
    , Ord t, Num t, Pretty t
    ) => e -> m (EventT_RState t m, EventT_RWState t m)
runEventGraph e = do
    let rState  = newEventT_RState
        rwState = newEventT_RWState 0
    
    (_, rwState) <- runEventT (scheduleEventIn 0 e) rState rwState
    rwState <- runEventGraphWithState rState rwState
    
    return (rState, rwState)


{- the main state vectors -}
data EventT_RState t m = EventT_RState
    { currentEvent  :: Maybe (EventDescriptor (EventT t m) t)
    }

data EventT_RWState t m = EventT_RWState
    { currentTime   :: t
    , simRunning    :: Bool
    , nextEventId   :: EventID
    , eventQueue    :: EventQueue t m
    , handlers      :: EventTHandlers t m
    }

-- |A new instance of the read/write portion of the 'EventT' internal
-- state.  The parameter is the initial time value.
newEventT_RWState :: (Monad m, Ord t) => t -> EventT_RWState t m
newEventT_RWState t = EventT_RWState
    { currentTime   = t
    , simRunning    = True
    , nextEventId   = EventID 1
    , eventQueue    = newEventQueue
    , handlers      = newEventTHandlers
    }

-- |A new instance of the read-only portion of the 'EventT' internal
-- state.
newEventT_RState :: EventT_RState t m
newEventT_RState = EventT_RState
    { currentEvent  = Nothing
    }

{- Miscellaneous small state management functions -}
getNextEventId :: Monad m => EventT t m EventID
getNextEventId = do
    rwState <- EventT get
    let eID = nextEventId rwState
    EventT (put (rwState {nextEventId = succ eID}))
    return eID

{- Support for debugging event handlers -}
data EventTHandlers t m = EventTHandlers
    { eventScheduleHandler  :: HandlerSet (EventT t m) (EventDescriptor (EventT t m) t) ()
    , eventCancelHandler    :: HandlerSet (EventT t m) (Either EventID (EventDescriptor (EventT t m) t)) ()
    , eventDispatchHandler  :: HandlerSet (EventT t m) (EventDescriptor (EventT t m) t) ()
    , eventCompleteHandler  :: HandlerSet (EventT t m) (EventDescriptor (EventT t m) t) ()
    , clockChangedHandler   :: HandlerSet (EventT t m) (t, t) () -- (old time, new time)
    }

type HandlerAccessor t m a b = 
    ( EventTHandlers t m -> HandlerSet (EventT t m) a b
    , HandlerSet (EventT t m) a b -> EventTHandlers t m -> EventTHandlers t m
    )

-- |Fires after an event is scheduled.  Is passed an 'EventDescriptor' for
-- the event.
onEventSchedule :: HandlerAccessor t m (EventDescriptor (EventT t m) t) ()
onEventSchedule = (eventScheduleHandler, \h hs -> hs {eventScheduleHandler = h})

-- |Fires after an event is canceled.  Is passed either an 'EventID' (if the
-- cancellation failed) or an 'EventDescriptor' for the event that was canceled.
onEventCancel   :: HandlerAccessor t m (Either EventID (EventDescriptor (EventT t m) t)) ()
onEventCancel = (eventCancelHandler, \h hs -> hs {eventCancelHandler = h})

-- |Fires just before an event is dispatched.  Is passed an 'EventDescriptor'
-- describing the event about to be run.
onEventDispatch :: HandlerAccessor t m (EventDescriptor (EventT t m) t) ()
onEventDispatch = (eventDispatchHandler, \h hs -> hs {eventDispatchHandler = h})

-- |Fires after an event returns.  Is passed an 'EventDescriptor' for the
-- event that just finished.
onEventComplete :: HandlerAccessor t m (EventDescriptor (EventT t m) t) ()
onEventComplete = (eventCompleteHandler, \h hs -> hs {eventCompleteHandler = h})

-- |Fires whenever the clock changes, and is passed a tuple containing
-- (old time, new time)
onClockChanged  :: HandlerAccessor t m (t, t) ()
onClockChanged = (clockChangedHandler, \h hs -> hs {clockChangedHandler = h})

newEventTHandlers :: Monad m => EventTHandlers t m
newEventTHandlers = EventTHandlers
    { eventScheduleHandler  = emptyHandlerSet
    , eventCancelHandler    = emptyHandlerSet
    , eventDispatchHandler  = emptyHandlerSet
    , eventCompleteHandler  = emptyHandlerSet
    , clockChangedHandler   = emptyHandlerSet
    }

-- |Add an event handler to be called when the specified event happens.
addHandler :: Monad m => HandlerAccessor t m a b -> (a -> EventT t m b) -> EventT t m HandlerID
addHandler (getter, setter) h = do
    hSet <- EventT (gets (getter.handlers))
    
    let (newHSet, hId) = addHandlerToSet h hSet
    
    EventT (modify (\s -> s {handlers = setter newHSet (handlers s)}))
    return hId

-- |Remove an event handler given its ID, and return it if it was in the set.
removeHandler :: Monad m => HandlerAccessor t m a b -> HandlerID -> EventT t m (Maybe (a -> EventT t m b))
removeHandler (getter, setter) hId = do
    hSet <- EventT (gets (getter.handlers))
    
    let (newHSet, h) = removeHandlerFromSet hId hSet
    
    EventT (modify (\s -> s {handlers = setter newHSet (handlers s)}))
    return h

invokeHandler :: Monad m => (EventTHandlers t m -> HandlerSet (EventT t m) a b) -> a -> EventT t m b
invokeHandler h args = do
    hSet <- EventT (gets (h.handlers))
    invokeHandlers hSet args


{- internal support relating to the event queue -}
type EventQueue t m = PQ (EventDescriptor (EventT t m) t)
type EventQueueRef t m = PriorityQueue (EventT t m) (EventDescriptor (EventT t m) t)

newEventQueue :: Ord t => EventQueue t m
newEventQueue = emptyPQ eventTime

withEventQueue :: (Monad m, Ord t) => (EventQueueRef t m -> EventT t m a) -> EventT t m a
withEventQueue f = f eventQueueRef

eventQueueRef :: (Monad m, Ord t) => EventQueueRef t m
eventQueueRef = mkPriorityQueue (UnsafeModifyRef (Accessor (Getter getEventQueue, Setter setEventQueue)))
    where 
        getEventQueue = EventT (gets eventQueue)
        setEventQueue eq = EventT $ do
            modify (\rwState -> rwState {eventQueue = eq})

dequeueNextEvent :: (Monad m, Ord t) => EventT t m (Maybe (EventDescriptor (EventT t m) t))
dequeueNextEvent = do
    running <- isSimulationRunning
    if running
        then withEventQueue dequeue
        else return Nothing

{- Some handy event handlers for logging and speed regulation -}
{- at some point these may be extracted to a type class -}
-- |(warning: very likely to change or disappear)
-- 
-- A collection of 'Bool's telling 'SetDebugHandlers' which handlers to install.
data DebugHandlerOptions = DebugHandlerOptions
    { -- |Run roughly in sync with the wall clock (using seconds as the simulation time unit)
      sync              :: Bool
    , -- |Print event dispatches to 'stderr' as they occur
      showDispatches    :: Bool 
    , -- |Print event cancellations to 'stderr' as they occur
      showCancels       :: Bool
    , -- |Print event scheduling actions to 'stderr' as they occur
      showSchedules     :: Bool
    , -- |Dump the entire contents of the event queue to
      -- 'stderr' every time an event is dispatched
      dumpQueue         :: Bool
    } deriving Show

defaultHandlerOptions = DebugHandlerOptions
    { sync              = False
    , showDispatches    = False
    , showCancels       = False
    , showSchedules     = False
    , dumpQueue         = False
    }

-- |(warning: very likely to change or disappear)
-- 
-- An event which installs some very primitive logging event handlers
-- useful for watching what's going on inside the event queue.
data SetDebugHandlers t (m :: * -> *) = SetDebugHandlers DebugHandlerOptions
instance Pretty (SetDebugHandlers t m) where
    pretty (SetDebugHandlers opts) = text "Set debug handlers, using options" <+> colon <> text (show opts)
instance (MonadIO m, RealFrac t, Pretty t) => MonadEvent (EventT t m) (SetDebugHandlers t m) where
    describeEvent e = return (pretty e)
    
    runEvent (SetDebugHandlers opts) = do
        let addHandlerIf p hSel h = if p opts then addHandler hSel h >> return () else return ()
        
        -- very rough wallclock synchronization
        addHandlerIf sync onClockChanged $ \(t1, t2) -> do
            let dt = t2 - t1
            let delayTime = max 0 (dt * 1000000)
            liftIO (threadDelay (floor delayTime))
        
        addHandlerIf showDispatches onEventDispatch $ \e -> do
            description <- describeEvent e
            liftIO (hPutDoc stderr (text "Dispatch: " <> description <> linebreak))
        
        addHandlerIf showCancels onEventCancel $ \e -> case e of
            Right e  -> do
                description <- describeEvent e
                liftIO (hPutDoc stderr (text "Canceled" <> colon <+> description <> linebreak))
            Left eid -> liftIO (hPutDoc stderr (text "Cancel" <> colon <+> text "event" <+> pretty eid <+> text "failed to be canceled - it probably already ran" <> linebreak))
        
        addHandlerIf showSchedules onEventSchedule $ \e -> do
            sz <- eventQueueSize
            qc <- eventQueueContents
            
            description <- describeEvent e
            liftIO (hPutDoc stderr (text "Schedule" <> colon <+> align description <> linebreak))
        
        addHandlerIf dumpQueue onEventDispatch $ \e -> do
            qc <- eventQueueContents
            descriptions <- mapM describeEvent qc
            liftIO (hPutDoc stderr (text "Event Queue" <> colon <+> align (vcat descriptions) <> linebreak))
        
        return ()