module Mit.Monad
  ( Mit,
    runMit,
    io,
    getEnv,
    withEnv,
    Goto,
    Label,
    label,
    with,
    with_,
    X,
  )
where

import Control.Monad qualified
import Mit.Prelude

newtype Mit r x a
  = Mit (r -> (a -> IO x) -> IO x)
  deriving stock (forall a b. a -> Mit r x b -> Mit r x a
forall a b. (a -> b) -> Mit r x a -> Mit r x b
forall r x a b. a -> Mit r x b -> Mit r x a
forall r x a b. (a -> b) -> Mit r x a -> Mit r x b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Mit r x b -> Mit r x a
$c<$ :: forall r x a b. a -> Mit r x b -> Mit r x a
fmap :: forall a b. (a -> b) -> Mit r x a -> Mit r x b
$cfmap :: forall r x a b. (a -> b) -> Mit r x a -> Mit r x b
Functor)

instance Applicative (Mit r x) where
  pure :: forall a. a -> Mit r x a
pure a
x = forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
_ a -> IO x
k -> a -> IO x
k a
x
  <*> :: forall a b. Mit r x (a -> b) -> Mit r x a -> Mit r x b
(<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad (Mit r x) where
  return :: forall a. a -> Mit r x a
return = forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Mit r -> (a -> IO x) -> IO x
mx >>= :: forall a b. Mit r x a -> (a -> Mit r x b) -> Mit r x b
>>= a -> Mit r x b
f =
    forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
r b -> IO x
k ->
      r -> (a -> IO x) -> IO x
mx r
r (\a
a -> forall r x a. Mit r x a -> r -> (a -> IO x) -> IO x
unMit (a -> Mit r x b
f a
a) r
r b -> IO x
k)

instance MonadIO (Mit r x) where
  liftIO :: forall a. IO a -> Mit r x a
liftIO = forall a r x. IO a -> Mit r x a
io

unMit :: Mit r x a -> r -> (a -> IO x) -> IO x
unMit :: forall r x a. Mit r x a -> r -> (a -> IO x) -> IO x
unMit (Mit r -> (a -> IO x) -> IO x
k) =
  r -> (a -> IO x) -> IO x
k

runMit :: r -> Mit r a a -> IO a
runMit :: forall r a. r -> Mit r a a -> IO a
runMit r
r Mit r a a
m =
  forall r x a. Mit r x a -> r -> (a -> IO x) -> IO x
unMit Mit r a a
m r
r forall (f :: * -> *) a. Applicative f => a -> f a
pure

return :: x -> Mit r x a
return :: forall x r a. x -> Mit r x a
return x
x =
  forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
_ a -> IO x
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure x
x

io :: IO a -> Mit r x a
io :: forall a r x. IO a -> Mit r x a
io IO a
m =
  forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
_ a -> IO x
k -> do
    a
x <- IO a
m
    a -> IO x
k a
x

getEnv :: Mit r x r
getEnv :: forall r x. Mit r x r
getEnv =
  forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
r r -> IO x
k -> r -> IO x
k r
r

withEnv :: (r -> s) -> Mit s x a -> Mit r x a
withEnv :: forall r s x a. (r -> s) -> Mit s x a -> Mit r x a
withEnv r -> s
f Mit s x a
m =
  forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
r a -> IO x
k -> forall r x a. Mit r x a -> r -> (a -> IO x) -> IO x
unMit Mit s x a
m (r -> s
f r
r) a -> IO x
k

type Goto r x a =
  forall xx void. Label (X x a) xx => a -> Mit r xx void

label :: forall r x a. (Goto r x a -> Mit r (X x a) a) -> Mit r x a
label :: forall r x a. (Goto r x a -> Mit r (X x a) a) -> Mit r x a
label Goto r x a -> Mit r (X x a) a
f =
  forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
r a -> IO x
k -> do
    forall a b. (a -> IO b) -> IO (X b a) -> IO b
unX a -> IO x
k (forall r x a. Mit r x a -> r -> (a -> IO x) -> IO x
unMit (Goto r x a -> Mit r (X x a) a
f \a
a -> forall x r a. x -> Mit r x a
return (forall a b. Label a b => a -> b
bury @(X x a) (forall a b. b -> X a b
XR a
a))) r
r (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. b -> X a b
XR))

with :: (forall v. (a -> IO v) -> IO v) -> (a -> Mit r (X x b) b) -> Mit r x b
with :: forall a r x b.
(forall v. (a -> IO v) -> IO v)
-> (a -> Mit r (X x b) b) -> Mit r x b
with forall v. (a -> IO v) -> IO v
f a -> Mit r (X x b) b
action =
  forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
r b -> IO x
k ->
    forall a b. (a -> IO b) -> IO (X b a) -> IO b
unX b -> IO x
k (forall v. (a -> IO v) -> IO v
f \a
a -> forall r x a. Mit r x a -> r -> (a -> IO x) -> IO x
unMit (a -> Mit r (X x b) b
action a
a) r
r (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. b -> X a b
XR))

with_ :: (forall v. IO v -> IO v) -> Mit r (X x a) a -> Mit r x a
with_ :: forall r x a.
(forall v. IO v -> IO v) -> Mit r (X x a) a -> Mit r x a
with_ forall v. IO v -> IO v
f Mit r (X x a) a
action =
  forall r x a. (r -> (a -> IO x) -> IO x) -> Mit r x a
Mit \r
r a -> IO x
k ->
    forall a b. (a -> IO b) -> IO (X b a) -> IO b
unX a -> IO x
k (forall v. IO v -> IO v
f (forall r x a. Mit r x a -> r -> (a -> IO x) -> IO x
unMit Mit r (X x a) a
action r
r (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. b -> X a b
XR)))

-- instance Label (X a b) (X a b)
-- instance Label (X a b) (X (X a b) c)
-- instance Label (X a b) (X (X (X a b) c) d)
-- etc...

data X a b
  = XL a
  | XR b

unX :: (a -> IO b) -> IO (X b a) -> IO b
unX :: forall a b. (a -> IO b) -> IO (X b a) -> IO b
unX a -> IO b
k IO (X b a)
mx =
  IO (X b a)
mx forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    XL b
b -> forall (f :: * -> *) a. Applicative f => a -> f a
pure b
b
    XR a
a -> a -> IO b
k a
a

class Label a b where
  bury :: a -> b
  default bury :: a ~ b => a -> b
  bury = forall a. a -> a
id

-- FIXME I don't really think this is correct...
instance {-# INCOHERENT #-} Label (X a b) (X a b)

instance Label (X a b) (X c d) => Label (X a b) (X (X c d) e) where
  bury :: X a b -> X (X c d) e
bury = forall a b. a -> X a b
XL forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Label a b => a -> b
bury