Factor out continuation running to DejaFu.Common

dejafu/Test/DejaFu/Common.hs

@@ -56,10 +56,14 @@ module Test.DejaFu.Common
 
   -- * Memory models
   , MemType(..)
+
+  -- * Miscellaneous
+  , runRefCont
   ) where
 
 import           Control.DeepSeq    (NFData(..))
 import           Control.Exception  (MaskingState(..))
+import           Control.Monad.Ref  (MonadRef(..))
 import           Data.List          (intercalate, nub, sort)
 import           Data.List.NonEmpty (NonEmpty)
 import           Data.Maybe         (fromMaybe, mapMaybe)
@@ -858,6 +862,18 @@ data MemType =
 instance NFData MemType where
   rnf m = m `seq` ()
 
+-------------------------------------------------------------------------------
+-- Miscellaneous
+
+-- | Run with a continuation that writes its value into a reference,
+-- returning the computation and the reference.  Using the reference
+-- is non-blocking, it is up to you to ensure you wait sufficiently.
+runRefCont :: MonadRef r n => (n () -> x) -> (a -> Maybe b) -> ((a -> x) -> x) -> n (x, r (Maybe b))
+runRefCont act f k = do
+  ref <- newRef Nothing
+  let c = k (act . writeRef ref . f)
+  pure (c, ref)
+
 -------------------------------------------------------------------------------
 -- Utilities
 

dejafu/Test/DejaFu/Conc/Internal.hs

@@ -57,16 +57,13 @@ runConcurrency :: MonadRef r n
                -> M n r a
                -> n (Either Failure a, Context n r g, SeqTrace, Maybe (ThreadId, ThreadAction))
 runConcurrency sched memtype g idsrc caps ma = do
-  ref <- newRef Nothing
-
-  let c = runCont ma (AStop . writeRef ref . Just . Right)
+  (c, ref) <- runRefCont AStop (Just . Right) (runM ma)
   let ctx = Context { cSchedState = g
                     , cIdSource   = idsrc
                     , cThreads    = launch' Unmasked initialThread (const c) M.empty
                     , cWriteBuf   = emptyBuffer
                     , cCaps       = caps
                     }
-
   (finalCtx, trace, finalAction) <- runThreads sched memtype ref ctx
   out <- readRef ref
   pure (fromJust out, finalCtx, trace, finalAction)

dejafu/Test/DejaFu/STM.hs

@@ -71,7 +71,7 @@ instance MonadThrow (STMLike n r) where
 instance MonadCatch (STMLike n r) where
   catch (S stm) handler = toSTM (SCatch (runSTM . handler) stm)
 
-instance Monad n => C.MonadSTM (STMLike n r) where
+instance C.MonadSTM (STMLike n r) where
   type TVar (STMLike n r) = TVar r
 
   retry = toSTM (const SRetry)

dejafu/Test/DejaFu/STM/Internal.hs

@@ -101,12 +101,8 @@ instance Foldable Result where
 -- | Run a STM transaction, returning an action to undo its effects.
 doTransaction :: MonadRef r n => M n r a -> IdSource -> n (Result a, n (), IdSource, TTrace)
 doTransaction ma idsource = do
-  ref <- newRef Nothing
-
-  let c = runCont ma (SStop . writeRef ref . Just . Right)
-
+  (c, ref) <- runRefCont SStop (Just . Right) (runCont ma)
   (idsource', undo, readen, written, trace) <- go ref c (pure ()) idsource [] [] []
-
   res <- readRef ref
 
   case res of