ZIO in Haskell

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UnboxedTuples #-}

import Control.Monad.Reader.Class (MonadReader (..))
import Data.Void (Void)
import GHC.IO (IO (..))
import GHC.Prim (RealWorld, State#)

newtype ZIO s r e a =
  ZIO ((# State# s, r #) -> (# State# s, a #))

instance Functor (ZIO s r e) where
  fmap f (ZIO z) =
    ZIO $ \(# s, r #) ->
      let (# s', a #) = z (# s, r #) in
      (# s', f a #)

instance Applicative (ZIO s r e) where
  pure x = ZIO $ \(# s, _ #) -> (# s, x #)
  ZIO z <*> ZIO aa =
    ZIO $ \(# s, r #) ->
      let (# s', z' #) = z (# s, r #) in
      let (# s'', aa' #) = aa (# s', r #) in
      (# s'', z' aa' #)

instance Monad (ZIO s r e) where
  ZIO z >>= k =
    ZIO $ \(# s, r #) ->
      let (# s', a #) = z (# s, r #) in
      let ZIO aa = k a in
      aa (# s', r #)

instance MonadReader r (ZIO s r e) where
  ask = ZIO $ \(# s, r #) -> (# s, r #)
  local f (ZIO z) =
    ZIO $ \(# s, r #) ->
      z (# s, f r #)

runZIO :: ZIO RealWorld r Void a -> r -> IO a
runZIO (ZIO z) r =
  IO $ \s -> z (# s, r #)

main' :: ZIO s r e ()
main' = pure ()

main :: IO ()
main = runZIO main' ()