TH.hs

{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns    #-}
-- | Template Haskell derivation of 'ZipMatchK' instances.
--
-- The generic instance (the one you get by writing @instance ZipMatchK TermSig@
-- with no body) converts a node into its "Generics.Kind" representation on
-- every comparison, and converts the result back. The representation of a
-- constructor is a chain of @L1@\/@R1@ wrappers as long as that constructor's
-- index, so the cost grows with the number of constructors in the signature,
-- and comparing terms is most of what a typechecker does.
--
-- The derivers here generate the instance that one would otherwise write out by
-- hand: a @case@ over the two nodes, allocating only its result. On a
-- 44-constructor signature that is worth a factor of 1.8 in time and 2.3 in
-- allocation on 'Control.Monad.Free.Foil.alphaEquiv' (see the @zipmatchk@
-- benchmark), and the derived instance does not get slower as the signature
-- grows.
--
-- The module the splice appears in needs at least
--
-- > {-# LANGUAGE GADTs #-}
-- > {-# LANGUAGE TemplateHaskell #-}
-- > {-# LANGUAGE TypeFamilies #-}
--
-- Signatures that refer to one another — as the ones generated by
-- 'Control.Monad.Free.Foil.TH.MkFreeFoil.mkFreeFoil' from a grammar with
-- several syntactic categories do — have to be derived in a __single splice__,
-- since a top-level splice ends a declaration group and an instance from a later
-- group is not visible to an earlier one:
--
-- > concat <$> traverse deriveZipMatchK2 [''Term'Sig, ''OpArg'Sig, ''Type'Sig]
module Data.ZipMatchK.TH (
  deriveZipMatchK,
  deriveZipMatchK1,
  deriveZipMatchK2,
  deriveZipMatchKWith,
) where

import           Data.List             (foldl')
import qualified Data.Set              as Set
import           Language.Haskell.TH

import           Control.Monad.Foil.TH.Util (removeName, tvarName)
import           Data.ZipMatchK.Generic    (ZipMatchK (..))
import           Data.ZipMatchK.Mappings   (Mappings (..))

-- | Derive a 'ZipMatchK' instance, zipping /all/ type parameters.
--
-- For a signature bifunctor
--
-- > data TermSig scope term = AppSig term term | LamSig scope
-- > deriveZipMatchK ''TermSig
--
-- this generates
--
-- > instance ZipMatchK TermSig where
-- >   zipMatchWithK (f :^: g :^: M0) x y = case (x, y) of
-- >     (AppSig l1 l2, AppSig r1 r2) -> AppSig <$> g l1 r1 <*> g l2 r2
-- >     (LamSig l1, LamSig r1)       -> LamSig <$> f l1 r1
-- >     _                            -> Nothing
--
-- Use 'deriveZipMatchK2' when the type has extra parameters that should stay
-- fixed (an annotation, say), as a signature generated by
-- "Control.Monad.Free.Foil.TH.MkFreeFoil" does.
deriveZipMatchK :: Name -> Q [Dec]
deriveZipMatchK = deriveZipMatchKWith Nothing

-- | Derive a 'ZipMatchK' instance for a functor, zipping the last type parameter
-- and fixing the rest.
deriveZipMatchK1 :: Name -> Q [Dec]
deriveZipMatchK1 = deriveZipMatchKWith (Just 1)

-- | Derive a 'ZipMatchK' instance for a signature bifunctor, zipping the last two
-- type parameters (the scoped terms and the terms) and fixing the rest.
--
-- For a signature with an extra parameter (a source position, for instance)
--
-- > data Term'Sig a scope term = AppSig a term term | LamSig a scope
-- > deriveZipMatchK2 ''Term'Sig
--
-- this generates
--
-- > instance ZipMatchK a => ZipMatchK (Term'Sig a) where
-- >   zipMatchWithK (f :^: g :^: M0) x y = case (x, y) of
-- >     (AppSig l1 l2 l3, AppSig r1 r2 r3) ->
-- >       AppSig <$> zipMatchWithK M0 l1 r1 <*> g l2 r2 <*> g l3 r3
-- >     ...
deriveZipMatchK2 :: Name -> Q [Dec]
deriveZipMatchK2 = deriveZipMatchKWith (Just 2)

-- | Derive a 'ZipMatchK' instance, zipping the last @n@ type parameters and
-- fixing the rest. 'Nothing' zips all of them.
--
-- Every fixed parameter that occurs in a field gets a 'ZipMatchK' constraint in
-- the instance context.
deriveZipMatchKWith :: Maybe Int -> Name -> Q [Dec]
deriveZipMatchKWith arity typeName = do
  (tvars, cons) <- reifyDataType typeName
  let params = map tvarName tvars
      n = case arity of
            Nothing -> length params
            Just k  -> k
  if n > length params
    then fail (show typeName ++ " has " ++ show (length params)
                ++ " type parameter(s), cannot zip the last " ++ show n)
    else do
      let (fixed, zipped) = splitAt (length params - n) params
          zippedSet = Set.fromList zipped

      -- One zipping function per zipped parameter, in order.
      mappingVars <- mapM (const (newName "_f")) zipped
      let mappings = zip zipped (map VarE mappingVars)
          mappingsPat = foldr (\v p -> InfixP (VarP v) '(:^:) p) (ConP 'M0 [] []) mappingVars

      fieldsOf <- concat <$> mapM (constructorFields params zippedSet) cons
      clauses <- mapM (conClause zippedSet mappings) fieldsOf
      let fallthrough
            | length fieldsOf > 1 = [Match WildP (NormalB (ConE 'Nothing)) []]
            | otherwise           = []

      x <- newName "x"
      y <- newName "y"
      let body = CaseE (TupE [Just (VarE x), Just (VarE y)]) (clauses ++ fallthrough)
          impl = FunD 'zipMatchWithK
            [Clause [mappingsPat, VarP x, VarP y] (NormalB body) []]

          headType = foldl' AppT (ConT typeName) (map VarT fixed)
          usedVars = foldMap (foldMap freeVarsOfType . snd) fieldsOf
          context = [ AppT (ConT ''ZipMatchK) (VarT v)
                    | v <- fixed, v `Set.member` usedVars ]

      return [InstanceD Nothing context (AppT (ConT ''ZipMatchK) headType) [impl]]

-- | The type parameters and constructors of a @data@ or @newtype@ declaration.
reifyDataType :: Name -> Q ([TyVarBndr BndrVis], [Con])
reifyDataType typeName = reify typeName >>= \case
  TyConI (DataD _ _ tvars _ cons _)      -> return (tvars, cons)
  TyConI (NewtypeD _ _ tvars _ con _)    -> return (tvars, [con])
  _ -> fail (show typeName ++ " is not a data type or a newtype")

-- | A constructor's name and the types of its fields, in terms of the type
-- parameters of the declaration it belongs to.
--
-- 'Control.Monad.Free.Foil.TH.MkFreeFoil.mkFreeFoil' declares its signatures in
-- GADT syntax, and 'reify' returns such a constructor wrapped in a 'ForallC',
-- with variable names of its own. So the fields are renamed through the
-- constructor's return type before anything else looks at them.
constructorFields :: [Name] -> Set.Set Name -> Con -> Q [(Name, [Type])]
constructorFields params zipped = \case
  NormalC conName types      -> return [(conName, map snd types)]
  RecC conName types         -> return [(conName, map (snd . removeName) types)]
  InfixC l conName r         -> return [(conName, [snd l, snd r])]
  RecGadtC conNames types r  -> constructorFields params zipped
                                  (GadtC conNames (map removeName types) r)
  ForallC _ context_ con
    | null context_ -> constructorFields params zipped con
    | otherwise -> fail "constructor contexts are not supported by deriveZipMatchK"
  GadtC conNames types retType -> do
    rename <- renaming retType
    let fields = map (substTypeVars rename . snd) types
        escaping = foldMap freeVarsOfType fields `Set.difference` Set.fromList params
    if Set.null escaping
      then return [ (conName, fields) | conName <- conNames ]
      else fail ("existentials are not supported by deriveZipMatchK: "
                  ++ show (Set.toList escaping))
  where
    -- The constructor's variables, named as the declaration names them.
    renaming retType = do
      let (_, args) = unApply retType
      if length args /= length params
        then fail "the constructor does not return the type being derived for"
        else sequence
          [ case stripType arg of
              VarT v -> return (v, param)
              _ | param `Set.member` zipped ->
                    fail ("the constructor fixes the type parameter "
                           ++ show param ++ ", which is being zipped")
                | otherwise -> return (param, param)  -- a fixed index; no field can mention it
          | (arg, param) <- zip args params ]

-- | One @case@ alternative, zipping a constructor with itself.
conClause :: Set.Set Name -> [(Name, Exp)] -> (Name, [Type]) -> Q Match
conClause zipped mappings (conName, types) = do
  lvars <- mapM (const (newName "l")) types
  rvars <- mapM (const (newName "r")) types
  zippers <- mapM (fieldZipper zipped mappings) types
  -- Con <$> z1 l1 r1 <*> z2 l2 r2 <*> ...
  let zipField z l r = AppE (AppE z (VarE l)) (VarE r)
      apply acc (op, (z, (l, r))) = InfixE (Just acc) op (Just (zipField z l r))
      fields = zip (VarE '(<$>) : repeat (VarE '(<*>)))
                   (zip zippers (zip lvars rvars))
      body
        | null fields = AppE (ConE 'Just) (ConE conName)
        | otherwise   = foldl' apply (ConE conName) fields
  return (Match
            (TupP [ ConP conName [] (map VarP lvars)
                  , ConP conName [] (map VarP rvars) ])
            (NormalB body) [])

-- | Rename type variables.
substTypeVars :: [(Name, Name)] -> Type -> Type
substTypeVars rename = go
  where
    go = \case
      VarT v        -> VarT (maybe v id (lookup v rename))
      AppT f x      -> AppT (go f) (go x)
      AppKindT t k  -> AppKindT (go t) k
      SigT t k      -> SigT (go t) k
      ParensT t     -> ParensT (go t)
      InfixT l n r  -> InfixT (go l) n (go r)
      UInfixT l n r -> UInfixT (go l) n (go r)
      t             -> t

-- | Compile a field type into a zipping function @a -> b -> Maybe c@.
--
-- This mirrors what @ZipMatchFields@ does for the generic instance, but at
-- compile time, so nothing is reflected into a representation type:
--
-- * a zipped type parameter becomes the corresponding zipping function;
-- * any other type is zipped by its own 'ZipMatchK' instance, with the
--   arguments that mention zipped parameters passed as further zipping
--   functions.
--
-- So a field @[term]@ becomes @zipMatchWithK (g :^: M0)@ (asking for
-- @ZipMatchK []@), a field @Either a term@ becomes @zipMatchWithK (g :^: M0)@
-- (asking for @ZipMatchK (Either a)@), and a field @VarIdent@ becomes
-- @zipMatchWithK M0@ (asking for @ZipMatchK VarIdent@).
fieldZipper :: Set.Set Name -> [(Name, Exp)] -> Type -> Q Exp
fieldZipper zipped mappings type_ = case stripType type_ of
  VarT v | Just f <- lookup v mappings -> return f
  t -> do
    let (_headType, args) = unApply t
        zippedArgs = dropWhile (not . mentionsZipped zipped) args
    argZippers <- mapM (fieldZipper zipped mappings) zippedArgs
    let ms = foldr (\z m -> InfixE (Just z) (ConE '(:^:)) (Just m)) (ConE 'M0) argZippers
    return (AppE (VarE 'zipMatchWithK) ms)

-- | Does the type mention any of the zipped type parameters?
mentionsZipped :: Set.Set Name -> Type -> Bool
mentionsZipped zipped t = not (Set.null (Set.intersection zipped (freeVarsOfType t)))

-- | Split a type into its head and its arguments.
unApply :: Type -> (Type, [Type])
unApply = go []
  where
    go args (AppT f x)  = go (x : args) f
    go args (ParensT t) = go args t
    go args (SigT t _)  = go args t
    go args t           = (t, args)

-- | Drop parentheses and kind signatures.
stripType :: Type -> Type
stripType (ParensT t) = stripType t
stripType (SigT t _)  = stripType t
stripType t           = t

freeVarsOfType :: Type -> Set.Set Name
freeVarsOfType = \case
  VarT v      -> Set.singleton v
  AppT f x    -> freeVarsOfType f <> freeVarsOfType x
  AppKindT t _-> freeVarsOfType t
  SigT t _    -> freeVarsOfType t
  ParensT t   -> freeVarsOfType t
  InfixT l _ r-> freeVarsOfType l <> freeVarsOfType r
  UInfixT l _ r -> freeVarsOfType l <> freeVarsOfType r
  _           -> Set.empty