module ZooKeeper.Internal.Utils
(
Resource (..)
, initResource
, initResource_
, withResource
, withResource'
) where
import Control.Exception (onException)
import Control.Monad (when)
import qualified Control.Monad.Catch as MonadCatch
import Control.Monad.IO.Class (MonadIO, liftIO)
import GHC.Stack (HasCallStack)
import Z.Data.PrimRef (atomicOrCounter, newCounter)
newtype Resource a = Resource { forall a. Resource a -> IO (a, IO ())
acquire :: IO (a, IO ()) }
initResource :: IO a -> (a -> IO ()) -> Resource a
{-# INLINABLE initResource #-}
initResource :: forall a. IO a -> (a -> IO ()) -> Resource a
initResource IO a
create a -> IO ()
release = IO (a, IO ()) -> Resource a
forall a. IO (a, IO ()) -> Resource a
Resource (IO (a, IO ()) -> Resource a) -> IO (a, IO ()) -> Resource a
forall a b. (a -> b) -> a -> b
$ do
a
r <- IO a
create
(a, IO ()) -> IO (a, IO ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
r, a -> IO ()
release a
r)
initResource_ :: IO a -> IO () -> Resource a
{-# INLINABLE initResource_ #-}
initResource_ :: forall a. IO a -> IO () -> Resource a
initResource_ IO a
create IO ()
release = IO (a, IO ()) -> Resource a
forall a. IO (a, IO ()) -> Resource a
Resource (IO (a, IO ()) -> Resource a) -> IO (a, IO ()) -> Resource a
forall a b. (a -> b) -> a -> b
$ do
a
r <- IO a
create
(a, IO ()) -> IO (a, IO ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
r, IO ()
release)
instance Functor Resource where
{-# INLINE fmap #-}
fmap :: forall a b. (a -> b) -> Resource a -> Resource b
fmap a -> b
f Resource a
resource = IO (b, IO ()) -> Resource b
forall a. IO (a, IO ()) -> Resource a
Resource (IO (b, IO ()) -> Resource b) -> IO (b, IO ()) -> Resource b
forall a b. (a -> b) -> a -> b
$ do
(a
a, IO ()
release) <- Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
resource
(b, IO ()) -> IO (b, IO ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b
f a
a, IO ()
release)
instance Applicative Resource where
{-# INLINE pure #-}
pure :: forall a. a -> Resource a
pure a
a = IO (a, IO ()) -> Resource a
forall a. IO (a, IO ()) -> Resource a
Resource ((a, IO ()) -> IO (a, IO ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()))
{-# INLINE (<*>) #-}
Resource (a -> b)
resource1 <*> :: forall a b. Resource (a -> b) -> Resource a -> Resource b
<*> Resource a
resource2 = IO (b, IO ()) -> Resource b
forall a. IO (a, IO ()) -> Resource a
Resource (IO (b, IO ()) -> Resource b) -> IO (b, IO ()) -> Resource b
forall a b. (a -> b) -> a -> b
$ do
(a -> b
f, IO ()
release1) <- Resource (a -> b) -> IO (a -> b, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource (a -> b)
resource1
(a
x, IO ()
release2) <- Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
resource2 IO (a, IO ()) -> IO () -> IO (a, IO ())
forall a b. IO a -> IO b -> IO a
`onException` IO ()
release1
(b, IO ()) -> IO (b, IO ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b
f a
x, IO ()
release2 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
release1)
instance Monad Resource where
{-# INLINE return #-}
return :: forall a. a -> Resource a
return = a -> Resource a
forall a. a -> Resource a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
{-# INLINE (>>=) #-}
Resource a
m >>= :: forall a b. Resource a -> (a -> Resource b) -> Resource b
>>= a -> Resource b
f = IO (b, IO ()) -> Resource b
forall a. IO (a, IO ()) -> Resource a
Resource (IO (b, IO ()) -> Resource b) -> IO (b, IO ()) -> Resource b
forall a b. (a -> b) -> a -> b
$ do
(a
m', IO ()
release1) <- Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
m
(b
x , IO ()
release2) <- Resource b -> IO (b, IO ())
forall a. Resource a -> IO (a, IO ())
acquire (a -> Resource b
f a
m') IO (b, IO ()) -> IO () -> IO (b, IO ())
forall a b. IO a -> IO b -> IO a
`onException` IO ()
release1
(b, IO ()) -> IO (b, IO ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (b
x, IO ()
release2 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
release1)
instance MonadIO Resource where
{-# INLINE liftIO #-}
liftIO :: forall a. IO a -> Resource a
liftIO IO a
f = IO (a, IO ()) -> Resource a
forall a. IO (a, IO ()) -> Resource a
Resource (IO (a, IO ()) -> Resource a) -> IO (a, IO ()) -> Resource a
forall a b. (a -> b) -> a -> b
$ (a -> (a, IO ())) -> IO a -> IO (a, IO ())
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ a
a -> (a
a, IO ()
dummyRelease)) IO a
f
where dummyRelease :: IO ()
dummyRelease = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
withResource :: (MonadCatch.MonadMask m, MonadIO m, HasCallStack)
=> Resource a -> (HasCallStack => a -> m b) -> m b
{-# INLINABLE withResource #-}
withResource :: forall (m :: * -> *) a b.
(MonadMask m, MonadIO m, HasCallStack) =>
Resource a -> (HasCallStack => a -> m b) -> m b
withResource Resource a
resource HasCallStack => a -> m b
k = m (a, IO ()) -> ((a, IO ()) -> m ()) -> ((a, IO ()) -> m b) -> m b
forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
MonadCatch.bracket
(IO (a, IO ()) -> m (a, IO ())
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
resource))
(\(a
_, IO ()
release) -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
release)
(\(a
a, IO ()
_) -> a -> m b
HasCallStack => a -> m b
k a
a)
withResource' :: (MonadCatch.MonadMask m, MonadIO m, HasCallStack)
=> Resource a -> (HasCallStack => a -> m () -> m b) -> m b
{-# INLINABLE withResource' #-}
withResource' :: forall (m :: * -> *) a b.
(MonadMask m, MonadIO m, HasCallStack) =>
Resource a -> (HasCallStack => a -> m () -> m b) -> m b
withResource' Resource a
resource HasCallStack => a -> m () -> m b
k = do
Counter
c <- IO Counter -> m Counter
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Int -> IO Counter
newCounter Int
0)
m (a, IO ()) -> ((a, IO ()) -> m ()) -> ((a, IO ()) -> m b) -> m b
forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
MonadCatch.bracket
(IO (a, IO ()) -> m (a, IO ())
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (a, IO ()) -> m (a, IO ())) -> IO (a, IO ()) -> m (a, IO ())
forall a b. (a -> b) -> a -> b
$ do
(a
a, IO ()
release) <- Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
resource
let release' :: IO ()
release' = do
Int
c' <- Counter -> Int -> IO Int
atomicOrCounter Counter
c Int
1
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
c' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) IO ()
release
(a, IO ()) -> IO (a, IO ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, IO ()
release'))
(\(a
_, IO ()
release) -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
release)
(\(a
a, IO ()
release) -> a -> m () -> m b
HasCallStack => a -> m () -> m b
k a
a (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
release))