Add the ReturnWith effect for early return from a computation

effectful-core/CHANGELOG.md

@@ -7,6 +7,9 @@
   `Effectful.Output.Static.Local.List`, `Effectful.Output.Static.Shared.Array`,
   `Effectful.Output.Static.Shared.List` and `Effectful.Labeled.Output`) for
   accumulation of values.
+* Add the `ReturnWith` effect (`Effectful.ReturnWith.Dynamic`,
+  `Effectful.ReturnWith.Static` and `Effectful.Labeled.ReturnWith`) for early
+  return from a computation.
 * Make the `Provider` and `ProviderList` effects dynamically dispatched and
   export their operations.
 * Add `Effectful.Labeled.Provider` and `Effectful.Labeled.Provider.List` with

effectful-core/effectful-core.cabal

@@ -92,6 +92,7 @@ library
                      Effectful.Labeled.Provider
                      Effectful.Labeled.Provider.List
                      Effectful.Labeled.Reader
+                     Effectful.Labeled.ReturnWith
                      Effectful.Labeled.State
                      Effectful.Labeled.Writer
                      Effectful.NonDet
@@ -106,6 +107,8 @@ library
                      Effectful.Provider.List
                      Effectful.Reader.Dynamic
                      Effectful.Reader.Static
+                     Effectful.ReturnWith.Dynamic
+                     Effectful.ReturnWith.Static
                      Effectful.State.Dynamic
                      Effectful.State.Static.Local
                      Effectful.State.Static.Shared

effectful-core/src/Effectful/Labeled/ReturnWith.hs

@@ -0,0 +1,43 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+-- | Convenience functions for the 'Labeled' 'ReturnWith' effect.
+--
+-- @since 2.7.0.0
+module Effectful.Labeled.ReturnWith
+  ( -- * Effect
+    ReturnWith(..)
+
+    -- ** Handlers
+  , runReturnWith
+
+    -- ** Operations
+  , returnWith
+
+    -- * Re-exports
+  , Labeled(..)
+  ) where
+
+import GHC.Stack (withFrozenCallStack)
+
+import Effectful
+import Effectful.Dispatch.Dynamic
+import Effectful.Labeled
+import Effectful.ReturnWith.Dynamic (ReturnWith(..))
+import Effectful.ReturnWith.Dynamic qualified as R
+
+-- | Run a computation that can return early with a value of type @r@ (via
+-- "Effectful.ReturnWith.Static").
+runReturnWith
+  :: forall label r es
+   . HasCallStack
+  => Eff (Labeled label (ReturnWith r) : es) r
+  -> Eff es r
+runReturnWith = runLabeled @label R.runReturnWith
+
+-- | Return early with the given value.
+returnWith
+  :: forall label r es a
+   . (HasCallStack, Labeled label (ReturnWith r) :> es)
+  => r
+  -- ^ The value.
+  -> Eff es a
+returnWith = withFrozenCallStack send . Labeled @label . ReturnWith

effectful-core/src/Effectful/ReturnWith/Dynamic.hs

@@ -0,0 +1,46 @@
+-- | The dynamically dispatched variant of the 'ReturnWith' effect.
+--
+-- /Note:/ unless you plan to change interpretations at runtime, it's
+-- recommended to use the statically dispatched variant,
+-- i.e. "Effectful.ReturnWith.Static".
+--
+-- @since 2.7.0.0
+module Effectful.ReturnWith.Dynamic
+  ( -- * Effect
+    ReturnWith(..)
+
+    -- ** Handlers
+  , runReturnWith
+
+    -- ** Operations
+  , returnWith
+  ) where
+
+import GHC.Stack (withFrozenCallStack)
+
+import Effectful
+import Effectful.Dispatch.Dynamic
+import Effectful.ReturnWith.Static qualified as R
+
+-- | Provide the ability to return early with a value of type @r@.
+data ReturnWith r :: Effect where
+  ReturnWith :: r -> ReturnWith r m a
+
+type instance DispatchOf (ReturnWith r) = Dynamic
+
+-- | Run a computation that can return early with a value of type @r@ (via
+-- "Effectful.ReturnWith.Static").
+runReturnWith
+  :: HasCallStack
+  => Eff (ReturnWith r : es) r
+  -> Eff es r
+runReturnWith = reinterpret_ R.runReturnWith $ \case
+  ReturnWith r -> R.returnWith r
+
+-- | Return early with the given value.
+returnWith
+  :: (HasCallStack, ReturnWith r :> es)
+  => r
+  -- ^ The value.
+  -> Eff es a
+returnWith = withFrozenCallStack send . ReturnWith

effectful-core/src/Effectful/ReturnWith/Static.hs

@@ -0,0 +1,92 @@
+-- | Support for early return from a computation.
+--
+-- >>> import Control.Monad (when)
+--
+-- >>> :{
+--   classify :: ReturnWith String :> es => Int -> Eff es String
+--   classify n = do
+--     when (n < 0) $ returnWith "negative"
+--     when (n == 0) $ returnWith "zero"
+--     pure "positive"
+-- :}
+--
+-- >>> runEff . runReturnWith $ classify 5
+-- "positive"
+--
+-- >>> runEff . runReturnWith $ classify (-5)
+-- "negative"
+--
+-- @since 2.7.0.0
+module Effectful.ReturnWith.Static
+  ( -- * Effect
+    ReturnWith
+
+    -- ** Handlers
+  , runReturnWith
+
+    -- ** Operations
+  , returnWith
+  ) where
+
+import Data.Kind
+import GHC.Stack
+
+import Effectful
+import Effectful.Dispatch.Static
+import Effectful.Exception
+import Effectful.Internal.Utils
+
+-- | Provide the ability to return early with a value of type @r@.
+data ReturnWith (r :: Type) :: Effect
+
+type instance DispatchOf (ReturnWith r) = Static NoSideEffects
+newtype instance StaticRep (ReturnWith r) = ReturnWith ReturnWithId
+
+-- | Run a computation that can return early with a value of type @r@.
+runReturnWith
+  :: forall r es
+   . HasCallStack
+  => Eff (ReturnWith r : es) r
+  -> Eff es r
+runReturnWith action = do
+  rid <- unsafeEff_ newReturnWithId
+  evalStaticRep (ReturnWith @r rid) $ do
+    catchJust (matchReturnWith rid) action pure
+
+-- | Return early with the given value.
+returnWith
+  :: forall r es a. (HasCallStack, ReturnWith r :> es)
+  => r
+  -- ^ The value.
+  -> Eff es a
+returnWith r = do
+  ReturnWith rid <- getStaticRep @(ReturnWith r)
+  withFrozenCallStack throwIO $ ReturnWithWrapper rid callStack (toAny r)
+
+----------------------------------------
+-- Helpers
+
+newtype ReturnWithId = ReturnWithId Unique
+  deriving newtype Eq
+
+-- | A unique is picked so that distinct 'ReturnWith' handlers for the same
+-- type don't catch each other's values.
+newReturnWithId :: IO ReturnWithId
+newReturnWithId = ReturnWithId <$> newUnique
+
+data ReturnWithWrapper = ReturnWithWrapper !ReturnWithId CallStack Any
+
+instance Show ReturnWithWrapper where
+  showsPrec _ (ReturnWithWrapper _ cs _)
+    = ("Effectful.ReturnWith.Static.ReturnWithWrapper\n" ++)
+    . (prettyCallStack cs ++)
+
+instance Exception ReturnWithWrapper where
+  -- See discussion in https://github.com/haskell-effectful/effectful/pull/232.
+  toException = asyncExceptionToException
+  fromException = asyncExceptionFromException
+
+matchReturnWith :: ReturnWithId -> ReturnWithWrapper -> Maybe r
+matchReturnWith rid (ReturnWithWrapper rtag _ r)
+  | rid == rtag = Just (fromAny r)
+  | otherwise = Nothing

effectful/CHANGELOG.md

@@ -14,6 +14,9 @@
   `Effectful.Output.Static.Local.List`, `Effectful.Output.Static.Shared.Array`,
   `Effectful.Output.Static.Shared.List` and `Effectful.Labeled.Output`) for
   accumulation of values.
+* Add the `ReturnWith` effect (`Effectful.ReturnWith.Dynamic`,
+  `Effectful.ReturnWith.Static` and `Effectful.Labeled.ReturnWith`) for early
+  return from a computation.
 * Re-export `Labeled(..)` from all `Effectful.Labeled.*` modules.
 * Drop support for GHC < 9.6.
 * Add definitions of `rethrowM` to `MonadThrow` and `catchNoPropagate` to

effectful/effectful.cabal

@@ -117,6 +117,7 @@ library
                          , Effectful.Labeled.Input
                          , Effectful.Labeled.Output
                          , Effectful.Labeled.Reader
+                         , Effectful.Labeled.ReturnWith
                          , Effectful.Labeled.State
                          , Effectful.Labeled.Writer
                          , Effectful.NonDet
@@ -131,6 +132,8 @@ library
                          , Effectful.Provider.List
                          , Effectful.Reader.Dynamic
                          , Effectful.Reader.Static
+                         , Effectful.ReturnWith.Dynamic
+                         , Effectful.ReturnWith.Static
                          , Effectful.State.Dynamic
                          , Effectful.State.Static.Local
                          , Effectful.State.Static.Shared
@@ -180,6 +183,7 @@ test-suite test
                     OutputTests
                     PrimTests
                     ReaderTests
+                    ReturnWithTests
                     StateTests
                     TimeoutTests
                     UnliftTests

effectful/tests/Main.hs

@@ -15,6 +15,7 @@ import NonDetTests
 import OutputTests
 import PrimTests
 import ReaderTests
+import ReturnWithTests
 import StateTests
 import TimeoutTests
 import UnliftTests
@@ -35,6 +36,7 @@ main = defaultMain $ testGroup "effectful"
   , outputTests
   , primTests
   , readerTests
+  , returnWithTests
   , stateTests
   , timeoutTests
   , unliftTests

effectful/tests/ReturnWithTests.hs

@@ -0,0 +1,88 @@
+module ReturnWithTests (returnWithTests) where
+
+import Test.Tasty
+import Test.Tasty.HUnit
+
+import Effectful
+import Effectful.Dispatch.Dynamic
+import Effectful.Exception (finally)
+import Effectful.Labeled.ReturnWith qualified as LR
+import Effectful.ReturnWith.Dynamic qualified as RD
+import Effectful.ReturnWith.Static qualified as RS
+import Effectful.State.Static.Local
+import Utils qualified as U
+
+returnWithTests :: TestTree
+returnWithTests = testGroup "ReturnWith"
+  [ testCase "static short-circuits" test_static
+  , testCase "static falls through" test_staticFallThrough
+  , testCase "different handlers are independent" test_independentHandlers
+  , testCase "cleanup actions run on early return" test_cleanup
+  , testCase "dynamic short-circuits" test_dynamic
+  , testCase "labeled handlers are targeted correctly" test_labeled
+  ]
+
+test_static :: Assertion
+test_static = runEff . evalState @Int 0 $ do
+  r <- RS.runReturnWith @String $ do
+    modify @Int (+1)
+    _ <- RS.returnWith "early"
+    modify @Int (+1)
+    pure "late"
+  U.assertEqual "result" "early" r
+  U.assertEqual "state changes before returnWith persist" 1 =<< get @Int
+
+test_staticFallThrough :: Assertion
+test_staticFallThrough = runEff $ do
+  r <- RS.runReturnWith @String $ pure "done"
+  U.assertEqual "result" "done" r
+
+test_independentHandlers :: Assertion
+test_independentHandlers = runEff $ do
+  r <- RS.runReturnWith @String . runOuterReturn $ do
+    inner <- RS.runReturnWith @String $ do
+      outerReturn
+      pure "inner"
+    pure $ "inner handler caught " ++ inner
+  U.assertEqual "correct value returned" "outer" r
+
+test_cleanup :: Assertion
+test_cleanup = runEff . evalState @Int 0 $ do
+  r <- RS.runReturnWith @String $ do
+    (RS.returnWith "early" >> pure "late") `finally` modify @Int (+1)
+  U.assertEqual "result" "early" r
+  U.assertEqual "cleanup ran" 1 =<< get @Int
+
+test_dynamic :: Assertion
+test_dynamic = runEff . evalState @Int 0 $ do
+  r <- RD.runReturnWith @String $ do
+    modify @Int (+1)
+    _ <- RD.returnWith "early"
+    modify @Int (+1)
+    pure "late"
+  U.assertEqual "result" "early" r
+  U.assertEqual "state changes before returnWith persist" 1 =<< get @Int
+
+test_labeled :: Assertion
+test_labeled = runEff $ do
+  r <- LR.runReturnWith @"outer" @String $ do
+    n <- LR.runReturnWith @"inner" @Int $ do
+      _ <- LR.returnWith @"outer" "outer wins"
+      pure 0
+    pure $ "inner returned " ++ show n
+  U.assertEqual "value caught by the outer handler" "outer wins" r
+
+----------------------------------------
+-- Helpers
+
+data OuterReturn :: Effect where
+  OuterReturn :: OuterReturn m ()
+
+type instance DispatchOf OuterReturn = Dynamic
+
+outerReturn :: OuterReturn :> es => Eff es ()
+outerReturn = send OuterReturn
+
+runOuterReturn :: RS.ReturnWith String :> es => Eff (OuterReturn : es) a -> Eff es a
+runOuterReturn = interpret_ $ \case
+  OuterReturn -> RS.returnWith "outer"