{-# OPTIONS_GHC -fno-warn-type-defaults      #-}
{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns    #-}
module Control.Monad.Free.Foil.TH.ZipMatch where

import           Language.Haskell.TH

import           Control.Monad.Foil.TH.Util
import           Control.Monad.Free.Foil

-- | Generate 'ZipMatch' instance for a given bifunctor.
deriveZipMatch
  :: Name -- ^ Type name for the signature bifunctor.
  -> Q [Dec]
deriveZipMatch :: Name -> Q [Dec]
deriveZipMatch Name
signatureT = do
  TyConI (DataD Cxt
_ctx Name
_name [TyVarBndr BndrVis]
signatureTVars Maybe Kind
_kind [Con]
signatureCons [DerivClause]
_deriv) <- Name -> Q Info
reify Name
signatureT

  case [TyVarBndr BndrVis] -> [TyVarBndr BndrVis]
forall a. [a] -> [a]
reverse [TyVarBndr BndrVis]
signatureTVars of
    (TyVarBndr BndrVis -> Name
forall a. TyVarBndr a -> Name
tvarName -> Name
term) : (TyVarBndr BndrVis -> Name
forall a. TyVarBndr a -> Name
tvarName -> Name
scope) : ([TyVarBndr BndrVis] -> [TyVarBndr BndrVis]
forall a. [a] -> [a]
reverse -> [TyVarBndr BndrVis]
params) -> do
      let signatureType :: Kind
signatureType = Name -> Cxt -> Kind
PeelConT Name
signatureT ((TyVarBndr BndrVis -> Kind) -> [TyVarBndr BndrVis] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Kind
VarT (Name -> Kind)
-> (TyVarBndr BndrVis -> Name) -> TyVarBndr BndrVis -> Kind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVarBndr BndrVis -> Name
forall a. TyVarBndr a -> Name
tvarName) [TyVarBndr BndrVis]
params)
      [Clause]
clauses <- [[Clause]] -> [Clause]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Clause]] -> [Clause]) -> Q [[Clause]] -> Q [Clause]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Con -> Q [Clause]) -> [Con] -> Q [[Clause]]
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 -> Name -> Con -> Q [Clause]
toClause Name
scope Name
term) [Con]
signatureCons
      let defaultClause :: Clause
defaultClause = [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
WildP, Pat
WildP] (Exp -> Body
NormalB (Name -> Exp
ConE 'Nothing)) []
      let instType :: Kind
instType = Kind -> Kind -> Kind
AppT (Name -> Kind
ConT ''ZipMatch) Kind
signatureType

      [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return
        [ Maybe Overlap -> Cxt -> Kind -> [Dec] -> Dec
InstanceD Maybe Overlap
forall a. Maybe a
Nothing [] Kind
instType
          [ Name -> [Clause] -> Dec
FunD 'zipMatch ([Clause]
clauses [Clause] -> [Clause] -> [Clause]
forall a. [a] -> [a] -> [a]
++ [Clause
defaultClause]) ]
        ]
    [TyVarBndr BndrVis]
_ -> String -> Q [Dec]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"cannot generate pattern synonyms"

  where
    toClause :: Name -> Name -> Con -> Q [Clause]
    toClause :: Name -> Name -> Con -> Q [Clause]
toClause Name
scope Name
term = Con -> Q [Clause]
go
      where
        go :: Con -> Q [Clause]
go = \case
          NormalC Name
conName [BangType]
types -> Name -> [BangType] -> Q [Clause]
mkClause Name
conName [BangType]
types
          RecC Name
conName [VarBangType]
types -> Con -> Q [Clause]
go (Name -> [BangType] -> Con
NormalC Name
conName ((VarBangType -> BangType) -> [VarBangType] -> [BangType]
forall a b. (a -> b) -> [a] -> [b]
map VarBangType -> BangType
removeName [VarBangType]
types))
          InfixC BangType
l Name
conName BangType
r -> Con -> Q [Clause]
go (Name -> [BangType] -> Con
NormalC Name
conName [BangType
l, BangType
r])
          ForallC [TyVarBndr Specificity]
_ Cxt
_ Con
con -> Con -> Q [Clause]
go Con
con
          GadtC [Name]
conNames [BangType]
types Kind
_retType -> [[Clause]] -> [Clause]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Clause]] -> [Clause]) -> Q [[Clause]] -> Q [Clause]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Name -> Q [Clause]) -> [Name] -> Q [[Clause]]
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
conName -> Name -> [BangType] -> Q [Clause]
mkClause Name
conName [BangType]
types) [Name]
conNames
          RecGadtC [Name]
conNames [VarBangType]
types Kind
retType -> Con -> Q [Clause]
go ([Name] -> [BangType] -> Kind -> Con
GadtC [Name]
conNames ((VarBangType -> BangType) -> [VarBangType] -> [BangType]
forall a b. (a -> b) -> [a] -> [b]
map VarBangType -> BangType
removeName [VarBangType]
types) Kind
retType)

        mkClause :: Name -> [BangType] -> Q [Clause]
        mkClause :: Name -> [BangType] -> Q [Clause]
mkClause Name
conName [BangType]
types = [Clause] -> Q [Clause]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return
          [[Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Cxt -> [Pat] -> Pat
ConP Name
conName [] [Pat]
lpats, Name -> Cxt -> [Pat] -> Pat
ConP Name
conName [] [Pat]
rpats]
            (Exp -> Body
NormalB (Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'Just) ((Exp -> Exp -> Exp) -> Exp -> [Exp] -> 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
AppE (Name -> Exp
ConE Name
conName) [Exp]
args))) []]
          where
            ([Pat]
lpats, [Pat]
rpats, [Exp]
args) = [(Pat, Pat, Exp)] -> ([Pat], [Pat], [Exp])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3
              [ case Kind
type_ of
                  VarT Name
typeName
                    | Name
typeName Name -> [Name] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Name
scope, Name
term] -> (Name -> Pat
VarP Name
l, Name -> Pat
VarP Name
r, [Maybe Exp] -> Exp
TupE [Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
VarE Name
l), Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
VarE Name
r)])
                  Kind
_ -> (Name -> Pat
VarP Name
l, Pat
WildP, Name -> Exp
VarE Name
l)
              | (Integer
i, (Bang
_bang, Kind
type_)) <- [Integer] -> [BangType] -> [(Integer, BangType)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Integer
0..] [BangType]
types
              , let l :: Name
l = String -> Name
mkName (String
"l" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
i)
              , let r :: Name
r = String -> Name
mkName (String
"r" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
i)
              ]