Phases with any Ord key type

Phases.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DeriveFunctor #-}
module Phases where
import Control.Applicative ((<**>))
import Data.Foldable (sequenceA_)

type f ~> g = forall x. f x -> g x

data Phases key f a where
  Pure :: a -> Phases key f a
  Phase :: key -> f a -> Phases key f (a -> b) -> Phases key f b

deriving stock instance Functor f => Functor (Phases key f)

instance (Ord key, Applicative f) => Applicative (Phases key f) where
  pure = Pure
  liftA2 f (Pure x) (Pure y) = Pure (f x y)
  liftA2 f (Pure x) (Phase k fx f') = Phase k fx (fmap (f x .) f')
  liftA2 f (Phase k fx f') (Pure x) = Phase k fx (fmap (\g y -> f (g y) x) f')
  liftA2 f (Phase k fx f') (Phase k' fy f'') =
    case compare k k' of
      LT -> Phase k fx (fmap (\g b y -> f (g y) b) f' <*> Phase k' fy f'')
      GT -> Phase k' fy (fmap (\g a y -> f a (g y)) f'' <*> Phase k fx f')
      EQ -> Phase k (liftA2 (,) fx fy) (liftA2 (\l r (x, y) -> f (l x) (r y)) f' f'')


runPhases :: Applicative f => Phases key f a -> f a
runPhases (Pure a) = pure a
runPhases (Phase _ fx pf) = fx <**> runPhases pf

phase :: Applicative f => key -> f ~> Phases key f
phase k fa = Phase k fa (Pure id)

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

-- >> runPhases (one × two)
-- "initializing .."
-- "beep boop"
-- "extra work"
-- "handle"
-- ((),())
one :: Phases Phase IO ()
one = sequenceA_
  [ phase Setup   (print "initializing ..")
  , phase Run     (print "beep boop")
  , phase Cleanup (print "handle")
  ]

two :: Phases Phase IO ()
two = sequenceA_
  [ phase Run (print "extra work")
  ]

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