Skip to content

Commit f695250

Browse files
treeowlandrewthad
treeowl
authored and
andrewthad
committed
Add MVar stuff
Add `UnliftedMVar` and a full suite of `MVar` operations.
1 parent b7d63c0 commit f695250

File tree

4 files changed

+487
-0
lines changed

4 files changed

+487
-0
lines changed

‎primitive-unlifted.cabal‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,9 @@ library
3838
Data.Primitive.Unlifted.Weak.IO
3939
Data.Primitive.Unlifted.Weak.Primops
4040
Data.Primitive.TArray.Classic
41+
Data.Primitive.Unlifted.MVar
42+
Data.Primitive.Unlifted.MVar.ST
43+
Data.Primitive.Unlifted.MVar.Primops
4144
build-depends:
4245
, base >=4.14.0.0 && <5
4346
, bytestring >=0.10.8.2 && <0.11

‎src/Data/Primitive/Unlifted/MVar.hs‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
{-# language UnboxedTuples #-}
2+
{-# language UnboxedSums #-}
3+
{-# language RoleAnnotations #-}
4+
{-# language ScopedTypeVariables #-}
5+
{-# language TypeFamilies #-}
6+
{-# language MagicHash #-}
7+
{-# language RankNTypes #-}
8+
{-# language PatternSynonyms #-}
9+
{-# language ViewPatterns #-}
10+
{-# language BangPatterns #-}
11+
{- options_ghc -ddump-simpl #-}
12+
13+
-- | This module includes all the features of "Control.Concurrent.MVar", except
14+
-- that the functions in "Data.Primitive.Unlifted.Weak" subsume the functionality
15+
-- of @mkWeakMV@ and @addMVarFinalizer@, so we do not include analogues of those
16+
-- functions.
17+
module Data.Primitive.Unlifted.MVar
18+
( UnliftedMVar_ (..)
19+
, UnliftedMVar
20+
, newUnliftedMVar
21+
, newEmptyUnliftedMVar
22+
, takeUnliftedMVar
23+
, tryTakeUnliftedMVar
24+
, putUnliftedMVar
25+
, tryPutUnliftedMVar
26+
, readUnliftedMVar
27+
, tryReadUnliftedMVar
28+
, isEmptyUnliftedMVar
29+
, swapUnliftedMVar
30+
, withUnliftedMVar
31+
, withUnliftedMVarMasked
32+
, modifyUnliftedMVar
33+
, modifyUnliftedMVar_
34+
, modifyUnliftedMVarMasked
35+
, modifyUnliftedMVarMasked_
36+
) where
37+
import qualified Data.Primitive.Unlifted.MVar.ST as MV
38+
import Data.Primitive.Unlifted.MVar.ST
39+
( UnliftedMVar_ (..), type UnliftedMVar )
40+
import Data.Primitive.Unlifted.Class (PrimUnlifted (..))
41+
import GHC.Exts (RealWorld)
42+
import Control.Monad.Primitive (stToPrim, PrimMonad (..), PrimBase, primToST)
43+
44+
newUnliftedMVar
45+
:: (PrimUnlifted a, PrimMonad m)
46+
=> a -> m (UnliftedMVar (PrimState m) a)
47+
newUnliftedMVar a = stToPrim $ MV.newUnliftedMVar a
48+
49+
newEmptyUnliftedMVar
50+
:: PrimMonad m
51+
=> m (UnliftedMVar (PrimState m) a)
52+
{-# INLINE newEmptyUnliftedMVar #-}
53+
newEmptyUnliftedMVar = stToPrim $ MV.newEmptyUnliftedMVar
54+
55+
takeUnliftedMVar
56+
:: (PrimMonad m, PrimUnlifted a)
57+
=> UnliftedMVar (PrimState m) a -> m a
58+
{-# INLINE takeUnliftedMVar #-}
59+
takeUnliftedMVar mv = stToPrim $ MV.takeUnliftedMVar mv
60+
61+
tryTakeUnliftedMVar
62+
:: (PrimMonad m, PrimUnlifted a)
63+
=> UnliftedMVar (PrimState m) a -> m (Maybe a)
64+
{-# INLINE tryTakeUnliftedMVar #-}
65+
tryTakeUnliftedMVar mv = stToPrim $ MV.tryTakeUnliftedMVar mv
66+
67+
putUnliftedMVar
68+
:: (PrimMonad m, PrimUnlifted a)
69+
=> UnliftedMVar (PrimState m) a -> a -> m ()
70+
{-# INLINE putUnliftedMVar #-}
71+
putUnliftedMVar mv a = stToPrim $ MV.putUnliftedMVar mv a
72+
73+
tryPutUnliftedMVar
74+
:: (PrimMonad m, PrimUnlifted a)
75+
=> UnliftedMVar (PrimState m) a -> a -> m Bool
76+
{-# INLINE tryPutUnliftedMVar #-}
77+
tryPutUnliftedMVar mv a = stToPrim $ MV.tryPutUnliftedMVar mv a
78+
79+
readUnliftedMVar
80+
:: (PrimMonad m, PrimUnlifted a)
81+
=> UnliftedMVar (PrimState m) a -> m a
82+
{-# INLINE readUnliftedMVar #-}
83+
readUnliftedMVar mv = stToPrim $ MV.readUnliftedMVar mv
84+
85+
tryReadUnliftedMVar
86+
:: (PrimMonad m, PrimUnlifted a)
87+
=> UnliftedMVar (PrimState m) a -> m (Maybe a)
88+
{-# INLINE tryReadUnliftedMVar #-}
89+
tryReadUnliftedMVar mv = stToPrim $ MV.tryReadUnliftedMVar mv
90+
91+
isEmptyUnliftedMVar
92+
:: PrimMonad m
93+
=> UnliftedMVar (PrimState m) a -> m Bool
94+
{-# INLINE isEmptyUnliftedMVar #-}
95+
isEmptyUnliftedMVar mv = stToPrim $ MV.isEmptyUnliftedMVar mv
96+
97+
swapUnliftedMVar
98+
:: (PrimMonad m, PrimState m ~ RealWorld, PrimUnlifted a)
99+
=> UnliftedMVar RealWorld a -> a -> m a
100+
{-# INLINE swapUnliftedMVar #-}
101+
swapUnliftedMVar mvar new = stToPrim $ MV.swapUnliftedMVar mvar new
102+
103+
withUnliftedMVar
104+
:: (PrimBase m, PrimState m ~ RealWorld, PrimUnlifted a)
105+
=> UnliftedMVar RealWorld a -> (a -> m b) -> m b
106+
{-# INLINE withUnliftedMVar #-}
107+
withUnliftedMVar m f = stToPrim $ MV.withUnliftedMVar m (primToST . f)
108+
109+
withUnliftedMVarMasked
110+
:: (PrimBase m, PrimState m ~ RealWorld, PrimUnlifted a)
111+
=> UnliftedMVar RealWorld a -> (a -> m b) -> m b
112+
{-# INLINE withUnliftedMVarMasked #-}
113+
withUnliftedMVarMasked m st = stToPrim $ MV.withUnliftedMVarMasked m (primToST . st)
114+
115+
modifyUnliftedMVar
116+
:: (PrimBase m, PrimState m ~ RealWorld, PrimUnlifted a)
117+
=> UnliftedMVar RealWorld a -> (a -> m (a, b)) -> m b
118+
{-# INLINE modifyUnliftedMVar #-}
119+
modifyUnliftedMVar m st = stToPrim $ MV.modifyUnliftedMVar m (primToST . st)
120+
121+
modifyUnliftedMVar_
122+
:: (PrimBase m, PrimState m ~ RealWorld, PrimUnlifted a)
123+
=> UnliftedMVar RealWorld a -> (a -> m a) -> m ()
124+
{-# INLINE modifyUnliftedMVar_ #-}
125+
modifyUnliftedMVar_ m st = stToPrim $ MV.modifyUnliftedMVar_ m (primToST . st)
126+
127+
modifyUnliftedMVarMasked
128+
:: (PrimBase m, PrimState m ~ RealWorld, PrimUnlifted a)
129+
=> UnliftedMVar RealWorld a -> (a -> m (a, b)) -> m b
130+
{-# INLINE modifyUnliftedMVarMasked #-}
131+
modifyUnliftedMVarMasked m st = stToPrim $ MV.modifyUnliftedMVarMasked m (primToST . st)
132+
133+
modifyUnliftedMVarMasked_
134+
:: (PrimBase m, PrimState m ~ RealWorld, PrimUnlifted a)
135+
=> UnliftedMVar RealWorld a -> (a -> m a) -> m ()
136+
{-# INLINE modifyUnliftedMVarMasked_ #-}
137+
modifyUnliftedMVarMasked_ m st = stToPrim $ MV.modifyUnliftedMVarMasked_ m (primToST . st)

‎src/Data/Primitive/Unlifted/MVar/Primops.hs‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
{-# language ScopedTypeVariables #-}
2+
{-# language MagicHash #-}
3+
{-# language KindSignatures #-}
4+
{-# language UnboxedTuples #-}
5+
{-# language UnboxedSums #-}
6+
{-# language UnliftedNewtypes #-}
7+
{-# language RoleAnnotations #-}
8+
{-# language DataKinds #-}
9+
10+
module Data.Primitive.Unlifted.MVar.Primops
11+
( UnliftedMVar#
12+
, newUnliftedMVar#
13+
, takeUnliftedMVar#
14+
, tryTakeUnliftedMVar#
15+
, putUnliftedMVar#
16+
, tryPutUnliftedMVar#
17+
, readUnliftedMVar#
18+
, tryReadUnliftedMVar#
19+
, sameUnliftedMVar#
20+
, isEmptyUnliftedMVar#
21+
) where
22+
import GHC.Exts
23+
24+
newtype UnliftedMVar# s (a :: TYPE 'UnliftedRep) = UnliftedMVar# (MVar# s Any)
25+
type role UnliftedMVar# nominal representational
26+
27+
newUnliftedMVar# :: State# s -> (# State# s, UnliftedMVar# s a #)
28+
{-# INLINE newUnliftedMVar# #-}
29+
newUnliftedMVar# s = case newMVar# s of
30+
(# s', mv #) -> (# s', UnliftedMVar# mv #)
31+
32+
takeUnliftedMVar# :: UnliftedMVar# s a -> State# s -> (# State# s, a #)
33+
{-# NOINLINE takeUnliftedMVar# #-}
34+
takeUnliftedMVar# (UnliftedMVar# mv) s = unsafeCoerce# (takeMVar# mv s)
35+
36+
tryTakeUnliftedMVar# :: UnliftedMVar# s a -> State# s -> (# State# s, (# (##) | a #) #)
37+
{-# NOINLINE tryTakeUnliftedMVar# #-}
38+
tryTakeUnliftedMVar# (UnliftedMVar# mv) s =
39+
case unsafeCoerce# (tryTakeMVar# mv s) of
40+
(# s', 0#, _ #) -> (# s', (#(##)| #)#)
41+
(# s', _, a #) -> (# s', (#|a #) #)
42+
43+
putUnliftedMVar# :: UnliftedMVar# s a -> a -> State# s -> State# s
44+
{-# NOINLINE putUnliftedMVar# #-}
45+
putUnliftedMVar# (UnliftedMVar# mv) a s
46+
= putMVar# mv (unsafeCoerce# a) s
47+
48+
tryPutUnliftedMVar# :: UnliftedMVar# s a -> a -> State# s -> (# State# s, Int# #)
49+
{-# NOINLINE tryPutUnliftedMVar# #-}
50+
tryPutUnliftedMVar# (UnliftedMVar# mv) a s
51+
= tryPutMVar# mv (unsafeCoerce# a) s
52+
53+
readUnliftedMVar# :: UnliftedMVar# s a -> State# s -> (# State# s, a #)
54+
{-# NOINLINE readUnliftedMVar# #-}
55+
readUnliftedMVar# (UnliftedMVar# mv) s = unsafeCoerce# (readMVar# mv s)
56+
57+
tryReadUnliftedMVar# :: UnliftedMVar# s a -> State# s -> (# State# s, (# (##) | a #) #)
58+
{-# NOINLINE tryReadUnliftedMVar# #-}
59+
tryReadUnliftedMVar# (UnliftedMVar# mv) s =
60+
case unsafeCoerce# (tryReadMVar# mv s) of
61+
(# s', 0#, _ #) -> (# s', (#(##)| #)#)
62+
(# s', _, a #) -> (# s', (#|a #) #)
63+
64+
sameUnliftedMVar# :: UnliftedMVar# s a -> UnliftedMVar# s a -> Int#
65+
{-# INLINE sameUnliftedMVar# #-}
66+
sameUnliftedMVar# (UnliftedMVar# mv1) (UnliftedMVar# mv2)
67+
= sameMVar# mv1 mv2
68+
69+
isEmptyUnliftedMVar# :: UnliftedMVar# s a -> State# s -> (# State# s, Int# #)
70+
{-# INLINE isEmptyUnliftedMVar# #-}
71+
isEmptyUnliftedMVar# (UnliftedMVar# mv) s
72+
= isEmptyMVar# mv s

0 commit comments

Comments
 (0)