How to terminate a computation that runs in the `IO` monad? - haskell

Question

There is a library that provides a data type F and a function of type
ffoldlIO :: (b -> a -> IO b) -> b -> F a -> IO b
The function is similar to
foldlIO :: (b -> a -> IO b) -> b -> [a] -> IO b
foldlIO f a = \xs -> foldr (\x r (!a') -> f a' x >>= r) return xs a
I wonder whether foldlIO (and thus ffoldlIO) can run in a short-circuit fashion.
Consider this example:
example1 :: IO Int
example1 = foldlIO (\a x -> if a < 4 then return (a + x) else return a) 0 [1..5]
Here foldlIO traverses the entire list, but what if we throw an exception to stop the computation and then catch it? Something like this:
data Terminate = Terminate
deriving (Show)
instance Exception Terminate
example2 :: IO Int
example2 = do
ra <- newIORef 0
let step a x
| a' < 4 = return a'
| otherwise = writeIORef ra a' >> throwIO Terminate
where a' = a + x
foldlIO step 0 [1..] `catch` \(_ :: Terminate) -> readIORef ra
Is this reliable? Is there a better way to terminate a computation that runs in the IO monad (and no other monad) or am I not supposed to do this at all?

Answer 1

For example, you can use ContT monad transformer like this:
example3 :: IO Int
example3 = flip runContT return . callCC $ \exit -> do
let step a x
| a' < 4 = return a'
| otherwise = exit a'
where a' = a + x
foldM step 0 [1..]
Also, you can define you own version of foldM with posibility of termination.
termFoldM :: (Monad m, Foldable t) =>
((b -> ContT b m c) -> b -> a -> ContT b m b) -> b -> t a -> m b
termFoldM f a t = flip runContT return . callCC $ \exit -> foldM (f exit) a xs
example4 :: IO Int
example4 = termFoldM step 0 [1..]
where
step exit a x
| a' < 4 = return a'
| otherwise = exit a'
where a' = a + x
But this way (with ContT) has one problem. You can't easy do some IO actions. For example, this code will not be compiled, because step function must return value of type ContT Int IO Int not IO Int.
let step a x
| a' < 4 = putStrLn ("'a = " ++ show a') >> return a'
| otherwise = exit a'
where a' = a + x
Fortunately, you can solve this by the lift function, like this:
let step a x
| a' < 4 = lift (putStrLn ("'a = " ++ show a')) >> return a'
| otherwise = exit a'
where a' = a + x

Answer 2

My first answer was not correct. So, I'll try to improve.
I think that the use of exceptions to terminate in IO monad is not a hack but it does not look clean. I propose to define the instance MonadCont IO like this:
data Terminate = forall a . Terminate a deriving (Typeable)
instance Show Terminate where show = const "Terminate"
instance Exception Terminate
instance MonadCont IO where
callCC f = f exit `catch` (\(Terminate x) -> return . unsafeCoerce $ x)
where exit = throwIO . Terminate
Then you can rewrite your example more cleaner.
example :: IO Int
example = callCC $ \exit -> do
let step a x
| a' < 4 = return a'
| otherwise = exit a'
where a' = a + x
foldlIO step 0 [1..]
Variant with IOREf.
data Terminate = Terminate deriving (Show, Typeable)
instance Exception Terminate
instance MonadCont IO where
callCC f = do
ref <- newIORef undefined
let exit a = writeIORef ref a >> throwIO Terminate
f exit `catch` (\Terminate -> readIORef ref)