Add evalEmulatorTrace

Author: sjoerdvisscher

plutus-contract/src/Plutus/Trace/Emulator.hs

@@ -62,6 +62,8 @@ module Plutus.Trace.Emulator(
     , runEmulatorStream
     , TraceConfig(..)
     , runEmulatorTrace
+    , runEmulatorTrace'
+    , evalEmulatorTrace
     , PrintEffect(..)
     , runEmulatorTraceEff
     , runEmulatorTraceIO
@@ -96,7 +98,7 @@ import Wallet.Emulator.MultiAgent (EmulatorEvent,
                                    EmulatorState (_chainState, _walletStates), MultiAgentControlEffect,
                                    MultiAgentEffect, _eteEmulatorTime, _eteEvent, schedulerEvent)
 import Wallet.Emulator.Stream (EmulatorConfig (_initialChainState, _params), EmulatorErr, foldEmulatorStreamM,
-                               initialChainState, initialDist, params, runTraceStream)
+                               initialChainState, initialDist, params, runTraceStream, runTraceStream')
 import Wallet.Emulator.Stream qualified
 import Wallet.Emulator.Wallet (Entity, balances)
 import Wallet.Emulator.Wallet qualified as Wallet
@@ -161,9 +163,9 @@ handleEmulatorTrace ::
     )
     => Params
     -> EmulatorTrace a
-    -> Eff (Reader ThreadId ': Yield (EmSystemCall effs EmulatorMessage) (Maybe EmulatorMessage) ': effs) ()
+    -> Eff (Reader ThreadId ': Yield (EmSystemCall effs EmulatorMessage) (Maybe EmulatorMessage) ': effs) a
 handleEmulatorTrace Params{pNetworkId, pSlotConfig} action = do
-    _ <- subsume @(Error EmulatorRuntimeError)
+    result <- subsume @(Error EmulatorRuntimeError)
             . interpret (mapLog (UserThreadEvent . UserLog))
             . flip handleError (throwError . EmulatedWalletError)
             . reinterpret handleEmulatedWalletAPI
@@ -174,13 +176,21 @@ handleEmulatorTrace Params{pNetworkId, pSlotConfig} action = do
             . interpret (handleStartContract @_ @effs pNetworkId)
             $ raiseEnd action
     void $ exit @effs @EmulatorMessage
+    pure result
 
 -- | Run a 'Trace Emulator', streaming the log messages as they arrive
 runEmulatorStream :: forall effs a.
     EmulatorConfig
     -> EmulatorTrace a
     -> Stream (Of (LogMessage EmulatorEvent)) (Eff effs) (Maybe EmulatorErr, EmulatorState)
-runEmulatorStream conf = runTraceStream conf . interpretEmulatorTrace conf
+runEmulatorStream conf = runTraceStream conf . void . interpretEmulatorTrace conf
+
+-- | Run a 'Trace Emulator', streaming the log messages as they arrive
+runEmulatorStream' :: forall effs a.
+    EmulatorConfig
+    -> EmulatorTrace a
+    -> Stream (Of (LogMessage EmulatorEvent)) (Eff effs) (Either EmulatorErr a, EmulatorState)
+runEmulatorStream' conf = runTraceStream' conf . interpretEmulatorTrace conf
 
 -- | Interpret a 'Trace Emulator' action in the multi agent and emulated
 --   blockchain effects.
@@ -194,12 +204,16 @@ interpretEmulatorTrace :: forall effs a.
     )
     => EmulatorConfig
     -> EmulatorTrace a
-    -> Eff effs ()
+    -> Eff effs (Maybe a)
 interpretEmulatorTrace conf action =
     -- add a wait action to the beginning to ensure that the
     -- initial transaction gets validated before the wallets
     -- try to spend their funds
-    let action' = Waiting.nextSlot >> action >> Waiting.nextSlot
+    let action' = do
+          void Waiting.nextSlot
+          res <- action
+          void Waiting.nextSlot
+          pure res
         wallets = fromMaybe (Wallet.toMockWallet <$> CW.knownMockWallets) (preview (initialChainState . _Left . to Map.keys) conf)
     in
     evalState @EmulatorThreads mempty
@@ -236,11 +250,18 @@ defaultShowEvent = \case
   WalletEvent _ _                                                      -> Nothing
   ev                                                                   -> Just . renderString . layoutPretty defaultLayoutOptions . pretty $ ev
 
+evalEmulatorTrace
+    :: EmulatorConfig
+    -> EmulatorTrace a
+    -> Either EmulatorErr a
+evalEmulatorTrace cfg trace = case runEmulatorTrace' cfg trace of
+  (_, r, _) -> r
+
 -- | Run an emulator trace to completion, returning a tuple of the final state
 -- of the emulator, the events, and any error, if any.
 runEmulatorTrace
     :: EmulatorConfig
-    -> EmulatorTrace ()
+    -> EmulatorTrace a
     -> ([EmulatorEvent], Maybe EmulatorErr, EmulatorState)
 runEmulatorTrace cfg trace =
     (\(xs :> (y, z)) -> (xs, y, z))
@@ -249,6 +270,17 @@ runEmulatorTrace cfg trace =
     $ foldEmulatorStreamM (generalize list)
     $ runEmulatorStream cfg trace
 
+runEmulatorTrace'
+    :: EmulatorConfig
+    -> EmulatorTrace a
+    -> ([EmulatorEvent], Either EmulatorErr a, EmulatorState)
+runEmulatorTrace' cfg trace =
+    (\(xs :> (y, z)) -> (xs, y, z))
+    $ run
+    $ runReader (initialDist cfg)
+    $ foldEmulatorStreamM (generalize list)
+    $ runEmulatorStream' cfg trace
+
 -- | Run the emulator trace returning an effect that can be evaluated by
 -- interpreting the 'PrintEffect's.
 runEmulatorTraceEff :: forall effs. Member PrintEffect effs

plutus-contract/src/Plutus/Trace/Playground.hs

@@ -151,7 +151,8 @@ interpretPlaygroundTrace :: forall w s e effs a.
     -> PlaygroundTrace a
     -> Eff effs ()
 interpretPlaygroundTrace conf contract wallets action =
-    evalState @EmulatorThreads mempty
+    void
+        $ evalState @EmulatorThreads mempty
         $ evalState @(Map Wallet ContractInstanceId) Map.empty
         $ handleDeterministicIds
         $ interpret (mapLog (review schedulerEvent))

plutus-contract/src/Plutus/Trace/Scheduler.hs

@@ -15,6 +15,7 @@
 {-# LANGUAGE TypeApplications    #-}
 {-# LANGUAGE TypeFamilies        #-}
 {-# LANGUAGE TypeOperators       #-}
+{-# OPTIONS_GHC -Wno-orphans #-}
 {-
 
 Implements a scheduler for cooperative multitasking. The scheduler supports
@@ -53,6 +54,7 @@ module Plutus.Trace.Scheduler(
     ) where
 
 
+import Control.Applicative ((<|>))
 import Control.Lens hiding (Empty)
 import Control.Monad.Freer
 import Control.Monad.Freer.Coroutine
@@ -72,6 +74,8 @@ import Plutus.Trace.Tag (Tag)
 import Prettyprinter
 import Prettyprinter.Extras (PrettyShow (..), Tagged (..))
 
+deriving instance Functor (Status effs a b)
+
 {- Note [Thread Tag]
 
 Within the scheduler, threads are identified by their 'ThreadId'. The thread
@@ -157,24 +161,27 @@ data WithPriority t
         , _thread   :: t
         } deriving Functor
 
-type SuspendedThread effs systemEvent = WithPriority (EmThread effs systemEvent)
+type SuspendedThread effs systemEvent a = WithPriority (EmThread effs systemEvent a)
 
 type EmSystemCall effs systemEvent = WithPriority (SysCall effs systemEvent)
 
 type AgentSystemCall systemEvent = WithPriority (MessageCall systemEvent)
 
 -- | Thread that can be run by the scheduler
-data EmThread effs systemEvent =
+data EmThread effs systemEvent a =
     EmThread
-        { _continuation :: Maybe systemEvent -> Eff effs (Status effs (EmSystemCall effs systemEvent) (Maybe systemEvent) ()) -- ^ The continuation to be run when the thread is resumed.
+        { _continuation :: Maybe systemEvent -> Eff effs (Status effs (EmSystemCall effs systemEvent) (Maybe systemEvent) (Maybe a)) -- ^ The continuation to be run when the thread is resumed.
         , _threadId     :: ThreadId -- ^ Thread ID
         , _tag          :: Tag -- ^ Tag of the thread. See note [Thread Tag]
         }
 
+mapEmThread :: (Maybe a -> Maybe b) -> EmThread effs systemEvent a -> EmThread effs systemEvent b
+mapEmThread f EmThread{_continuation, _threadId, _tag} = EmThread{ _threadId, _tag, _continuation = fmap (fmap (fmap f)) _continuation }
+
 -- | The system calls we can make to the scheduler, affecting the the threads
 --   that are currently running.
 data ThreadCall effs systemEvent
-    = Fork (ThreadId -> SuspendedThread effs systemEvent) -- ^ Start a new thread with a new thread ID.
+    = Fork (ThreadId -> SuspendedThread effs systemEvent ()) -- ^ Start a new thread with a new thread ID.
     | Thaw ThreadId -- ^ Unfreeze a thread.
     | Exit -- ^ Terminate the scheduler.
 
@@ -190,11 +197,11 @@ makePrisms ''MessageCall
 makePrisms ''ThreadCall
 
 -- | Scheduler state
-data SchedulerState effs systemEvent
+data SchedulerState effs systemEvent a
     = SchedulerState
-        { _normalPrio    :: Seq (EmThread effs systemEvent) -- ^  Threads running at normal priority
-        , _sleeping      :: Seq (EmThread effs systemEvent) -- ^ Sleeping threads (waiting for an external event)
-        , _frozen        :: Seq (EmThread effs systemEvent) -- ^ Frozen threads (will not be resumed until they are explicitly unfrozen)
+        { _normalPrio    :: Seq (EmThread effs systemEvent a) -- ^  Threads running at normal priority
+        , _sleeping      :: Seq (EmThread effs systemEvent a) -- ^ Sleeping threads (waiting for an external event)
+        , _frozen        :: Seq (EmThread effs systemEvent a) -- ^ Frozen threads (will not be resumed until they are explicitly unfrozen)
         , _lastThreadId  :: ThreadId -- ^ Last thread id assigned to a thread
         , _mailboxes     :: HashMap ThreadId (Seq systemEvent) -- ^ The mailboxes of all active threads.
         , _activeThreads :: Map Tag (HashSet ThreadId) -- ^ Map of tags to thread IDs. See note [Thread Tag]
@@ -204,16 +211,16 @@ makeLenses ''SchedulerState
 
 -- | Remove a thread from the set of active threads. Usually called when the
 --   thread is finished.
-removeActiveThread :: ThreadId -> SchedulerState effs systemEvent -> SchedulerState effs systemEvent
+removeActiveThread :: ThreadId -> SchedulerState effs systemEvent a -> SchedulerState effs systemEvent a
 removeActiveThread tid = over (activeThreads . mapped) (HashSet.delete tid)
 
 -- | A suspended thread with a priority and the thread itself.
-suspendThread :: Priority -> EmThread effs systemEvent -> SuspendedThread effs systemEvent
+suspendThread :: Priority -> EmThread effs systemEvent a -> SuspendedThread effs systemEvent a
 suspendThread = WithPriority
 
 -- | Make a thread with the given priority from an action. This is a
 --   convenience for defining 'SimulatorInterpreter' values.
-mkThread :: Tag -> Priority -> Eff (Reader ThreadId ': Yield (EmSystemCall effs systemEvent) (Maybe systemEvent) ': effs) () -> ThreadId -> SuspendedThread effs systemEvent
+mkThread :: Tag -> Priority -> Eff (Reader ThreadId ': Yield (EmSystemCall effs systemEvent) (Maybe systemEvent) ': effs) (Maybe a) -> ThreadId -> SuspendedThread effs systemEvent a
 mkThread tag prio action tid =
     let action' = runReader tid action
     in WithPriority
@@ -248,7 +255,7 @@ fork :: forall effs systemEvent effs2.
     -> Priority -- ^ Priority of the new thread.
     -> Eff (Reader ThreadId ': Yield (EmSystemCall effs systemEvent) (Maybe systemEvent) ': effs) ()
     -> Eff effs2 (Maybe systemEvent)
-fork tag prio action = mkSysCall prio (Right $ Fork $ mkThread tag prio action)
+fork tag prio action = mkSysCall prio (Right $ Fork $ mkThread tag prio $ Just <$> action)
 
 -- | Suspend the current thread
 sleep :: forall effs systemEvent effs2.
@@ -271,59 +278,60 @@ initialThreadTag = "initial thread"
 --   effect using the scheduler, see note [Scheduler]. 'runThreads' only
 --   returns when all threads are finished.
 runThreads ::
-    forall effs systemEvent.
+    forall a effs systemEvent.
     ( Eq systemEvent
     , Member (LogMsg SchedulerLog) effs
     )
-    => Eff (Reader ThreadId ': Yield (EmSystemCall effs systemEvent) (Maybe systemEvent) ': effs) ()
-    -> Eff effs ()
+    => Eff (Reader ThreadId ': Yield (EmSystemCall effs systemEvent) (Maybe systemEvent) ': effs) a
+    -> Eff effs (Maybe a)
 runThreads e = do
     k <- runC $ runReader initialThreadId e
     case k of
-        Done () -> pure ()
+        Done a -> pure $ Just a
         Continue _ k' ->
-            let initialThread = EmThread{_continuation = k', _threadId = initialThreadId, _tag = initialThreadTag}
-            in loop
+            let initialThread = EmThread{_continuation = fmap (fmap Just) . k', _threadId = initialThreadId, _tag = initialThreadTag}
+            in loop Nothing
                 $ initialState
                     & activeThreads . at initialThreadTag . non mempty %~ HashSet.insert initialThreadId
                     & mailboxes . at initialThreadId ?~ Seq.empty
                     & (fst . nextThreadId)
                     & enqueue (suspendThread Normal initialThread)
 
 -- | Run the threads that are scheduled in a 'SchedulerState' to completion.
-loop :: forall effs systemEvent.
+loop :: forall a effs systemEvent.
     ( Eq systemEvent
     , Member (LogMsg SchedulerLog) effs
     )
-    => SchedulerState effs systemEvent
-    -> Eff effs ()
-loop s = do
+    => Maybe a
+    -> SchedulerState effs systemEvent a
+    -> Eff effs (Maybe a)
+loop ma s = do
     case dequeue s of
         AThread EmThread{_continuation, _threadId, _tag} event schedulerState prio -> do
             let mkLog e = SchedulerLog{slEvent=e, slThread=_threadId, slPrio=prio, slTag = _tag}
             result <- _continuation event
             case result of
-                Done () -> loop $ schedulerState & removeActiveThread _threadId
+                Done ma' -> loop (ma <|> ma') $ schedulerState & removeActiveThread _threadId
                 Continue WithPriority{_priority, _thread=sysCall} k -> do
                     let thisThread = suspendThread _priority EmThread{_threadId=_threadId, _continuation=k, _tag = _tag}
                     newState <- schedulerState & enqueue thisThread & handleSysCall sysCall
                     case newState of
-                        Left r          -> logDebug (mkLog $ Stopped r)
-                        Right newState' -> loop newState'
-        _ -> pure ()
+                        Left r          -> logDebug (mkLog $ Stopped r) >> pure ma
+                        Right newState' -> loop ma newState'
+        _ -> pure ma
 
 -- | Deal with a system call from a running thread.
 handleSysCall ::
     ( Eq systemEvent
     , Member (LogMsg SchedulerLog) effs
     )
     => SysCall effs systemEvent
-    -> SchedulerState effs systemEvent
-    -> Eff effs (Either StopReason (SchedulerState effs systemEvent))
+    -> SchedulerState effs systemEvent a
+    -> Eff effs (Either StopReason (SchedulerState effs systemEvent a))
 handleSysCall sysCall schedulerState = case sysCall of
     Right (Fork newThread) -> do
         let (schedulerState', tid) = nextThreadId schedulerState
-            t = newThread tid
+            t = mapEmThread (const Nothing) <$> newThread tid
             tag = _tag $ _thread t
             newState = enqueue t schedulerState'
                         & activeThreads . at tag . non mempty %~ HashSet.insert tid
@@ -344,15 +352,15 @@ handleSysCall sysCall schedulerState = case sysCall of
 
 
 -- | Return a fresh thread ID and increment the counter
-nextThreadId :: SchedulerState effs systemEvent -> (SchedulerState effs systemEvent, ThreadId)
+nextThreadId :: SchedulerState effs systemEvent a -> (SchedulerState effs systemEvent a, ThreadId)
 nextThreadId s = (s & lastThreadId %~ ThreadId . succ . unThreadId, s ^. lastThreadId)
 
 -- | State of the scheduler before any threads are run.
-initialState :: SchedulerState effs systemEvent
+initialState :: SchedulerState effs systemEvent a
 initialState = SchedulerState Seq.empty Seq.empty Seq.empty initialThreadId HashMap.empty Map.empty
 
 -- | Add a suspended thread to the queue.
-enqueue :: SuspendedThread effs systemEvent -> SchedulerState effs systemEvent -> SchedulerState effs systemEvent
+enqueue :: SuspendedThread effs systemEvent a -> SchedulerState effs systemEvent a -> SchedulerState effs systemEvent a
 enqueue WithPriority {_priority, _thread} s =
     case _priority of
         Normal   -> s & normalPrio %~ (|> _thread)
@@ -361,11 +369,11 @@ enqueue WithPriority {_priority, _thread} s =
 
 -- | Result of calling 'dequeue'. Either a thread that is ready to receive a
 --   message, or no more threads.
-data SchedulerDQResult effs systemEvent
-    = AThread (EmThread effs systemEvent) (Maybe systemEvent) (SchedulerState effs systemEvent) Priority
+data SchedulerDQResult effs systemEvent a
+    = AThread (EmThread effs systemEvent a) (Maybe systemEvent) (SchedulerState effs systemEvent a) Priority
     | NoMoreThreads
 
-dequeue :: SchedulerState effs systemEvent -> SchedulerDQResult effs systemEvent
+dequeue :: SchedulerState effs systemEvent a -> SchedulerDQResult effs systemEvent a
 dequeue s = case dequeueThread s of
     Nothing -> NoMoreThreads
     Just (s', thread, prio) -> case dequeueMessage s' (_threadId thread) of
@@ -374,7 +382,7 @@ dequeue s = case dequeueThread s of
 
 -- | Find the next thread that is ready to be resumed.
 --   See note [Thread Priority]
-dequeueThread :: SchedulerState effs systemEvent -> Maybe (SchedulerState effs systemEvent, EmThread effs systemEvent, Priority)
+dequeueThread :: SchedulerState effs systemEvent  a-> Maybe (SchedulerState effs systemEvent a, EmThread effs systemEvent a, Priority)
 dequeueThread s =
     case s ^. normalPrio of
         x :<| xs -> Just (s & normalPrio .~ xs, x, Normal)
@@ -383,7 +391,7 @@ dequeueThread s =
                 Empty    -> Nothing
 
 -- | Get the first message for the thread.
-dequeueMessage :: SchedulerState effs systemEvent -> ThreadId -> Maybe (SchedulerState effs systemEvent, systemEvent)
+dequeueMessage :: SchedulerState effs systemEvent a -> ThreadId -> Maybe (SchedulerState effs systemEvent a, systemEvent)
 dequeueMessage s i = do
     mailbox <- s ^. mailboxes . at i
     (x, xs) <- case mailbox of { Empty -> Nothing; x :<| xs -> Just (x, xs) }

plutus-contract/src/Wallet/Emulator/Stream.hs

@@ -19,6 +19,7 @@ module Wallet.Emulator.Stream(
     , initialState
     , params
     , runTraceStream
+    , runTraceStream'
     -- * Stream manipulation
     , takeUntilSlot
     , filterLogLevel
@@ -120,8 +121,26 @@ runTraceStream :: forall effs.
             , Error EmulatorRuntimeError
             ] ()
     -> Stream (Of (LogMessage EmulatorEvent)) (Eff effs) (Maybe EmulatorErr, EmulatorState)
-runTraceStream conf@EmulatorConfig{_params} =
-    fmap (first (either Just (const Nothing)))
+runTraceStream conf =
+    fmap (first handleResult) . runTraceStream' conf . fmap Just
+    where
+        handleResult (Right _)                          = Nothing
+        handleResult (Left InitialThreadFailedToReturn) = Nothing
+        handleResult (Left err)                         = Just err
+
+runTraceStream' :: forall effs a.
+    EmulatorConfig
+    -> Eff '[ State EmulatorState
+            , LogMsg EmulatorEvent'
+            , MultiAgentEffect
+            , MultiAgentControlEffect
+            , ChainEffect
+            , ChainControlEffect
+            , Error EmulatorRuntimeError
+            ] (Maybe a)
+    -> Stream (Of (LogMessage EmulatorEvent)) (Eff effs) (Either EmulatorErr a, EmulatorState)
+runTraceStream' conf@EmulatorConfig{_params} =
+    fmap (first $ either Left $ maybe (Left InitialThreadFailedToReturn) Right)
     . S.hoist (pure . run)
     . runStream @(LogMessage EmulatorEvent) @_ @'[]
     . runState (initialState conf)
@@ -183,6 +202,7 @@ data EmulatorErr =
     | ChainIndexErr ChainIndexError
     | AssertionErr EM.AssertionError
     | InstanceErr EmulatorRuntimeError
+    | InitialThreadFailedToReturn
     deriving (Show)
 
 handleLogCoroutine :: forall e effs.