Stateful loop with different types of breaks

Question

I am trying to convert the following stateful imperative code into Haskell.

while (true) {
  while (get()) {
    if (put1()) {
      failImmediately();
    }
  }
  if (put2()) {
    succeedImmediately();
  }
}

Both the put1 and put2 read a state of the system and modify it. get can for simplicity just read the state. failImmediately should break out of the endless-loop and present one type of result, succeedImmediately should also break out but present a different result.

What I tried to use was State Env Result where Env represented the state of environment and Result was something like Either Failure Success for some custom Failure and Success.

I struggle with the requirement that the whole resulting expression should collapse into the Failure/Success once one of them is produced (breaking the loop) and otherwise keep going.

One idea I had was use Either Exit () where data Exit = Success | Failure and use StateT somehow to behave upon Left of the Either as if Either was the monad being chained, i.e. ignoring any subsequent actions.

I would really appreciate any inspiration or sample of haskell code that would achieve the same behaviour as the snippet above.

Edit: refined version moved to a separate question "Stateful computation with different types of short-circuit (Maybe, Either)".

Answer 1

Using the kit from @chi's answer, just highlighting that you don't need the full power of ContT, the direct-short-circuiting semantics of EitherT is enough:

import Control.Monad.Trans.Either

data Result a = Failure | Success a

foo :: EitherT (Result Int) IO Int
foo = forever $ do
    whileM get $ do
        whenM put1 $ do
            left Failure
    whenM put2 $ do
        left $ Success 42

run :: (Monad m) => EitherT (Result a) m a -> m (Maybe a)
run act = do
    res <- runEitherT act
    return $ case res of
        Left Failure -> Nothing
        Left (Success x) -> Just x
        Right x -> Just x

-- whenM / whileM and get/put1/put2 as per @chi's answeer

Answer 2

An almost literal, non elegant but effective translation.

We exploit the ContT monad transformer to achieve the effect of "early return". I.e., we want to be able to break our loops at any point. This is achieved by using callCC $ \exit -> ... which roughly makes exit our magic function which let us escape from the inner blocks immediately.

import Control.Monad.Cont

action :: IO String
action = flip runContT return $ callCC $ \exit -> 
   forever $ do                   -- while (true)
      let loop = do
             r1 <- lift $ get     -- if (get())
             when r1 $ do
                r2 <- lift $ put1
                when r2 $         -- if (put1())
                   exit "failImmediately"
                loop              -- "repeat while"
      loop
      r3 <- lift $ put2
      when r3 $
         exit "succeedImmediately"

get :: IO Bool
get = readLn

put1 :: IO Bool
put1 = putStrLn "put1 here" >> readLn

put2 :: IO Bool
put2 = putStrLn "put2 here" >> readLn

main :: IO ()
main = action >>= putStrLn

We can also define some custom helpers to prettify the code:

action2 :: IO String
action2 = flip runContT return $ callCC $ \exit -> 
   forever $ do                -- while (true)
      whileM get $             -- while(get())
         whenM put1 $          -- if (put1())
            exit "failImmediately"
      whenM put2 $             -- if (put2())
         exit "succeedImmediately"

whenM :: (MonadTrans t, Monad m, Monad (t m)) => m Bool -> t m () -> t m ()
whenM condition a = do
   r <- lift condition
   when r a

whileM :: (MonadTrans t, Monad m, Monad (t m)) => m Bool -> t m () -> t m ()
whileM condition a = whenM condition (a >> whileM condition a)