Transactional.hs

module Sq.Transactional
   ( Transactional
   , embed
   , transactionalRetry
   , one
   , maybe
   , zero
   , some
   , list
   , fold
   , foldM
   , Ref
   , Retry (..)
   , retry
   , orElse
   ) where

import Control.Applicative hiding (some)
import Control.Concurrent
import Control.Concurrent.STM hiding (orElse, retry)
import Control.Exception.Safe qualified as Ex
import Control.Foldl qualified as F
import Control.Monad hiding (foldM)
import Control.Monad.Catch qualified as Cx
import Control.Monad.IO.Class
import Control.Monad.Ref hiding (Ref)
import Control.Monad.Ref qualified
import Control.Monad.Trans.Reader (ReaderT (ReaderT))
import Control.Monad.Trans.Resource qualified as R
import Control.Monad.Trans.Resource.Extra qualified as R hiding (runResourceT)
import Data.Acquire qualified as A
import Data.Coerce
import Data.Int
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Kind
import Data.List.NonEmpty (NonEmpty)
import Prelude hiding (Read, maybe, read)

import Sq.Connection
import Sq.Mode
import Sq.Statement
import Sq.Support

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

-- | Used as the @r@ type-parameter in @'Transactional' g r t a@.
--
-- * If the 'Transactional' uses any 'Alternative' or 'MonadPlus' feature, then
-- @r@ must be 'Retry', and the 'Transactional' can only be executed through
-- 'Sq.read', 'Sq.commit' or 'Sq.rollback'.
--
-- * Otherwise, @r@ can be 'NoRetry'. In that case, 'embed' can
-- also be used to execute the 'Transactional'.
data Retry = NoRetry | Retry
   deriving (Eq, Ord, Show)

data Env (g :: k) (r :: Retry) (t :: Mode) = Env
   { unique :: STM Int
   -- ^ Next unique 'Int' within the 'Transactional' to be used as key in 'refs'
   , refs :: TVar (IntMap (SomeRef g))
   -- ^ Currently valid 'Ref's. We keep track of them in order to implement
   -- 'catch'. The 'IntMap' is just for fast diffing purposes.
   , tx :: Transaction t
   -- ^ Current transaction.
   }

acquireEnv :: Transaction t -> A.Acquire (Env g r t)
acquireEnv tx = do
   unique :: STM Int <- liftIO do
      tv <- newTVarIO 0
      pure $ atomicModifyRef' tv \i -> (i + 1, i)
   refs :: TVar (IntMap (SomeRef g)) <-
      R.mkAcquire1 (newTVarIO mempty) \tvsrs ->
         atomically do
            srs <- swapTVar tvsrs mempty
            forM_ srs \(SomeRef (Ref tv)) ->
               writeTVar tv Nothing
   pure Env{..}

-- | @'Transactional' g r t a@ groups together multiple interactions with a same
-- @'Transaction' t@ that finally produce a value of type @a@. Think of
-- 'Transactional' as if it was 'STM'.
--
-- * @g@ is an ephemeral tag for the whole inteaction group that prevents
-- 'Ref's and 'stream's from escaping its intended scope (like 'Data.STRef.ST'
-- does it). Just ignore it, it will always be polymorphic.
--
-- * @r@ says whether the 'Transactional' could potentially be retried from
-- scratch in order to observe a new snapshot of the database (like 'STM' does
-- it).  Learn more about this in 'Retry'.
--
-- * @t@ says whether the 'Transactional' could potentially perform 'Write'
-- or 'Read'-only operations.
--
-- * @a@ is the Haskell value finally produced by a successfu execution of
-- the 'Transactional'.
--
-- __To execute a 'Transactional'__ you will normally use one of 'Sq.read' or
-- 'Sq.commit' (or 'Sq.rollback' or 'Sq.embed', but those are less common).
--
-- @
-- /-- We are using 'Sq.commit' to execute the 'Transactional'. This means/
-- /-- that the 'Transactional' will have read and 'Write' capabilities, that/
-- /-- it can 'retry', and that ultimately, unless there are unhandled/
-- /-- exceptions, the changes will be commited to the database./
-- __"Sq".'Sq.commit' pool do__
--
--    /-- We can execute 'Write' 'Statement's:/
--    __userId1 <- "Sq".'Sq.one' /insertUser/ \"haskell\@example.com\"__
--
--    /-- And 'Read' 'Statement's:/
--    __userId2 \<- "Sq".'Sq.one' /getUserIdByEmail/ \"haskell\@example.com\"__
--
--    /-- We have 'MonadFail' too:/
--    __'when' (userId1 /= userId2) do__
--        __'fail' \"Something unexpected happened!\"__
--
--    /-- We also have 'Ref's, which work just like 'TVar's:/
--    __ref \<- 'newRef' (0 :: 'Int')__
--
--    /-- 'Ex.catch' behaves like 'catchSTM', undoing changes to 'Ref's/
--    /-- and to the database itself when the original action fails:/
--    __userId3 \<- 'Ex.catch'__
--        /-- Something will fail .../
--        __(do 'modifyRef' ref (+ 1)__
--            __\_ \<- "Sq".'Sq.one' /insertUser/ \"sqlite\@example.com\"__
--            __'Ex.throwM' FakeException123)__
--        /-- ... but there is a catch!/
--        __(\\FakeException123 -> do__
--            /-- The observable universe has been reset to what it/
--            /-- was before the 'Ex.catch':/
--            __"Sq".'Sq.zero' /getUserIdByEmail/ \"sqlite\@example.com\"__
--            __'modifyRef' ref (+ 10))__
--
--    /-- Only the effects from the exception handling function were preserved:/
--    __"Sq".'Sq.zero' /getUserIdByEmail/ \"sqlite\@example.com\"__
--    __10 <- 'readRef' ref__
--
--    /-- 'retry' and its synonyms 'mzero' and 'empty' not only discard changes as/
--    /-- 'Ex.catch' does, but they also cause the ongoing 'Transaction' to be/
--    /-- discarded, and the entire 'Transactional' to be executed again on a/
--    /-- brand new 'Transaction' observing a new snapshot of the database. For/
--    /-- example, the following code will keep retrying the whole 'Transactional'/
--    /-- until the user with the specified email exists./
--    __userId4 \<- "Sq".'maybe' /getUserIdByEmail/ \"nix@example.com\" >>= \\case__
--        __'Just' x -> 'pure' x__
--        __'Nothing' -> 'retry'__
--
--    /-- Presumably, this example was waiting for a concurrent connection to/
--    /-- insert said user. If we got here, it is because that happened./
--
--    /-- As usual, 'mzero' and 'empty' can be handled by means of '<|>' and 'mplus',/
--    /-- or its synonym 'orElse'./
--    __'False' \<- 'mplus' ('writeRef' ref 8 >> 'mzero' >> 'pure' 'True')__
--                   __('pure' 'False')__
--
--    /-- The recent 'writeRef' to 8 on the 'retry'ied 'Transactional' was discarded:/
--    __10 <- 'readRef' ref__
--
--    __'pure' ()__
-- @
newtype Transactional (g :: k) (r :: Retry) (t :: Mode) (a :: Type)
   = Transactional (Env g r t -> R.ResourceT IO a)
   deriving
      ( Functor
      , Applicative
      , Monad
      , Cx.MonadThrow
      , Cx.MonadMask
      , MonadFail
      )
      via (ReaderT (Env g r t) (R.ResourceT IO))

-- | INTERNAL only. This doesn't deal with @g@.
un :: Transactional g r t a -> Env g r t -> R.ResourceT IO a
un = coerce
{-# INLINE un #-}

mk :: (Env g r t -> R.ResourceT IO a) -> Transactional g r t a
mk = coerce
{-# INLINE mk #-}

-- | INTERNAL. Used to implement 'Sq.read', 'Sq.commit' and 'Sq.rollback'.
--
-- Run all the actions in a 'Transactional' as part of a single 'Transaction'.
transactionalRetry
   :: forall m r t a
    . (MonadIO m)
   => A.Acquire (Transaction t)
   -> (forall g. Transactional g r t a)
   -> m a
transactionalRetry atx ta = liftIO (go 0)
  where
   go :: Word -> IO a
   go !n = Ex.catch once \ErrRetry -> do
      -- TODO: Wait with `sqlite3_commit_hook` instead of just retrying.
      let ms = logBase 2 (fromIntegral (max 1 n) :: Double)
      threadDelay $ truncate (1_000 * ms)
      go (n + 1)
   once :: IO a
   once = R.runResourceT do
      (_, env) <- A.allocateAcquire $ acquireEnv =<< atx
      un ta env

-- | Embeds all the actions in a 'Transactional' as part of an ongoing
-- 'Transaction'.
--
-- * __NOTICE__ Contrary to 'Sq.read', 'Sq.commit' or 'Sq.rollback',
-- this 'Transactional' cannot 'retry', as doing so would require
-- cancelling the ongoing 'Transaction'.
embed
   :: forall m t a
    . (MonadIO m)
   => Transaction t
   -- ^ Ongoing transaction.
   -> (forall g. Transactional g 'NoRetry t a)
   -> m a
embed tx ta =
   liftIO $ R.runResourceT do
      (_, env) <- A.allocateAcquire $ acquireEnv tx
      un ta env

-- | __Impurely fold__ the output rows.
--
-- * For a non-'Transactional' version of this function, see 'Sq.foldIO'.
foldM
   :: forall o z i t s g r
    . (SubMode t s)
   => F.FoldM (Transactional g r t) o z
   -> Statement s i o
   -> i
   -> Transactional g r t z
foldM f st i = mk \env ->
   foldIO (F.hoists (flip un env) f) (pure env.tx) st i

-- | 'Ex.catch' behaves like "STM"'s 'catchSTM'.
--
-- In @'Ex.catch' ma f@, if an exception is thrown by @ma@, then any
-- database or 'Ref' changes made by @ma@ will be discarded. Furthermore, if
-- @f@ can handle said exception, then the action resulting from applying @f@
-- will be executed. Otherwise, if @f@ can't handle the exception, it will
-- bubble up.
--
-- Note: This instance's 'Cx.catch' catches async exceptions because that's
-- what 'Cx.MonadCatch' instances normaly do. As a user of this instance, you
-- probably want to use "Control.Exceptions.Safe" to make sure you don't catch
-- async exceptions unless you really want to.
instance Ex.MonadCatch (Transactional g r t) where
   catch act f = mk \env -> do
      refsRollback <- liftIO $ atomically $ saveSomeRefs env.refs
      case env.tx.smode of
         SRead ->
            Ex.catchAsync (un act env) \se -> do
               liftIO $ atomically refsRollback
               case Ex.fromException se of
                  Nothing -> Ex.throwM se
                  Just e -> un (f e) env
         SWrite -> Ex.mask \restore -> do
            sp <- savepoint env.tx
            Ex.tryAsync (restore (un act env)) >>= \case
               Right a -> do
                  -- savepointRelease is not critical.
                  void $ Ex.tryAny $ savepointRelease sp
                  pure a
               Left se -> do
                  liftIO $ atomically refsRollback
                  savepointRollback sp
                  -- savepointRelease is not critical. Just making sure we
                  -- don't accumulate many savepoints in case there is some
                  -- recursion going on.
                  void $ Ex.tryAny $ savepointRelease sp
                  case Ex.fromException se of
                     Nothing -> Ex.throwM se
                     Just e -> restore $ un (f e) env

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

-- | INTERNAL.
data ErrRetry = ErrRetry
   deriving stock (Show)
   deriving anyclass (Ex.Exception)

-- | 'retry' behaves like 'STM'\'s 'Control.Concurrent.STM.retry'. It causes
-- the current 'Transaction' to be cancelled so that a new one can take its
-- place and the entire 'Transactional' is executed again. This allows the
-- 'Transactional' to observe a new snapshot of the database.
--
-- * 'retry', 'empty' and 'mzero' all do fundamentally the same thing,
-- however 'retry' leads to better type inferrence because it forces the
-- @r@ type-parameter to be 'Retry'.
--
-- * __NOTICE__ You only need to use 'mzero' if you need access to a newer
-- database snapshot. If all you want to do is undo some 'Ref' transformation
-- effects, or undo database changes, then use 'catch' which doesn't abandon
-- the 'Transaction'.
--
-- * __WARNING__ If we keep 'retry'ing and the database never changes, then
-- we will be stuck in a loop forever. To mitigate this, when executing the
-- 'Transactional' through 'Sq.read', 'Sq.commit' or 'Sq.rollback', you may
-- want to use 'System.Timeout.timeout' to abort at some point in the future.
retry :: Transactional g 'Retry t a
retry = Ex.throwM ErrRetry
{-# INLINE retry #-}

-- | @'orElse' ma mb@ behaves like 'STM'\'s @'Control.Concurrent.STM.orElse' ma
-- mb@.  If @ma@ completes without executing 'retry', then that constitutes the
-- entirety of @'orElse' ma mb@. Otherwise, if @ma@ executed 'retry', then all
-- the effects from @ma@ are discared and @mb@ is tried in its place.
--
-- * 'orElse', '<|>' and 'mplus' all do the same thing, but 'orElse' has a more
-- general type because it doesn't force the @r@ type-parameter to be 'Retry'.
orElse
   :: Transactional g r t a
   -> Transactional g r t a
   -> Transactional g r t a
orElse tl tr = Ex.catch tl \ErrRetry -> tr

-- | @
-- 'empty' = 'retry'
-- '(<|>)' = 'orElse'
-- @
instance Alternative (Transactional g 'Retry t) where
   empty = retry
   {-# INLINE empty #-}
   (<|>) = orElse
   {-# INLINE (<|>) #-}

-- | @
-- 'mzero' = 'retry'
-- 'mplus' = 'orElse'
-- @
instance MonadPlus (Transactional g 'Retry t) where
   mzero = retry
   {-# INLINE mzero #-}
   mplus = orElse
   {-# INLINE mplus #-}

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

data SomeRef g where
   SomeRef :: Ref g a -> SomeRef g

-- | Creates a “savepoint” with the current state of the given 'SomeRef's.
-- The produced 'STM' action can be used to rollback the 'SomeRef's current
-- state in the future.
saveSomeRefs :: TVar (IntMap (SomeRef g)) -> STM (STM ())
saveSomeRefs tvsrs = do
   srs0 <- readTVar tvsrs
   rollbacks <- forM srs0 \(SomeRef (Ref tv)) ->
      writeTVar tv <$> readTVar tv
   pure do
      srs1 <- swapTVar tvsrs srs0
      forM_ (IntMap.difference srs1 srs0) \(SomeRef (Ref tv)) ->
         writeTVar tv Nothing
      sequence_ rollbacks

-- | Like 'TVar', but you can use it inside 'Transactional' through the
-- 'MonadRef' and 'MonadAtomicRef' vocabulary.
newtype Ref g a = Ref (TVar (Maybe a))
   deriving newtype
      ( Eq
        -- ^ Pointer equality
      )

-- | All operations are atomic.
instance MonadRef (Transactional g r t) where
   type Ref (Transactional g r t) = Sq.Transactional.Ref g
   newRef a = mk \env -> liftIO $ atomically do
      i <- env.unique
      tv <- newTVar $! Just a
      let ref = Ref tv
      -- Note: We only explicitly remove things from the IntMap through
      -- saveSomeRefs, or when exiting Transactional. Maybe some day we
      -- optimize this.
      modifyTVar' env.refs $ IntMap.insert i $! SomeRef ref
      pure ref
   readRef (Ref tv) = mk \_ -> liftIO $ atomically do
      readTVar tv >>= \case
         Just a -> pure a
         Nothing -> Ex.throwM $ resourceVanishedWithCallStack "Ref"
   writeRef r a = atomicModifyRef r \_ -> (a, ())
   modifyRef r f = atomicModifyRef r \a -> (f a, ())
   modifyRef' r f = atomicModifyRef' r \a -> (f a, ())

instance MonadAtomicRef (Transactional g r t) where
   atomicModifyRef (Ref tv) f =
      mk \_ -> liftIO $ atomically do
         readTVar tv >>= \case
            Just a0 | (a1, b) <- f a0 -> do
               writeTVar tv $! Just a1
               pure b
            Nothing -> Ex.throwM $ resourceVanishedWithCallStack "Ref"
   atomicModifyRef' (Ref tv) f =
      mk \_ -> liftIO $ atomically do
         readTVar tv >>= \case
            Just a0 | (!a1, !b) <- f a0 -> do
               writeTVar tv $! Just a1
               pure b
            Nothing -> Ex.throwM $ resourceVanishedWithCallStack "Ref"

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

-- | Executes a 'Statement' expected to return __zero or one__ rows.
--
-- * Throws 'ErrRows_TooMany' if more than one row.
maybe
   :: forall o i t s g r
    . (SubMode t s)
   => Statement s i o
   -> i
   -> Transactional g r t (Maybe o)
maybe = foldM $ foldMaybeM ErrRows_TooMany
{-# INLINE maybe #-}

-- | Executes a 'Statement' expected to return exactly __one__ row.
--
-- * Throws 'ErrRows_TooFew' if zero rows, 'ErrRows_TooMany' if more than one row.
one
   :: forall o i t s g r
    . (SubMode t s)
   => Statement s i o
   -> i
   -> Transactional g r t o
one = foldM $ foldOneM ErrRows_TooFew ErrRows_TooMany
{-# INLINE one #-}

-- | Executes a 'Statement' expected to return exactly __zero__ rows.
--
-- * Throws 'ErrRows_TooMany' if more than zero rows.
zero
   :: forall o i t s g r
    . (SubMode t s)
   => Statement s i o
   -> i
   -> Transactional g r t ()
zero = foldM $ foldZeroM ErrRows_TooMany
{-# INLINE zero #-}

-- | Executes a 'Statement' expected to return __one or more__ rows.
--
-- * Returns the length of the 'NonEmpty' list, too.
--
-- * Throws 'ErrRows_TooFew' if zero rows.
some
   :: forall o i t s g r
    . (SubMode t s)
   => Statement s i o
   -> i
   -> Transactional g r t (Int64, NonEmpty o)
some = foldM $ foldNonEmptyM ErrRows_TooFew
{-# INLINE some #-}

-- | Executes a 'Statement' expected to return __zero or more__ rows.
--
-- * Returns the length of the list, too.
list
   :: forall o i t s g r
    . (SubMode t s)
   => Statement s i o
   -> i
   -> Transactional g r t (Int64, [o])
list = fold foldList
{-# INLINE list #-}

-- | __Purely fold__ all the output rows.
fold
   :: forall o z i t s g r
    . (SubMode t s)
   => F.Fold o z
   -> Statement s i o
   -> i
   -> Transactional g r t z
fold = foldM . F.generalize
{-# INLINE fold #-}