The Phases Applicative (but written as a heap)

Phases.hs

{-# LANGUAGE GHC2021, LambdaCase, RoleAnnotations #-}

module Phases (Phases, runPhases, phase) where

-- To maintain the heap invariant according to @Ord k@, the type role of @k@
-- must be nominal, and @Phases@ must be opaque.
type role Phases nominal representational representational
data Phases k f a where
  Pure  :: a -> Phases k f a
  Phase :: k -> (w -> x -> y -> z) -> f w -> Phases k f x -> Phases k f y
        -> Phases k f z
deriving instance Functor f => Functor (Phases k f)

instance (Ord k, Applicative f) => Applicative (Phases k f) where
  pure = Pure
  liftA2 f (Pure a) pb =      f a <$> pb
  liftA2 f pa (Pure b) = flip f b <$> pa
  liftA2 f pa@(Phase ka ga fwa pxa pya) pb@(Phase kb gb fwb pxb pyb)
    | ka <= kb  = Phase ka leq fwa       pxa (pya × pb)
    | otherwise = Phase kb gt  fwb (pa × pxb) pyb
   where
    leq w x (y, b) = f (ga w x y) b
    gt  w (a, x) y = f a (gb w x y)

runPhases :: (Ord k, Applicative f) => Phases k f a -> f a
runPhases = \case
  Pure x             -> pure x
  Phase _ g fw px py -> uncurry . g <$> fw <*> runPhases (px × py)

phase :: k -> f a -> Phases k f a
phase k fx = Phase k (const . const) fx (Pure ()) (Pure ())

(×) :: Applicative f => f a -> f b -> f (a, b)
(×) = liftA2 (,)


data Phase = Setup | Run | Cleanup
  deriving (Eq, Ord)

-- >> runPhases (_one × _two)
-- "initializing .."
-- "beep boop"
-- "extra work"
-- "handle"
-- ((),())
_one :: Phases Phase IO ()
_one = phase Setup   (print "initializing ..")
    *> phase Run     (print "beep boop")
    *> phase Cleanup (print "handle")

_two :: Phases Phase IO ()
_two = phase Run (print "extra work")