comonad-coactions-0.1.0.0: src/Control/Comonad/Coaction/TH.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# LANGUAGE TypeData #-}

module Control.Comonad.Coaction.TH (mkLowerBy) where

import Control.Comonad
import Control.Comonad.Trans.Class (ComonadTrans (..))
import Data.Kind qualified as K
import Language.Haskell.TH

infixl 5 #

(#) :: Type -> Type -> Type
(#) = AppT

(|->|) :: Type -> Type -> Type
a |->| b = ArrowT # a # b

mkLowerBy :: Q [Dec]
mkLowerBy =
  reify ''ComonadTrans
    >>= \case
      ClassI _ instances ->
        do
          decs <-
            [d|
              type data Nat = Z | S Nat

              class (Comonad w, Comonad q) => LowerBy (k :: Nat) (w :: K.Type -> K.Type) (q :: K.Type -> K.Type) | k q -> w where
                lowerBy :: q a -> w a

              instance (Comonad w) => LowerBy Z w w where
                lowerBy = id
              |]
          let famName = mkName "Steps"
          w <- newName "w"
          q <- newName "q"
          k <- newName "k"
          let famDec =
                ClosedTypeFamilyD
                  ( TypeFamilyHead
                      famName
                      [ KindedTV w BndrReq (StarT |->| StarT),
                        KindedTV q BndrReq (StarT |->| StarT)
                      ]
                      (KindSig . ConT $ mkName "Nat")
                      Nothing
                  )
                  $ TySynEqn Nothing (ConT famName # VarT w # VarT w) (ConT $ mkName "Z")
                    : ( instances >>= \case
                          InstanceD _ _ (AppT (ConT _) t) _ ->
                            [ TySynEqn
                                Nothing
                                (ConT famName # VarT w # (t # VarT q))
                                (ConT (mkName "S") # (ConT famName # VarT w # VarT q))
                            ]
                          _ -> []
                      )
          let inductiveInstances =
                instances >>= \case
                  InstanceD ov ct (AppT (ConT _) t) _ ->
                    pure $
                      InstanceD
                        ov
                        (ct ++ [ConT (mkName "LowerBy") # VarT k # VarT w # VarT q, ConT ''Comonad # (t # VarT q)])
                        (ConT (mkName "LowerBy") # (ConT (mkName "S") # VarT k) # VarT w # (t # VarT q))
                        [ ValD
                            (VarP $ mkName "lowerBy")
                            (NormalB $ UInfixE (AppTypeE (VarE $ mkName "lowerBy") (VarT k)) (VarE '(.)) (VarE 'lower))
                            []
                        ]
                  _ -> []
          pure $ decs ++ famDec : inductiveInstances
      _ -> pure []