{-
A crude implementation of an example for unrestricted while loops in
"Unguarded Recursion on Coinductive Resumptions" (https://lmcs.episciences.org/4784)
-}
{-# LANGUAGE TypeOperators, MultiParamTypeClasses #-}

import Control.Arrow (Kleisli(..), (&&&))
import Control.Monad ((>=>))

type a :+: b = Either a b
type a :*: b = (a,b)

-- Monad strength
tau :: Monad m => (a :*: m b) -> m (a :*: b)
tau (a, mb) = do
  b <- mb
  pure (a,b)

class (Monad m) => Action m a n where
  act :: a -> Kleisli m n n

class (Monad m) => ElgotMonad m where
  iter :: (x -> m (y :+: x)) -> x -> m y
  iter f = f >=> either pure (iter f) 

data While a =
  Act a
  | IfThenElse Bool (While a) (While a)
  | Seq (While a) (While a)
  | WhileDo Bool (While a)
  deriving (Show)

-- :*: distributes over :+:
dist :: n :*: (() :+: ()) -> (n :*: ()) :+: (n :*: ())
dist (n, c) = either (\_ -> Left (n, ())) (\_ -> Right (n, ())) c
  
denb :: (Monad m) => Bool -> Kleisli m n (() :+: ())
denb True = Kleisli (\_ -> pure $ Right ())
denb False = Kleisli (\_ -> pure $ Left ())

den :: (ElgotMonad m, Action m a n) => While a -> Kleisli m n n
den (Act a) = act a
den (Seq p q) = Kleisli $ runKleisli (den p) >=> runKleisli (den q)
den (IfThenElse b p q) = Kleisli $
  (fmap dist . tau . (id &&& runKleisli (denb b)))
  >=> either (runKleisli (den p) . fst) (runKleisli (den q) . fst)
den (WhileDo b p) = Kleisli $ iter $
  (fmap dist . tau . (id &&& runKleisli (denb b)))
  >=> either (pure . Left . fst) (fmap Right . runKleisli (den p) . fst)

-- example
data A = P

instance Action IO A () where
  act P = Kleisli $ \_ -> putStrLn "P"

instance ElgotMonad IO

printer :: While A
printer = WhileDo True (Act P)

run :: IO ()
run = runKleisli (den printer) ()