free-foil-0.6.0: src/Data/ZipMatchK/TH.hs
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
-- | Template Haskell derivation of 'ZipMatchK' and 'ZipMatchFoldK' instances.
--
-- The generic instance (the one obtained 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 large
-- signature that is worth roughly a factor of two in time and in allocation on
-- 'Control.Monad.Free.Foil.alphaEquiv' (see the @zipmatchk@ benchmark), and
-- unlike the generic instance the derived one 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 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. The signatures that
-- 'Control.Monad.Free.Foil.TH.MkFreeFoil.mkFreeFoil' generates from a grammar
-- with several syntactic categories are of that kind:
--
-- > concat <$> traverse deriveZipMatchK2 [''Term'Sig, ''OpArg'Sig, ''Type'Sig]
module Data.ZipMatchK.TH (
deriveZipMatchK,
deriveZipMatchK1,
deriveZipMatchK2,
deriveZipMatchKWith,
zipMatchKInstance,
zipMatchFoldKInstance,
) 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.Monoid (Ap (..))
import Data.ZipMatchK.Generic (ZipMatchFoldK (..), ZipMatchK (..))
import Data.ZipMatchK.Mappings (Folds (..), Mappings (..))
-- | Derive the 'ZipMatchK' and 'ZipMatchFoldK' instances, 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
-- >
-- > instance ZipMatchFoldK TermSig where
-- > zipMatchFoldK (f :%: g :%: F0) x y = case (x, y) of
-- > (AppSig l1 l2, AppSig r1 r2) -> Just (g l1 r1 <> g l2 r2)
-- > (LamSig l1, LamSig r1) -> Just (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.
--
-- @since 0.3.2
deriveZipMatchK :: Name -> Q [Dec]
deriveZipMatchK = deriveZipMatchKWith Nothing
-- | Derive the 'ZipMatchK' and 'ZipMatchFoldK' instances for a functor, zipping
-- the last type parameter and fixing the rest.
--
-- @since 0.3.2
deriveZipMatchK1 :: Name -> Q [Dec]
deriveZipMatchK1 = deriveZipMatchKWith (Just 1)
-- | Derive the 'ZipMatchK' and 'ZipMatchFoldK' instances 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
-- > ...
--
-- @since 0.3.2
deriveZipMatchK2 :: Name -> Q [Dec]
deriveZipMatchK2 = deriveZipMatchKWith (Just 2)
-- | Derive the 'ZipMatchK' and 'ZipMatchFoldK' instances, 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.
--
-- @since 0.3.2
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 zippedSet = Set.fromList (drop (length params - n) params)
fieldsOf <- concat <$> mapM (constructorFields params zippedSet) cons
(++) <$> zipMatchKInstance typeName params n fieldsOf
<*> zipMatchFoldKInstance typeName params n fieldsOf
-- | The 'ZipMatchK' instance of a data type, given its name, its type
-- parameters in order, how many of the last ones are zipped, and the fields
-- of each constructor in terms of those parameters. This is what
-- 'deriveZipMatchKWith' generates after reifying the type. It is for code
-- that declares the type in the same splice, such as
-- 'Control.Monad.Free.Foil.TH.MkFreeFoil.mkFreeFoil', which cannot reify it.
--
-- @since 0.6.0
zipMatchKInstance :: Name -> [Name] -> Int -> [(Name, [Type])] -> Q [Dec]
zipMatchKInstance typeName params n fieldsOf = 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
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 'ZipMatchFoldK' instance of a data type, from the same arguments as
-- 'zipMatchKInstance'. A field of a zipped parameter contributes the result of
-- its folding function, a field of a type applied to zipped parameters is
-- folded by that type's 'ZipMatchFoldK', and any other field is compared by
-- its 'ZipMatchK' and contributes nothing. The contributions are combined with
-- '<>', nested to the right, so that a lazy monoid such as 'Data.Monoid.All'
-- stops at the first field that fails. The method is @INLINE@, so that at a
-- call site with a known monoid and known folding functions the match
-- compiles to a @case@ on the two values with the functions applied in place.
--
-- @since 0.6.0
zipMatchFoldKInstance :: Name -> [Name] -> Int -> [(Name, [Type])] -> Q [Dec]
zipMatchFoldKInstance typeName params n fieldsOf = do
let (fixed, zipped) = splitAt (length params - n) params
zippedSet = Set.fromList zipped
-- One folding function per zipped parameter, in order.
foldVars <- mapM (const (newName "_f")) zipped
let folds = zip zipped (map VarE foldVars)
foldsPat = foldr (\v p -> InfixP (VarP v) '(:%:) p) (ConP 'F0 [] []) foldVars
clauses <- mapM (foldConClause zippedSet folds) 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 'zipMatchFoldK
[Clause [foldsPat, VarP x, VarP y] (NormalB body) []]
headType = foldl' AppT (ConT typeName) (map VarT fixed)
types = concatMap snd fieldsOf
needsZipMatchK = foldMap (zipMatchKVars zippedSet) types
needsZipMatchFoldK = foldMap (zipMatchFoldKVars zippedSet) types
context = [ AppT (ConT ''ZipMatchK) (VarT v) | v <- fixed, v `Set.member` needsZipMatchK ]
++ [ AppT (ConT ''ZipMatchFoldK) (VarT v) | v <- fixed, v `Set.member` needsZipMatchFoldK ]
return [InstanceD Nothing context (AppT (ConT ''ZipMatchFoldK) headType)
[PragmaD (InlineP 'zipMatchFoldK Inline FunLike AllPhases), impl]]
-- | How a field takes part in a folding match.
data FieldFolder
= Total Exp -- ^ @l -> r -> m@: a zipped parameter.
| Partial Exp -- ^ @l -> r -> Maybe m@: a type applied to zipped parameters.
| Check Exp -- ^ @l -> r -> Maybe c@: a type without zipped parameters.
-- | Compile a field type into a 'FieldFolder', as 'fieldZipper' does for
-- 'zipMatchWithK'. The components of a nested type can fail to match, so when
-- one of them does, all of them are folded in @'Ap' 'Maybe' m@.
fieldFolder :: Set.Set Name -> [(Name, Exp)] -> Type -> Q FieldFolder
fieldFolder zipped folds type_ = case stripType type_ of
VarT v | Just f <- lookup v folds -> return (Total f)
t | not (mentionsZipped zipped t) -> return (Check (AppE (VarE 'zipMatchWithK) (ConE 'M0)))
| otherwise -> do
let (_headType, args) = unApply t
zippedArgs = dropWhile (not . mentionsZipped zipped) args
argFolders <- mapM (fieldFolder zipped folds) zippedArgs
let list = foldr (\e m -> InfixE (Just e) (ConE '(:%:)) (Just m)) (ConE 'F0)
total = [ e | Total e <- argFolders ]
if length total == length argFolders
then return (Partial (AppE (VarE 'zipMatchFoldK) (list total)))
else do
wrapped <- mapM inAp argFolders
l <- newName "l"
r <- newName "r"
let call = AppE (AppE (AppE (VarE 'zipMatchFoldK) (list wrapped)) (VarE l)) (VarE r)
return (Partial (LamE [VarP l, VarP r]
(InfixE (Just (VarE 'getAp)) (VarE '(=<<)) (Just call))))
where
-- A folding function into @'Ap' 'Maybe' m@.
inAp folder = do
a <- newName "a"
b <- newName "b"
let applied e = AppE (AppE e (VarE a)) (VarE b)
ap e = LamE [VarP a, VarP b] (AppE (ConE 'Ap) e)
return $ case folder of
Total e -> ap (AppE (ConE 'Just) (applied e))
Partial e -> ap (applied e)
Check e -> ap (InfixE (Just (VarE 'mempty)) (VarE '(<$)) (Just (applied e)))
-- | One @case@ alternative of 'zipMatchFoldK', matching a constructor with
-- itself: the fields that can fail are matched first, then the contributions
-- are combined.
foldConClause :: Set.Set Name -> [(Name, Exp)] -> (Name, [Type]) -> Q Match
foldConClause zipped folds (conName, types) = do
lvars <- mapM (const (newName "l")) types
rvars <- mapM (const (newName "r")) types
folders <- mapM (fieldFolder zipped folds) types
steps <- sequence
[ case folder of
Total f -> return (Nothing, Just (AppE (AppE f (VarE l)) (VarE r)))
Partial f -> do
m <- newName "m"
return (Just (AppE (AppE f (VarE l)) (VarE r), VarP m), Just (VarE m))
Check f ->
return (Just (AppE (AppE f (VarE l)) (VarE r), WildP), Nothing)
| (folder, (l, r)) <- zip folders (zip lvars rvars) ]
let contributions = [ e | (_, Just e) <- steps ]
combined = case contributions of
[] -> VarE 'mempty
_ -> foldr1 (\a b -> InfixE (Just a) (VarE '(<>)) (Just b)) contributions
matchJust (scrutinee, pat) rest =
CaseE scrutinee
[ Match (ConP 'Nothing [] []) (NormalB (ConE 'Nothing)) []
, Match (ConP 'Just [] [pat]) (NormalB rest) [] ]
body = foldr matchJust (AppE (ConE 'Just) combined) [ c | (Just c, _) <- steps ]
return (Match
(TupP [ ConP conName [] (map VarP lvars)
, ConP conName [] (map VarP rvars) ])
(NormalB body) [])
-- | The fixed type parameters whose 'ZipMatchK' instance folding a field of
-- this type needs: those of a type without zipped parameters, and the fixed
-- arguments of a type applied to zipped ones.
zipMatchKVars :: Set.Set Name -> Type -> Set.Set Name
zipMatchKVars zipped type_ = case stripType type_ of
VarT v | v `Set.member` zipped -> Set.empty
t | not (mentionsZipped zipped t) -> freeVarsOfType t
| otherwise ->
let (_headType, args) = unApply t
(fixedArgs, zippedArgs) = break (mentionsZipped zipped) args
in foldMap freeVarsOfType fixedArgs <> foldMap (zipMatchKVars zipped) zippedArgs
-- | The fixed type parameters whose 'ZipMatchFoldK' instance folding a field of
-- this type needs: a parameter applied to zipped ones.
zipMatchFoldKVars :: Set.Set Name -> Type -> Set.Set Name
zipMatchFoldKVars zipped type_ = case stripType type_ of
t | not (mentionsZipped zipped t) -> Set.empty
| otherwise ->
let (headType, args) = unApply t
zippedArgs = dropWhile (not . mentionsZipped zipped) args
here = case stripType headType of
VarT v | null (takeWhile (not . mentionsZipped zipped) args) -> Set.singleton v
_ -> Set.empty
in here <> foldMap (zipMatchFoldKVars zipped) zippedArgs
-- | 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