Rewriting type families with GHC type-checking plugins

Introduction

Talk overview

When to reach for a type-checking plugin?
Lightning fast review of type-checker plugins.
New functionality: rewriting type family applications.
Illustration: intrinsically typed System F.
Presentation of a library for writing type-checking plugins.

When to reach for a type-checking plugin?

Equational theory: Peano arithmetic

GHC doesn’t understand basic arithmetic.


type (+) :: Nat -> Nat -> Nat
type family a + b where
  Zero   + b = b
  Succ a + b = Succ (a + b)

GHC can only reduce a + b when it knows what a is:
is it Zero or is it Succ i for some i?

Can’t reason about (+) parametrically:

Could not deduce n + Succ n ~ Succ (n + n)

Equational theory: row types

Row: unordered association map, field name ⇝ type, e.g.
( "field1" :: Int, "field2" :: Bool ).
Can be used to model extensible records (order doesn’t matter).

To account for re-ordering, we need:

Insert k v (Insert l w r) ~ Insert l w (Insert k v r)

when k and l are distinct field names that don’t appear in the row r.

Problem: in GHC, we can’t have type-families on the LHS of a type family equation:

• Illegal type synonym family application ‘Insert l w r’ in instance:
    Insert k v (Insert l w r)

Non-confluence

GHC can’t allow type-families on the LHS of a type family equation in general: no guarantee type-family reduction is confluent.


type F :: Type -> Type
type family F a where
  F (F a) = a       -- F[0]
  F Bool  = Float   -- F[1]
  F a     = Maybe a -- F[2]


F (F Bool) ~F[0]~> Bool
F (F Bool) ~(F F[1])~> F Float ~F[2]~> Maybe Float

Type-checking plugins to the rescue

Rewriter plugins


data TcPlugin = forall s. TcPlugin
  { tcPluginInit    :: TcPluginM s
  , tcPluginSolve   :: s -> TcPluginSolver
  , tcPluginRewrite :: s -> UniqFM TyCon TcPluginRewriter -- new!
  , tcPluginStop    :: s -> TcPluginM ()
  }


type TcPluginRewriter
  =  [Ct]   -- Givens
  -> [Type] -- Type family arguments
  -> TcPluginM TcPluginRewriteResult

data TcPluginRewriteResult
  = TcPluginNoRewrite
  | TcPluginRewriteTo
    { tcPluginReduction    :: !Reduction
    , tcRewriterNewWanteds :: [Ct]
    }

data Reduction = Reduction Coercion !Type

Stephanie Weirich’s challenge

Stephanie Weirich posed the challenge of implementing a strongly-typed representation of System F. [Talk] [GitHub]


-- | A term, in the given context, of the given type.
data Term ϕ a where
  -- | Function application.
  (:$) :: Term ϕ (a :-> b) -> Term ϕ a -> Term ϕ b
  -- | Type application.
  (:@) :: Term ϕ (Forall ty) -> Sing a -> Term ϕ (SubOne a ty)
  -- Elided: variables, literals, Lambda abstraction, Big Lambda.

Type-level substitution


type ApplySub ::
    forall (kϕ :: KContext) (kψ :: KContext) (k :: Kind).
    Sub kϕ kψ -> Type kϕ k -> Type kψ k

-- | Apply a substitution to a type.
type family ApplySub s ty where {..}

-- | Substitute a single type.
type SubOne a b = ApplySub ('Extend 'Id a) b

Problem


ApplySub s (ApplySub t a) ~ ApplySub ('Compose s t) a

Type family argument on the LHS of a type family equation.

Might also need to use Givens:

[G] ApplySub t a ~ f

ApplySub s f
  ~~>
  ApplySub ('Compose s t) a

Rewriting plugin


pluginRewrite :: PluginDefs -> UniqFM TyCon TcPluginRewriter
pluginRewrite defs@( PluginDefs { applySubTyCon } ) =
  listToUFM
    [ ( applySubTyCon, rewriteApplySub defs ) ]

rewriteApplySub
  :: PluginDefs
  -> [ Ct ]     -- Givens
  -> [ Type ]   -- Arguments to ApplySub
  -> TcPluginM TcPluginRewriteResult
rewriteApplySub defs givens [ kϕ, kψ, k, sub, sub_arg ]
  = ...


rewriteApplySub defs givens [ kϕ, kψ, k, sub, sub_arg ]
  = ...


[G] ApplySub kϕ0 kϕ l t a ~ sub_arg

rewriteApplySub defs givens [ kϕ, kψ, k, sub, sub_arg ]
  ~~>
  TcPluginRewriteTo
    { tcPluginReduction =
        mkPluginNomRedn "ApplySub s (ApplySub t a)" $
          mkTyConApp composeTyCon [kϕ0, kϕ, kψ, sub, t]
    , tcRewriterNewWanteds = [] }

Type-checking plugins in practice?

Typed holes


import GHC.TcPlugin.API

-- Want to use a 'Coercion' as equality constraint evidence.
help :: Coercion -> EvTerm
help = _


    * Found hole: _ :: Coercion -> EvTerm
    * Valid hole fits include
        evCoercion :: Coercion -> EvTerm

Typed holes (2)


import GHC.TcPlugin.API

-- Want to apply a 'TyCon' to arguments.
help :: TyCon -> [Type] -> Type
help = _


    * Found hole: _ :: TyCon -> [Type] -> Type
    * Valid hole fits include
        mkTyConApp :: TyCon -> [Type] -> Type

Hackage documentation

Custom type errors

Helpers for custom type errors.


myErrMsg :: Type -> TcPluginErrorMessage
myErrMsg ty =
   Txt "Couldn't handle type " :|: PrintType ty :|: Txt "."
  -- Same process as with the 'ErrorMessage' data kind.

errorWantedCt :: MonadTcPluginWork m => CtLoc -> Type -> m Ct
errorWantedCt loc ty =
   newWanted loc $ mkTcPluginErrorTy (myErrMsg ty)

Available now!

Get your hands dirty: ghc-tcplugin-api on Hackage.

Type-family rewriting functionality in GHC 9.4.
Compatibility layer in ghc-tcplugin-api for GHC 9.0 and 9.2.

ghc-tcplugin-api on GitHub contains the System F example.
Slides: github.com/sheaf/HIW-tcplugins

Upcoming Tweag blog post: writing and debugging a type-checking plugin.
In depth review of GHC constraint solving and type family reduction.