{-# LANGUAGE DataKinds    #-}
{-# LANGUAGE TypeFamilies #-}
-----------------------------------------------------------------------------

-- |

-- Module      : Effectful.Resource

-- Copyright   : (c) Michael Szvetits, 2023

-- License     : BSD-3-Clause (see the file LICENSE)

-- Maintainer  : typedbyte@qualified.name

-- Stability   : stable

-- Portability : portable

--

-- A region-based resource effect for the effectful ecosystem.

-----------------------------------------------------------------------------

module Effectful.Resource
  ( -- * Resource Effect

    Resource
  , runResource
  -- * Manage Regions

  , Region
  , withRegion
  , currentRegion
  -- * Manage Resources

  , Key
  , InvalidKey(..)
  , manage
  , allocate
  , free
  , freeAll
  , move
  , move_
  , defer
  ) where

-- base

import Control.Exception (Exception, bracket, finally, mask_, uninterruptibleMask_)
import Control.Monad     (join)
import Data.Functor      (void)

-- effectful-core

import Effectful                 (Dispatch(..), DispatchOf, Eff, Effect, IOE, (:>))
import Effectful.Dispatch.Static (SideEffects(..), StaticRep, evalStaticRep,
                                  getStaticRep, localStaticRep, unsafeEff_,
                                  unsafeSeqUnliftIO)

-- stm

import Control.Concurrent.STM (TVar, atomically, modifyTVar', newTVarIO, readTVar,
                               stateTVar, throwSTM, writeTVar)

-- | A region owns resources and frees them on close.

data Region = Region
  { Region -> TVar [(Key, IO ())]
resources :: TVar [(Key, IO ())]
  , Region -> TVar Int
nextID    :: TVar Int
  }
  deriving Region -> Region -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Region -> Region -> Bool
$c/= :: Region -> Region -> Bool
== :: Region -> Region -> Bool
$c== :: Region -> Region -> Bool
Eq

-- | Each resource is identified by a unique key.

data Key = Key
  { Key -> Int
_keyID    :: Int
  , Key -> Region
keyRegion :: Region
  }
  deriving Key -> Key -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Key -> Key -> Bool
$c/= :: Key -> Key -> Bool
== :: Key -> Key -> Bool
$c== :: Key -> Key -> Bool
Eq

open :: IO Region
open :: IO Region
open =
  TVar [(Key, IO ())] -> TVar Int -> Region
Region
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. a -> IO (TVar a)
newTVarIO []
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. a -> IO (TVar a)
newTVarIO Int
0

close :: Region -> IO ()
close :: Region -> IO ()
close Region
region = forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ do
  [(Key, IO ())]
rs <-
    forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$
      forall s a. TVar s -> (s -> (a, s)) -> STM a
stateTVar
        ( Region -> TVar [(Key, IO ())]
resources Region
region )
        ( \[(Key, IO ())]
r -> ([(Key, IO ())]
r, []) )
  forall {a}. [(a, IO ())] -> IO ()
freeList [(Key, IO ())]
rs
    where
      freeList :: [(a, IO ())] -> IO ()
freeList []         = forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      freeList ((a
_,IO ()
m):[(a, IO ())]
ms) = IO ()
m forall a b. IO a -> IO b -> IO a
`finally` [(a, IO ())] -> IO ()
freeList [(a, IO ())]
ms

manageIO
  :: Region
  -> IO a
  -> (a -> IO ())
  -> IO (a, Key)
manageIO :: forall a. Region -> IO a -> (a -> IO ()) -> IO (a, Key)
manageIO Region
region IO a
create a -> IO ()
destroy = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
  a
a <- IO a
create
  forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
    Int
next <- forall a. TVar a -> STM a
readTVar TVar Int
idTVar
    let key :: Key
key = Int -> Region -> Key
Key Int
next Region
region
    forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
idTVar forall a. Enum a => a -> a
succ
    forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar [(Key, IO ())]
rsTVar ((Key
key, a -> IO ()
destroy a
a) :)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, Key
key)
  where
    idTVar :: TVar Int
idTVar = Region -> TVar Int
nextID Region
region
    rsTVar :: TVar [(Key, IO ())]
rsTVar = Region -> TVar [(Key, IO ())]
resources Region
region

-- | An error which occurs if a key is freed\/moved that has already been freed\/moved.

data InvalidKey = InvalidKey
  deriving (InvalidKey -> InvalidKey -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: InvalidKey -> InvalidKey -> Bool
$c/= :: InvalidKey -> InvalidKey -> Bool
== :: InvalidKey -> InvalidKey -> Bool
$c== :: InvalidKey -> InvalidKey -> Bool
Eq, Eq InvalidKey
InvalidKey -> InvalidKey -> Bool
InvalidKey -> InvalidKey -> Ordering
InvalidKey -> InvalidKey -> InvalidKey
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: InvalidKey -> InvalidKey -> InvalidKey
$cmin :: InvalidKey -> InvalidKey -> InvalidKey
max :: InvalidKey -> InvalidKey -> InvalidKey
$cmax :: InvalidKey -> InvalidKey -> InvalidKey
>= :: InvalidKey -> InvalidKey -> Bool
$c>= :: InvalidKey -> InvalidKey -> Bool
> :: InvalidKey -> InvalidKey -> Bool
$c> :: InvalidKey -> InvalidKey -> Bool
<= :: InvalidKey -> InvalidKey -> Bool
$c<= :: InvalidKey -> InvalidKey -> Bool
< :: InvalidKey -> InvalidKey -> Bool
$c< :: InvalidKey -> InvalidKey -> Bool
compare :: InvalidKey -> InvalidKey -> Ordering
$ccompare :: InvalidKey -> InvalidKey -> Ordering
Ord, Int -> InvalidKey -> ShowS
[InvalidKey] -> ShowS
InvalidKey -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [InvalidKey] -> ShowS
$cshowList :: [InvalidKey] -> ShowS
show :: InvalidKey -> String
$cshow :: InvalidKey -> String
showsPrec :: Int -> InvalidKey -> ShowS
$cshowsPrec :: Int -> InvalidKey -> ShowS
Show)

instance Exception InvalidKey

extract :: Eq k => k -> [(k,v)] -> Maybe (v, [(k,v)])
extract :: forall k v. Eq k => k -> [(k, v)] -> Maybe (v, [(k, v)])
extract k
k = ([(k, v)] -> [(k, v)]) -> [(k, v)] -> Maybe (v, [(k, v)])
extract' forall a. a -> a
id
  where
    extract' :: ([(k, v)] -> [(k, v)]) -> [(k, v)] -> Maybe (v, [(k, v)])
extract' [(k, v)] -> [(k, v)]
_ [] = forall a. Maybe a
Nothing
    extract' [(k, v)] -> [(k, v)]
f (x :: (k, v)
x@(k
k',v
v):[(k, v)]
xs)
      | k
k forall a. Eq a => a -> a -> Bool
== k
k'   = forall a. a -> Maybe a
Just (v
v, [(k, v)] -> [(k, v)]
f [(k, v)]
xs)
      | Bool
otherwise = ([(k, v)] -> [(k, v)]) -> [(k, v)] -> Maybe (v, [(k, v)])
extract' ([(k, v)] -> [(k, v)]
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v)
x:)) [(k, v)]
xs

moveIO :: Key -> Region -> IO Key
moveIO :: Key -> Region -> IO Key
moveIO Key
key Region
region = forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
  [(Key, IO ())]
rs <- forall a. TVar a -> STM a
readTVar TVar [(Key, IO ())]
keyRsTVar
  case forall k v. Eq k => k -> [(k, v)] -> Maybe (v, [(k, v)])
extract Key
key [(Key, IO ())]
rs of
    Maybe (IO (), [(Key, IO ())])
Nothing -> forall e a. Exception e => e -> STM a
throwSTM InvalidKey
InvalidKey
    Just (IO ()
m, [(Key, IO ())]
rs') -> do
      forall a. TVar a -> a -> STM ()
writeTVar TVar [(Key, IO ())]
keyRsTVar [(Key, IO ())]
rs'
      Int
next <- forall a. TVar a -> STM a
readTVar TVar Int
idTVar
      forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
idTVar forall a. Enum a => a -> a
succ
      let newKey :: Key
newKey = Int -> Region -> Key
Key Int
next Region
region
      forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar [(Key, IO ())]
rsTVar ((Key
newKey, IO ()
m) :)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure Key
newKey
  where
    keyRsTVar :: TVar [(Key, IO ())]
keyRsTVar = Region -> TVar [(Key, IO ())]
resources forall a b. (a -> b) -> a -> b
$ Key -> Region
keyRegion Key
key
    idTVar :: TVar Int
idTVar    = Region -> TVar Int
nextID Region
region
    rsTVar :: TVar [(Key, IO ())]
rsTVar    = Region -> TVar [(Key, IO ())]
resources Region
region

-- | The region-based resource effect.

data Resource :: Effect

type instance DispatchOf Resource = Static WithSideEffects

newtype instance StaticRep Resource = Resource Region

-- | Runs the resource effect.

runResource :: IOE :> es => Eff (Resource : es) a -> Eff es a
runResource :: forall (es :: [Effect]) a.
(IOE :> es) =>
Eff (Resource : es) a -> Eff es a
runResource Eff (Resource : es) a
m =
  forall (es :: [Effect]) a.
HasCallStack =>
((forall r. Eff es r -> IO r) -> IO a) -> Eff es a
unsafeSeqUnliftIO forall a b. (a -> b) -> a -> b
$ \forall r. Eff es r -> IO r
run ->
    forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO Region
open Region -> IO ()
close forall a b. (a -> b) -> a -> b
$ \Region
emptyRegion ->
      forall r. Eff es r -> IO r
run forall a b. (a -> b) -> a -> b
$
        forall (e :: Effect) (sideEffects :: SideEffects) (es :: [Effect])
       a.
(DispatchOf e ~ 'Static sideEffects, MaybeIOE sideEffects es) =>
StaticRep e -> Eff (e : es) a -> Eff es a
evalStaticRep (Region -> StaticRep Resource
Resource Region
emptyRegion) Eff (Resource : es) a
m

-- | Runs a computation in a new region.

withRegion :: Resource :> es => Eff es a -> Eff es a
withRegion :: forall (es :: [Effect]) a. (Resource :> es) => Eff es a -> Eff es a
withRegion Eff es a
m =
  forall (es :: [Effect]) a.
HasCallStack =>
((forall r. Eff es r -> IO r) -> IO a) -> Eff es a
unsafeSeqUnliftIO forall a b. (a -> b) -> a -> b
$ \forall r. Eff es r -> IO r
run ->
    forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO Region
open Region -> IO ()
close forall a b. (a -> b) -> a -> b
$ \Region
emptyRegion ->
      forall r. Eff es r -> IO r
run forall a b. (a -> b) -> a -> b
$
        forall (e :: Effect) (sideEffects :: SideEffects) (es :: [Effect])
       a.
(DispatchOf e ~ 'Static sideEffects, e :> es) =>
(StaticRep e -> StaticRep e) -> Eff es a -> Eff es a
localStaticRep (\StaticRep Resource
_ -> Region -> StaticRep Resource
Resource Region
emptyRegion) Eff es a
m

-- | Gets the current region.

currentRegion :: Resource :> es => Eff es Region
currentRegion :: forall (es :: [Effect]). (Resource :> es) => Eff es Region
currentRegion = do
  Resource Region
region <- forall (e :: Effect) (sideEffects :: SideEffects) (es :: [Effect]).
(DispatchOf e ~ 'Static sideEffects, e :> es) =>
Eff es (StaticRep e)
getStaticRep
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Region
region

-- | Allocates a resource in the current region which can be moved and freed

-- manually using its key.

allocate
  :: Resource :> es
  => IO a            -- ^ The computation which acquires the resource.

  -> (a -> IO b)     -- ^ The computation which releases the resource.

  -> Eff es (a, Key) -- ^ The acquired resource and its corresponding key.

allocate :: forall (es :: [Effect]) a b.
(Resource :> es) =>
IO a -> (a -> IO b) -> Eff es (a, Key)
allocate IO a
create a -> IO b
destroy = do
  Resource Region
region <- forall (e :: Effect) (sideEffects :: SideEffects) (es :: [Effect]).
(DispatchOf e ~ 'Static sideEffects, e :> es) =>
Eff es (StaticRep e)
getStaticRep
  forall a (es :: [Effect]). IO a -> Eff es a
unsafeEff_ forall a b. (a -> b) -> a -> b
$ forall a. Region -> IO a -> (a -> IO ()) -> IO (a, Key)
manageIO Region
region IO a
create (forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> IO b
destroy)

-- | Allocates a resource in the current region which is automatically freed at

-- the end of the region.

manage
  :: Resource :> es
  => IO a        -- ^ The computation which acquires the resource.

  -> (a -> IO b) -- ^ The computation which releases the resource.

  -> Eff es a    -- ^ The acquired resource.

manage :: forall (es :: [Effect]) a b.
(Resource :> es) =>
IO a -> (a -> IO b) -> Eff es a
manage IO a
create a -> IO b
destroy =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]) a b.
(Resource :> es) =>
IO a -> (a -> IO b) -> Eff es (a, Key)
allocate IO a
create a -> IO b
destroy

-- | Moves a resource to the specified region, yielding a new key for the resource.

-- The old key is invalid after the movement.

move :: Resource :> es => Key -> Region -> Eff es Key
move :: forall (es :: [Effect]).
(Resource :> es) =>
Key -> Region -> Eff es Key
move Key
key Region
region = 
  forall a (es :: [Effect]). IO a -> Eff es a
unsafeEff_ forall a b. (a -> b) -> a -> b
$ Key -> Region -> IO Key
moveIO Key
key Region
region

-- | Moves a resource to the specified region. It is freed at the end of this region.

-- The key of the moved resource is invalid after the movement.

move_ :: Resource :> es => Key -> Region -> Eff es ()
move_ :: forall (es :: [Effect]).
(Resource :> es) =>
Key -> Region -> Eff es ()
move_ Key
key Region
region =
  forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (es :: [Effect]).
(Resource :> es) =>
Key -> Region -> Eff es Key
move Key
key Region
region

freeIO :: Key -> IO ()
freeIO :: Key -> IO ()
freeIO Key
key =
  forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
    [(Key, IO ())]
rs <- forall a. TVar a -> STM a
readTVar TVar [(Key, IO ())]
keyRsTVar
    case forall k v. Eq k => k -> [(k, v)] -> Maybe (v, [(k, v)])
extract Key
key [(Key, IO ())]
rs of
      Maybe (IO (), [(Key, IO ())])
Nothing -> forall e a. Exception e => e -> STM a
throwSTM InvalidKey
InvalidKey
      Just (IO ()
m, [(Key, IO ())]
rs') -> do
        forall a. TVar a -> a -> STM ()
writeTVar TVar [(Key, IO ())]
keyRsTVar [(Key, IO ())]
rs'
        forall (f :: * -> *) a. Applicative f => a -> f a
pure IO ()
m
  where
    keyRsTVar :: TVar [(Key, IO ())]
keyRsTVar = Region -> TVar [(Key, IO ())]
resources forall a b. (a -> b) -> a -> b
$ Key -> Region
keyRegion Key
key
{-# INLINE freeIO #-}

-- | Frees a resource manually.

free :: Key -> Eff es ()
free :: forall (es :: [Effect]). Key -> Eff es ()
free = forall a (es :: [Effect]). IO a -> Eff es a
unsafeEff_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IO a -> IO a
uninterruptibleMask_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> IO ()
freeIO

-- | Frees a collection of resources manually.

freeAll :: Foldable t => t Key -> Eff es ()
freeAll :: forall (t :: * -> *) (es :: [Effect]).
Foldable t =>
t Key -> Eff es ()
freeAll = forall a (es :: [Effect]). IO a -> Eff es a
unsafeEff_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IO a -> IO a
uninterruptibleMask_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Key -> IO ()
freeIO

-- | Associats a cleanup action with the current region which is executed when the

-- region is closed.

defer :: Resource :> es => IO a -> Eff es ()
defer :: forall (es :: [Effect]) a. (Resource :> es) => IO a -> Eff es ()
defer IO a
action =
  forall (es :: [Effect]) a b.
(Resource :> es) =>
IO a -> (a -> IO b) -> Eff es a
manage (forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (forall a b. a -> b -> a
const IO a
action)