Skip to content

Commit 486aeca

Browse files
committed
MonadSTM: added unGetTQueue & unGetTBQueue
1 parent 1a3f3bb commit 486aeca

File tree

4 files changed

+93
-26
lines changed

4 files changed

+93
-26
lines changed

io-classes/src/Control/Monad/Class/MonadSTM.hs

Lines changed: 66 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ class ( Monad m
132132
tryPeekTQueue :: TQueue m a -> STM m (Maybe a)
133133
writeTQueue :: TQueue m a -> a -> STM m ()
134134
isEmptyTQueue :: TQueue m a -> STM m Bool
135+
unGetTQueue :: TQueue m a -> a -> STM m ()
135136

136137
type TBQueue m :: Type -> Type
137138
newTBQueue :: Natural -> STM m (TBQueue m a)
@@ -145,6 +146,7 @@ class ( Monad m
145146
lengthTBQueue :: TBQueue m a -> STM m Natural
146147
isEmptyTBQueue :: TBQueue m a -> STM m Bool
147148
isFullTBQueue :: TBQueue m a -> STM m Bool
149+
unGetTBQueue :: TBQueue m a -> a -> STM m ()
148150

149151
-- Helpful derived functions with default implementations
150152

@@ -224,6 +226,10 @@ class ( Monad m
224226
=> TQueue m a -> STM m Bool
225227
isEmptyTQueue = isEmptyTQueueDefault
226228

229+
default unGetTQueue :: TQueue m ~ TQueueDefault m
230+
=> TQueue m a -> a -> STM m ()
231+
unGetTQueue = unGetTQueueDefault
232+
227233
default peekTQueue :: TQueue m ~ TQueueDefault m
228234
=> TQueue m a -> STM m a
229235
peekTQueue = peekTQueueDefault
@@ -272,6 +278,10 @@ class ( Monad m
272278
=> TBQueue m a -> STM m [a]
273279
flushTBQueue = flushTBQueueDefault
274280

281+
default unGetTBQueue :: TBQueue m ~ TBQueueDefault m
282+
=> TBQueue m a -> a -> STM m ()
283+
unGetTBQueue = unGetTBQueueDefault
284+
275285

276286
stateTVarDefault :: MonadSTM m => TVar m s -> (s -> (a, s)) -> STM m a
277287
stateTVarDefault var f = do
@@ -530,6 +540,7 @@ instance MonadSTM IO where
530540
flushTBQueue = STM.flushTBQueue
531541
writeTQueue = STM.writeTQueue
532542
isEmptyTQueue = STM.isEmptyTQueue
543+
unGetTQueue = STM.unGetTQueue
533544
newTBQueue = STM.newTBQueue
534545
readTBQueue = STM.readTBQueue
535546
tryReadTBQueue = STM.tryReadTBQueue
@@ -539,6 +550,7 @@ instance MonadSTM IO where
539550
lengthTBQueue = STM.lengthTBQueue
540551
isEmptyTBQueue = STM.isEmptyTBQueue
541552
isFullTBQueue = STM.isFullTBQueue
553+
unGetTBQueue = STM.unGetTBQueue
542554

543555
newTVarIO = STM.newTVarIO
544556
readTVarIO = STM.readTVarIO
@@ -758,6 +770,11 @@ tryPeekTQueueDefault (TQueue read _write) = do
758770
(x:_) -> return (Just x)
759771
_ -> return Nothing
760772

773+
unGetTQueueDefault :: MonadSTM m => TQueueDefault m a -> a -> STM m ()
774+
unGetTQueueDefault (TQueue read _write) a = modifyTVar read (a:)
775+
776+
777+
761778
--
762779
-- Default TBQueue implementation in terms of TVars
763780
--
@@ -879,6 +896,19 @@ flushTBQueueDefault (TBQueue rsize read wsize write size) = do
879896
writeTVar wsize size
880897
return (xs ++ reverse ys)
881898

899+
unGetTBQueueDefault :: MonadSTM m => TBQueueDefault m a -> a -> STM m ()
900+
unGetTBQueueDefault (TBQueue rsize read wsize _write _size) a = do
901+
r <- readTVar rsize
902+
if (r > 0)
903+
then do writeTVar rsize $! r - 1
904+
else do
905+
w <- readTVar wsize
906+
if (w > 0)
907+
then writeTVar wsize $! w - 1
908+
else retry
909+
xs <- readTVar read
910+
writeTVar read (a:xs)
911+
882912

883913
-- | 'throwIO' specialised to @stm@ monad.
884914
--
@@ -976,6 +1006,7 @@ instance MonadSTM m => MonadSTM (ContT r m) where
9761006
tryPeekTQueue = WrappedSTM . tryPeekTQueue
9771007
writeTQueue v = WrappedSTM . writeTQueue v
9781008
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1009+
unGetTQueue = WrappedSTM .: unGetTQueue
9791010

9801011
type TBQueue (ContT r m) = TBQueue m
9811012
newTBQueue = WrappedSTM . newTBQueue
@@ -988,6 +1019,7 @@ instance MonadSTM m => MonadSTM (ContT r m) where
9881019
lengthTBQueue = WrappedSTM . lengthTBQueue
9891020
isEmptyTBQueue = WrappedSTM . isEmptyTBQueue
9901021
isFullTBQueue = WrappedSTM . isFullTBQueue
1022+
unGetTBQueue = WrappedSTM .: unGetTBQueue
9911023

9921024

9931025
instance MonadSTM m => MonadSTM (ReaderT r m) where
@@ -1021,12 +1053,13 @@ instance MonadSTM m => MonadSTM (ReaderT r m) where
10211053

10221054
type TQueue (ReaderT r m) = TQueue m
10231055
newTQueue = WrappedSTM newTQueue
1024-
readTQueue = WrappedSTM . readTQueue
1025-
tryReadTQueue = WrappedSTM . tryReadTQueue
1026-
peekTQueue = WrappedSTM . peekTQueue
1027-
tryPeekTQueue = WrappedSTM . tryPeekTQueue
1028-
writeTQueue v = WrappedSTM . writeTQueue v
1029-
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1056+
readTQueue = WrappedSTM . readTQueue
1057+
tryReadTQueue = WrappedSTM . tryReadTQueue
1058+
peekTQueue = WrappedSTM . peekTQueue
1059+
tryPeekTQueue = WrappedSTM . tryPeekTQueue
1060+
writeTQueue v = WrappedSTM . writeTQueue v
1061+
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1062+
unGetTQueue = WrappedSTM .: unGetTQueue
10301063

10311064
type TBQueue (ReaderT r m) = TBQueue m
10321065
newTBQueue = WrappedSTM . newTBQueue
@@ -1039,6 +1072,7 @@ instance MonadSTM m => MonadSTM (ReaderT r m) where
10391072
lengthTBQueue = WrappedSTM . lengthTBQueue
10401073
isEmptyTBQueue = WrappedSTM . isEmptyTBQueue
10411074
isFullTBQueue = WrappedSTM . isFullTBQueue
1075+
unGetTBQueue = WrappedSTM .: unGetTBQueue
10421076

10431077

10441078
instance (Monoid w, MonadSTM m) => MonadSTM (WriterT w m) where
@@ -1072,12 +1106,13 @@ instance (Monoid w, MonadSTM m) => MonadSTM (WriterT w m) where
10721106

10731107
type TQueue (WriterT w m) = TQueue m
10741108
newTQueue = WrappedSTM newTQueue
1075-
readTQueue = WrappedSTM . readTQueue
1076-
tryReadTQueue = WrappedSTM . tryReadTQueue
1077-
peekTQueue = WrappedSTM . peekTQueue
1078-
tryPeekTQueue = WrappedSTM . tryPeekTQueue
1079-
writeTQueue v = WrappedSTM . writeTQueue v
1080-
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1109+
readTQueue = WrappedSTM . readTQueue
1110+
tryReadTQueue = WrappedSTM . tryReadTQueue
1111+
peekTQueue = WrappedSTM . peekTQueue
1112+
tryPeekTQueue = WrappedSTM . tryPeekTQueue
1113+
writeTQueue v = WrappedSTM . writeTQueue v
1114+
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1115+
unGetTQueue = WrappedSTM .: unGetTQueue
10811116

10821117
type TBQueue (WriterT w m) = TBQueue m
10831118
newTBQueue = WrappedSTM . newTBQueue
@@ -1090,6 +1125,7 @@ instance (Monoid w, MonadSTM m) => MonadSTM (WriterT w m) where
10901125
lengthTBQueue = WrappedSTM . lengthTBQueue
10911126
isEmptyTBQueue = WrappedSTM . isEmptyTBQueue
10921127
isFullTBQueue = WrappedSTM . isFullTBQueue
1128+
unGetTBQueue = WrappedSTM .: unGetTBQueue
10931129

10941130

10951131
instance MonadSTM m => MonadSTM (StateT s m) where
@@ -1129,6 +1165,7 @@ instance MonadSTM m => MonadSTM (StateT s m) where
11291165
tryPeekTQueue = WrappedSTM . tryPeekTQueue
11301166
writeTQueue v = WrappedSTM . writeTQueue v
11311167
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1168+
unGetTQueue = WrappedSTM .: unGetTQueue
11321169

11331170
type TBQueue (StateT s m) = TBQueue m
11341171
newTBQueue = WrappedSTM . newTBQueue
@@ -1141,6 +1178,7 @@ instance MonadSTM m => MonadSTM (StateT s m) where
11411178
lengthTBQueue = WrappedSTM . lengthTBQueue
11421179
isEmptyTBQueue = WrappedSTM . isEmptyTBQueue
11431180
isFullTBQueue = WrappedSTM . isFullTBQueue
1181+
unGetTBQueue = WrappedSTM .: unGetTBQueue
11441182

11451183

11461184
instance MonadSTM m => MonadSTM (ExceptT e m) where
@@ -1174,12 +1212,13 @@ instance MonadSTM m => MonadSTM (ExceptT e m) where
11741212

11751213
type TQueue (ExceptT e m) = TQueue m
11761214
newTQueue = WrappedSTM newTQueue
1177-
readTQueue = WrappedSTM . readTQueue
1178-
tryReadTQueue = WrappedSTM . tryReadTQueue
1179-
peekTQueue = WrappedSTM . peekTQueue
1180-
tryPeekTQueue = WrappedSTM . tryPeekTQueue
1181-
writeTQueue v = WrappedSTM . writeTQueue v
1182-
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1215+
readTQueue = WrappedSTM . readTQueue
1216+
tryReadTQueue = WrappedSTM . tryReadTQueue
1217+
peekTQueue = WrappedSTM . peekTQueue
1218+
tryPeekTQueue = WrappedSTM . tryPeekTQueue
1219+
writeTQueue v = WrappedSTM . writeTQueue v
1220+
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1221+
unGetTQueue = WrappedSTM .: unGetTQueue
11831222

11841223
type TBQueue (ExceptT e m) = TBQueue m
11851224
newTBQueue = WrappedSTM . newTBQueue
@@ -1192,6 +1231,7 @@ instance MonadSTM m => MonadSTM (ExceptT e m) where
11921231
lengthTBQueue = WrappedSTM . lengthTBQueue
11931232
isEmptyTBQueue = WrappedSTM . isEmptyTBQueue
11941233
isFullTBQueue = WrappedSTM . isFullTBQueue
1234+
unGetTBQueue = WrappedSTM .: unGetTBQueue
11951235

11961236

11971237
instance (Monoid w, MonadSTM m) => MonadSTM (RWST r w s m) where
@@ -1225,12 +1265,13 @@ instance (Monoid w, MonadSTM m) => MonadSTM (RWST r w s m) where
12251265

12261266
type TQueue (RWST r w s m) = TQueue m
12271267
newTQueue = WrappedSTM newTQueue
1228-
readTQueue = WrappedSTM . readTQueue
1229-
tryReadTQueue = WrappedSTM . tryReadTQueue
1230-
peekTQueue = WrappedSTM . peekTQueue
1231-
tryPeekTQueue = WrappedSTM . tryPeekTQueue
1232-
writeTQueue v = WrappedSTM . writeTQueue v
1233-
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1268+
readTQueue = WrappedSTM . readTQueue
1269+
tryReadTQueue = WrappedSTM . tryReadTQueue
1270+
peekTQueue = WrappedSTM . peekTQueue
1271+
tryPeekTQueue = WrappedSTM . tryPeekTQueue
1272+
writeTQueue v = WrappedSTM . writeTQueue v
1273+
isEmptyTQueue = WrappedSTM . isEmptyTQueue
1274+
unGetTQueue = WrappedSTM .: unGetTQueue
12341275

12351276
type TBQueue (RWST r w s m) = TBQueue m
12361277
newTBQueue = WrappedSTM . newTBQueue
@@ -1243,6 +1284,7 @@ instance (Monoid w, MonadSTM m) => MonadSTM (RWST r w s m) where
12431284
lengthTBQueue = WrappedSTM . lengthTBQueue
12441285
isEmptyTBQueue = WrappedSTM . isEmptyTBQueue
12451286
isFullTBQueue = WrappedSTM . isFullTBQueue
1287+
unGetTBQueue = WrappedSTM .: unGetTBQueue
12461288

12471289

12481290
(.:) :: (c -> d) -> (a -> b -> c) -> (a -> b -> d)

io-sim/src/Control/Monad/IOSim/STM.hs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,11 @@ tryPeekTQueueDefault (TQueue queue) = do
8484
x :_ -> Just x
8585
[] -> Nothing
8686

87+
unGetTQueueDefault :: MonadSTM m => TQueueDefault m a -> a -> STM m ()
88+
unGetTQueueDefault (TQueue queue) a = do
89+
(xs, ys) <- readTVar queue
90+
writeTVar queue (a : xs, ys)
91+
8792
--
8893
-- Default TBQueue implementation in terms of 'Seq' (used by sim)
8994
--
@@ -190,3 +195,13 @@ flushTBQueueDefault (TBQueue queue size) = do
190195
else do
191196
writeTVar queue $! ([], 0, [], size)
192197
return (xs ++ reverse ys)
198+
199+
unGetTBQueueDefault :: MonadSTM m => TBQueueDefault m a -> a -> STM m ()
200+
unGetTBQueueDefault (TBQueue queue _size) a = do
201+
(xs, r, ys, w) <- readTVar queue
202+
if (r > 0)
203+
then do writeTVar queue (a : xs, r - 1, ys, w)
204+
else do
205+
if (w > 0)
206+
then writeTVar queue (a : xs, r, ys, w - 1)
207+
else retry

io-sim/src/Control/Monad/IOSim/Types.hs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -407,6 +407,7 @@ instance MonadSTM (IOSim s) where
407407
tryPeekTQueue = tryPeekTQueueDefault
408408
writeTQueue = writeTQueueDefault
409409
isEmptyTQueue = isEmptyTQueueDefault
410+
unGetTQueue = unGetTQueueDefault
410411

411412
newTBQueue = newTBQueueDefault
412413
readTBQueue = readTBQueueDefault
@@ -418,6 +419,7 @@ instance MonadSTM (IOSim s) where
418419
lengthTBQueue = lengthTBQueueDefault
419420
isEmptyTBQueue = isEmptyTBQueueDefault
420421
isFullTBQueue = isFullTBQueueDefault
422+
unGetTBQueue = unGetTBQueueDefault
421423

422424
instance MonadInspectSTM (IOSim s) where
423425
type InspectMonad (IOSim s) = ST s

strict-stm/src/Control/Monad/Class/MonadSTM/Strict.hs

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ module Control.Monad.Class.MonadSTM.Strict
6969
, tryPeekTQueue
7070
, writeTQueue
7171
, isEmptyTQueue
72+
, unGetTQueue
7273
-- * 'StrictTBQueue'
7374
, StrictTBQueue
7475
, labelTBQueue
@@ -87,6 +88,7 @@ module Control.Monad.Class.MonadSTM.Strict
8788
, lengthTBQueue
8889
, isEmptyTBQueue
8990
, isFullTBQueue
91+
, unGetTBQueue
9092
-- ** Low-level API
9193
, checkInvariant
9294
-- * Deprecated API
@@ -110,8 +112,8 @@ import Control.Monad.Class.MonadSTM as X hiding (LazyTMVar, LazyTVar,
110112
traceTBQueue, traceTBQueueIO, traceTMVar, traceTMVarIO,
111113
traceTQueue, traceTQueueIO, traceTVar, traceTVarIO,
112114
tryPeekTBQueue, tryPeekTQueue, tryPutTMVar, tryReadTBQueue,
113-
tryReadTMVar, tryReadTQueue, tryTakeTMVar, writeTBQueue,
114-
writeTQueue, writeTVar)
115+
tryReadTMVar, tryReadTQueue, tryTakeTMVar, unGetTBQueue,
116+
unGetTQueue, writeTBQueue, writeTQueue, writeTVar)
115117
import qualified Control.Monad.Class.MonadSTM as Lazy
116118
import GHC.Stack
117119
import Numeric.Natural (Natural)
@@ -386,6 +388,9 @@ writeTQueue (StrictTQueue tqueue) !a = Lazy.writeTQueue tqueue a
386388
isEmptyTQueue :: MonadSTM m => StrictTQueue m a -> STM m Bool
387389
isEmptyTQueue = Lazy.isEmptyTQueue . toLazyTQueue
388390

391+
unGetTQueue :: MonadSTM m => StrictTQueue m a -> a -> STM m ()
392+
unGetTQueue (StrictTQueue queue) !a = Lazy.unGetTQueue queue a
393+
389394
{-------------------------------------------------------------------------------
390395
Strict TBQueue
391396
-------------------------------------------------------------------------------}
@@ -445,6 +450,9 @@ isEmptyTBQueue = Lazy.isEmptyTBQueue . toLazyTBQueue
445450
isFullTBQueue :: MonadSTM m => StrictTBQueue m a -> STM m Bool
446451
isFullTBQueue = Lazy.isFullTBQueue . toLazyTBQueue
447452

453+
unGetTBQueue :: MonadSTM m => StrictTBQueue m a -> a -> STM m ()
454+
unGetTBQueue (StrictTBQueue queue) !a = Lazy.unGetTBQueue queue a
455+
448456

449457
{-------------------------------------------------------------------------------
450458
Dealing with invariants

0 commit comments

Comments
 (0)