{-# 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 :: Name -> Q [Dec]
deriveZipMatchK = Maybe Int -> Name -> Q [Dec]
deriveZipMatchKWith Maybe Int
forall a. Maybe a
Nothing

-- | Derive a 'ZipMatchK' instance for a functor, zipping the last type parameter
-- and fixing the rest.
deriveZipMatchK1 :: Name -> Q [Dec]
deriveZipMatchK1 :: Name -> Q [Dec]
deriveZipMatchK1 = Maybe Int -> Name -> Q [Dec]
deriveZipMatchKWith (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
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 :: Name -> Q [Dec]
deriveZipMatchK2 = Maybe Int -> Name -> Q [Dec]
deriveZipMatchKWith (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
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 :: Maybe Int -> Name -> Q [Dec]
deriveZipMatchKWith Maybe Int
arity Name
typeName = do
  ([TyVarBndr BndrVis]
tvars, [Con]
cons) <- Name -> Q ([TyVarBndr BndrVis], [Con])
reifyDataType Name
typeName
  let params :: [Name]
params = (TyVarBndr BndrVis -> Name) -> [TyVarBndr BndrVis] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr BndrVis -> Name
forall a. TyVarBndr a -> Name
tvarName [TyVarBndr BndrVis]
tvars
      n :: Int
n = case Maybe Int
arity of
            Maybe Int
Nothing -> [Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
params
            Just Int
k  -> Int
k
  if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
params
    then String -> Q [Dec]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (Name -> String
forall a. Show a => a -> String
show Name
typeName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" has " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show ([Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
params)
                String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" type parameter(s), cannot zip the last " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n)
    else do
      let ([Name]
fixed, [Name]
zipped) = Int -> [Name] -> ([Name], [Name])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
params Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) [Name]
params
          zippedSet :: Set Name
zippedSet = [Name] -> Set Name
forall a. Ord a => [a] -> Set a
Set.fromList [Name]
zipped

      -- One zipping function per zipped parameter, in order.
      [Name]
mappingVars <- (Name -> Q Name) -> [Name] -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Q Name -> Name -> Q Name
forall a b. a -> b -> a
const (String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"_f")) [Name]
zipped
      let mappings :: [(Name, Exp)]
mappings = [Name] -> [Exp] -> [(Name, Exp)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
zipped ((Name -> Exp) -> [Name] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE [Name]
mappingVars)
          mappingsPat :: Pat
mappingsPat = (Name -> Pat -> Pat) -> Pat -> [Name] -> Pat
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Name
v Pat
p -> Pat -> Name -> Pat -> Pat
InfixP (Name -> Pat
VarP Name
v) '(:^:) Pat
p) (Name -> [Type] -> [Pat] -> Pat
ConP 'M0 [] []) [Name]
mappingVars

      [(Name, [Type])]
fieldsOf <- [[(Name, [Type])]] -> [(Name, [Type])]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Name, [Type])]] -> [(Name, [Type])])
-> Q [[(Name, [Type])]] -> Q [(Name, [Type])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Con -> Q [(Name, [Type])]) -> [Con] -> Q [[(Name, [Type])]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ([Name] -> Set Name -> Con -> Q [(Name, [Type])]
constructorFields [Name]
params Set Name
zippedSet) [Con]
cons
      [Match]
clauses <- ((Name, [Type]) -> Q Match) -> [(Name, [Type])] -> Q [Match]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Set Name -> [(Name, Exp)] -> (Name, [Type]) -> Q Match
conClause Set Name
zippedSet [(Name, Exp)]
mappings) [(Name, [Type])]
fieldsOf
      let fallthrough :: [Match]
fallthrough
            | [(Name, [Type])] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Name, [Type])]
fieldsOf Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 = [Pat -> Body -> [Dec] -> Match
Match Pat
WildP (Exp -> Body
NormalB (Name -> Exp
ConE 'Nothing)) []]
            | Bool
otherwise           = []

      Name
x <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"x"
      Name
y <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"y"
      let body :: Exp
body = Exp -> [Match] -> Exp
CaseE ([Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
VarE Name
x), Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
VarE Name
y)]) ([Match]
clauses [Match] -> [Match] -> [Match]
forall a. [a] -> [a] -> [a]
++ [Match]
fallthrough)
          impl :: Dec
impl = Name -> [Clause] -> Dec
FunD 'zipMatchWithK
            [[Pat] -> Body -> [Dec] -> Clause
Clause [Pat
mappingsPat, Name -> Pat
VarP Name
x, Name -> Pat
VarP Name
y] (Exp -> Body
NormalB Exp
body) []]

          headType :: Type
headType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Type -> Type -> Type
AppT (Name -> Type
ConT Name
typeName) ((Name -> Type) -> [Name] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Type
VarT [Name]
fixed)
          usedVars :: Set Name
usedVars = ((Name, [Type]) -> Set Name) -> [(Name, [Type])] -> Set Name
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Type -> Set Name) -> [Type] -> Set Name
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Type -> Set Name
freeVarsOfType ([Type] -> Set Name)
-> ((Name, [Type]) -> [Type]) -> (Name, [Type]) -> Set Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [Type]) -> [Type]
forall a b. (a, b) -> b
snd) [(Name, [Type])]
fieldsOf
          context :: [Type]
context = [ Type -> Type -> Type
AppT (Name -> Type
ConT ''ZipMatchK) (Name -> Type
VarT Name
v)
                    | Name
v <- [Name]
fixed, Name
v Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Name
usedVars ]

      [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
forall a. Maybe a
Nothing [Type]
context (Type -> Type -> Type
AppT (Name -> Type
ConT ''ZipMatchK) Type
headType) [Dec
impl]]

-- | The type parameters and constructors of a @data@ or @newtype@ declaration.
reifyDataType :: Name -> Q ([TyVarBndr BndrVis], [Con])
reifyDataType :: Name -> Q ([TyVarBndr BndrVis], [Con])
reifyDataType Name
typeName = Name -> Q Info
reify Name
typeName Q Info
-> (Info -> Q ([TyVarBndr BndrVis], [Con]))
-> Q ([TyVarBndr BndrVis], [Con])
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  TyConI (DataD [Type]
_ Name
_ [TyVarBndr BndrVis]
tvars Maybe Type
_ [Con]
cons [DerivClause]
_)      -> ([TyVarBndr BndrVis], [Con]) -> Q ([TyVarBndr BndrVis], [Con])
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([TyVarBndr BndrVis]
tvars, [Con]
cons)
  TyConI (NewtypeD [Type]
_ Name
_ [TyVarBndr BndrVis]
tvars Maybe Type
_ Con
con [DerivClause]
_)    -> ([TyVarBndr BndrVis], [Con]) -> Q ([TyVarBndr BndrVis], [Con])
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([TyVarBndr BndrVis]
tvars, [Con
con])
  Info
_ -> String -> Q ([TyVarBndr BndrVis], [Con])
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (Name -> String
forall a. Show a => a -> String
show Name
typeName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" 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 :: [Name] -> Set Name -> Con -> Q [(Name, [Type])]
constructorFields [Name]
params Set Name
zipped = \case
  NormalC Name
conName [BangType]
types      -> [(Name, [Type])] -> Q [(Name, [Type])]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [(Name
conName, (BangType -> Type) -> [BangType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map BangType -> Type
forall a b. (a, b) -> b
snd [BangType]
types)]
  RecC Name
conName [VarBangType]
types         -> [(Name, [Type])] -> Q [(Name, [Type])]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [(Name
conName, (VarBangType -> Type) -> [VarBangType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (BangType -> Type
forall a b. (a, b) -> b
snd (BangType -> Type)
-> (VarBangType -> BangType) -> VarBangType -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VarBangType -> BangType
removeName) [VarBangType]
types)]
  InfixC BangType
l Name
conName BangType
r         -> [(Name, [Type])] -> Q [(Name, [Type])]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [(Name
conName, [BangType -> Type
forall a b. (a, b) -> b
snd BangType
l, BangType -> Type
forall a b. (a, b) -> b
snd BangType
r])]
  RecGadtC [Name]
conNames [VarBangType]
types Type
r  -> [Name] -> Set Name -> Con -> Q [(Name, [Type])]
constructorFields [Name]
params Set Name
zipped
                                  ([Name] -> [BangType] -> Type -> Con
GadtC [Name]
conNames ((VarBangType -> BangType) -> [VarBangType] -> [BangType]
forall a b. (a -> b) -> [a] -> [b]
map VarBangType -> BangType
removeName [VarBangType]
types) Type
r)
  ForallC [TyVarBndr Specificity]
_ [Type]
context_ Con
con
    | [Type] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Type]
context_ -> [Name] -> Set Name -> Con -> Q [(Name, [Type])]
constructorFields [Name]
params Set Name
zipped Con
con
    | Bool
otherwise -> String -> Q [(Name, [Type])]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"constructor contexts are not supported by deriveZipMatchK"
  GadtC [Name]
conNames [BangType]
types Type
retType -> do
    [(Name, Name)]
rename <- Type -> Q [(Name, Name)]
forall {m :: * -> *}. MonadFail m => Type -> m [(Name, Name)]
renaming Type
retType
    let fields :: [Type]
fields = (BangType -> Type) -> [BangType] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map ([(Name, Name)] -> Type -> Type
substTypeVars [(Name, Name)]
rename (Type -> Type) -> (BangType -> Type) -> BangType -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BangType -> Type
forall a b. (a, b) -> b
snd) [BangType]
types
        escaping :: Set Name
escaping = (Type -> Set Name) -> [Type] -> Set Name
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Type -> Set Name
freeVarsOfType [Type]
fields Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` [Name] -> Set Name
forall a. Ord a => [a] -> Set a
Set.fromList [Name]
params
    if Set Name -> Bool
forall a. Set a -> Bool
Set.null Set Name
escaping
      then [(Name, [Type])] -> Q [(Name, [Type])]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [ (Name
conName, [Type]
fields) | Name
conName <- [Name]
conNames ]
      else String -> Q [(Name, [Type])]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"existentials are not supported by deriveZipMatchK: "
                  String -> String -> String
forall a. [a] -> [a] -> [a]
++ [Name] -> String
forall a. Show a => a -> String
show (Set Name -> [Name]
forall a. Set a -> [a]
Set.toList Set Name
escaping))
  where
    -- The constructor's variables, named as the declaration names them.
    renaming :: Type -> m [(Name, Name)]
renaming Type
retType = do
      let (Type
_, [Type]
args) = Type -> (Type, [Type])
unApply Type
retType
      if [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
args Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
params
        then String -> m [(Name, Name)]
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"the constructor does not return the type being derived for"
        else [m (Name, Name)] -> m [(Name, Name)]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence
          [ case Type -> Type
stripType Type
arg of
              VarT Name
v -> (Name, Name) -> m (Name, Name)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
v, Name
param)
              Type
_ | Name
param Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Name
zipped ->
                    String -> m (Name, Name)
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"the constructor fixes the type parameter "
                           String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
param String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
", which is being zipped")
                | Bool
otherwise -> (Name, Name) -> m (Name, Name)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
param, Name
param)  -- a fixed index; no field can mention it
          | (Type
arg, Name
param) <- [Type] -> [Name] -> [(Type, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Type]
args [Name]
params ]

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

-- | Rename type variables.
substTypeVars :: [(Name, Name)] -> Type -> Type
substTypeVars :: [(Name, Name)] -> Type -> Type
substTypeVars [(Name, Name)]
rename = Type -> Type
go
  where
    go :: Type -> Type
go = \case
      VarT Name
v        -> Name -> Type
VarT (Name -> (Name -> Name) -> Maybe Name -> Name
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Name
v Name -> Name
forall a. a -> a
id (Name -> [(Name, Name)] -> Maybe Name
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
v [(Name, Name)]
rename))
      AppT Type
f Type
x      -> Type -> Type -> Type
AppT (Type -> Type
go Type
f) (Type -> Type
go Type
x)
      AppKindT Type
t Type
k  -> Type -> Type -> Type
AppKindT (Type -> Type
go Type
t) Type
k
      SigT Type
t Type
k      -> Type -> Type -> Type
SigT (Type -> Type
go Type
t) Type
k
      ParensT Type
t     -> Type -> Type
ParensT (Type -> Type
go Type
t)
      InfixT Type
l Name
n Type
r  -> Type -> Name -> Type -> Type
InfixT (Type -> Type
go Type
l) Name
n (Type -> Type
go Type
r)
      UInfixT Type
l Name
n Type
r -> Type -> Name -> Type -> Type
UInfixT (Type -> Type
go Type
l) Name
n (Type -> Type
go Type
r)
      Type
t             -> Type
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 :: Set Name -> [(Name, Exp)] -> Type -> Q Exp
fieldZipper Set Name
zipped [(Name, Exp)]
mappings Type
type_ = case Type -> Type
stripType Type
type_ of
  VarT Name
v | Just Exp
f <- Name -> [(Name, Exp)] -> Maybe Exp
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
v [(Name, Exp)]
mappings -> Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
f
  Type
t -> do
    let (Type
_headType, [Type]
args) = Type -> (Type, [Type])
unApply Type
t
        zippedArgs :: [Type]
zippedArgs = (Type -> Bool) -> [Type] -> [Type]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Bool -> Bool
not (Bool -> Bool) -> (Type -> Bool) -> Type -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set Name -> Type -> Bool
mentionsZipped Set Name
zipped) [Type]
args
    [Exp]
argZippers <- (Type -> Q Exp) -> [Type] -> Q [Exp]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Set Name -> [(Name, Exp)] -> Type -> Q Exp
fieldZipper Set Name
zipped [(Name, Exp)]
mappings) [Type]
zippedArgs
    let ms :: Exp
ms = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Exp
z Exp
m -> Maybe Exp -> Exp -> Maybe Exp -> Exp
InfixE (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
z) (Name -> Exp
ConE '(:^:)) (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
m)) (Name -> Exp
ConE 'M0) [Exp]
argZippers
    Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'zipMatchWithK) Exp
ms)

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

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

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

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