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

import           Control.Monad              (forM_)
import           Control.Monad.Foil.TH.Util
import           Control.Monad.Free.Foil
import           Data.List                  (nub)
import           Language.Haskell.TH
import           Language.Haskell.TH.Syntax

-- | Generate helpful pattern synonyms given a signature bifunctor.
mkPatternSynonyms
  :: Name -- ^ Type name for the signature bifunctor.
  -> Q [Dec]
mkPatternSynonyms :: Name -> Q [Dec]
mkPatternSynonyms 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
      ([Name]
names, [Dec]
decs) <- [(Name, Dec)] -> ([Name], [Dec])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Name, Dec)] -> ([Name], [Dec]))
-> ([[(Name, Dec)]] -> [(Name, Dec)])
-> [[(Name, Dec)]]
-> ([Name], [Dec])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[(Name, Dec)]] -> [(Name, Dec)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Name, Dec)]] -> ([Name], [Dec]))
-> Q [[(Name, Dec)]] -> Q ([Name], [Dec])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Con -> Q [(Name, Dec)]) -> [Con] -> Q [[(Name, Dec)]]
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 (Kind -> Name -> Name -> Con -> Q [(Name, Dec)]
mkPatternSynonym (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)) Name
scope Name
term) [Con]
signatureCons
      [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> Q [Dec]) -> [Dec] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ [Dec]
decs [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++
        [ Pragma -> Dec
PragmaD ([Name] -> Maybe Name -> Pragma
CompleteP ('Var Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: [Name] -> [Name]
forall a. Eq a => [a] -> [a]
nub [Name]
names) Maybe Name
forall a. Maybe a
Nothing)]
    [TyVarBndr BndrVis]
_ -> String -> Q [Dec]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"cannot generate pattern synonyms"

mkPatternSynonym :: Type -> Name -> Name -> Con -> Q [(Name, Dec)]
mkPatternSynonym :: Kind -> Name -> Name -> Con -> Q [(Name, Dec)]
mkPatternSynonym Kind
signatureType Name
scope Name
term = \case
  NormalC Name
conName [BangType]
types -> Kind -> Name -> Name -> Con -> Q [(Name, Dec)]
mkPatternSynonym Kind
signatureType Name
scope Name
term
    ([Name] -> [BangType] -> Kind -> Con
GadtC [Name
conName] [BangType]
types (Kind -> Kind -> Kind
AppT (Kind -> Kind -> Kind
AppT Kind
signatureType (Name -> Kind
VarT Name
scope)) (Name -> Kind
VarT Name
term)))

  RecC Name
conName [VarBangType]
types -> Kind -> Name -> Name -> Con -> Q [(Name, Dec)]
mkPatternSynonym Kind
signatureType Name
scope Name
term (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 -> Kind -> Name -> Name -> Con -> Q [(Name, Dec)]
mkPatternSynonym Kind
signatureType Name
scope Name
term (Name -> [BangType] -> Con
NormalC Name
conName [BangType
l, BangType
r])

  ForallC [TyVarBndr Specificity]
params Cxt
ctx Con
con -> do
    [ (Name
name, PatSynSigD Name
patName Kind
patType), (Name, Dec)
patD ] <- Kind -> Name -> Name -> Con -> Q [(Name, Dec)]
mkPatternSynonym Kind
signatureType Name
scope Name
term Con
con
    [(Name, Dec)] -> Q [(Name, Dec)]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return
      [ (Name
name, Name -> Kind -> Dec
PatSynSigD Name
patName ([TyVarBndr Specificity] -> Cxt -> Kind -> Kind
ForallT [TyVarBndr Specificity]
params Cxt
ctx Kind
patType))
      , (Name, Dec)
patD
      ]

  GadtC [Name]
conNames [BangType]
types Kind
_retType -> do
    let argsWithTypes :: [Either ((Name, Kind), (Name, Kind)) (Name, Kind)]
argsWithTypes = (Integer
 -> BangType -> Either ((Name, Kind), (Name, Kind)) (Name, Kind))
-> [Integer]
-> [BangType]
-> [Either ((Name, Kind), (Name, Kind)) (Name, Kind)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer
-> BangType -> Either ((Name, Kind), (Name, Kind)) (Name, Kind)
forall {a} {a}.
Show a =>
a -> (a, Kind) -> Either ((Name, Kind), (Name, Kind)) (Name, Kind)
toPatternArgType [Integer
0..] [BangType]
types
        argsWithTypes' :: [(Name, Kind)]
argsWithTypes' = (Either ((Name, Kind), (Name, Kind)) (Name, Kind)
 -> [(Name, Kind)])
-> [Either ((Name, Kind), (Name, Kind)) (Name, Kind)]
-> [(Name, Kind)]
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Either ((Name, Kind), (Name, Kind)) (Name, Kind) -> [(Name, Kind)]
forall {a}. Either (a, a) a -> [a]
collapse [Either ((Name, Kind), (Name, Kind)) (Name, Kind)]
argsWithTypes
        pats :: [Pat]
pats   = (Either ((Name, Kind), (Name, Kind)) (Name, Kind) -> Pat)
-> [Either ((Name, Kind), (Name, Kind)) (Name, Kind)] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Either ((Name, Kind), (Name, Kind)) (Name, Kind) -> Pat
forall {b} {b} {b}. Either ((Name, b), (Name, b)) (Name, b) -> Pat
toArg [Either ((Name, Kind), (Name, Kind)) (Name, Kind)]
argsWithTypes
        args :: [Name]
args  = ((Name, Kind) -> Name) -> [(Name, Kind)] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Kind) -> Name
forall a b. (a, b) -> a
fst [(Name, Kind)]
argsWithTypes'
        types' :: Cxt
types' = ((Name, Kind) -> Kind) -> [(Name, Kind)] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Name, Kind) -> Kind
forall a b. (a, b) -> b
snd [(Name, Kind)]
argsWithTypes'
    [Name] -> (Name -> Q ()) -> Q ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Name]
conNames ((Name -> Q ()) -> Q ()) -> (Name -> Q ()) -> Q ()
forall a b. (a -> b) -> a -> b
$ \Name
conName ->
      Q () -> Q ()
addModFinalizer (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ DocLoc -> String -> Q ()
putDoc (Name -> DocLoc
DeclDoc (Name -> Name
mkPatternName Name
conName))
        (String
"/Generated/ with '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show 'mkPatternSynonyms String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'. Pattern synonym for an '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show ''AST String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"' node of type '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
conName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'.")
    [(Name, Dec)] -> Q [(Name, Dec)]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Name, Dec)] -> Q [(Name, Dec)])
-> [(Name, Dec)] -> Q [(Name, Dec)]
forall a b. (a -> b) -> a -> b
$ [[(Name, Dec)]] -> [(Name, Dec)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
      [ [ (Name
patternName, Name -> Kind -> Dec
PatSynSigD Name
patternName ((Kind -> Kind -> Kind) -> Kind -> Cxt -> Kind
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Kind -> Kind -> Kind
AppT (Kind -> Kind -> Kind) -> (Kind -> Kind) -> Kind -> Kind -> Kind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Kind -> Kind -> Kind
AppT Kind
ArrowT) Kind
termType Cxt
types'))
        , (Name
patternName, Name -> PatSynArgs -> PatSynDir -> Pat -> Dec
PatSynD  Name
patternName ([Name] -> PatSynArgs
PrefixPatSyn [Name]
args) PatSynDir
ImplBidir (Name -> Cxt -> [Pat] -> Pat
ConP 'Node [] [Name -> Cxt -> [Pat] -> Pat
ConP Name
conName [] [Pat]
pats]))
        ]
      | Name
conName <- [Name]
conNames
      , let patternName :: Name
patternName = Name -> Name
mkPatternName Name
conName
      ]

  RecGadtC [Name]
conNames [VarBangType]
types Kind
retType -> Kind -> Name -> Name -> Con -> Q [(Name, Dec)]
mkPatternSynonym Kind
signatureType Name
scope Name
term ([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)

  where
    n :: Name
n = String -> Name
mkName String
"n"
    binderT :: Kind
binderT = Name -> Kind
VarT (String -> Name
mkName String
"binder")
    termType :: Kind
termType = Name -> Cxt -> Kind
PeelConT ''AST [Kind
binderT, Kind
signatureType, Name -> Kind
VarT Name
n]
    toArg :: Either ((Name, b), (Name, b)) (Name, b) -> Pat
toArg = \case
      Left ((Name
b, b
_), (Name
x, b
_)) -> Name -> Cxt -> [Pat] -> Pat
ConP 'ScopedAST [] [Name -> Pat
VarP Name
b, Name -> Pat
VarP Name
x]
      Right (Name
x, b
_) -> Name -> Pat
VarP Name
x

    toPatternArgType :: a -> (a, Kind) -> Either ((Name, Kind), (Name, Kind)) (Name, Kind)
toPatternArgType a
i (a
_bang, VarT Name
typeName)
      | Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
scope =
          ((Name, Kind), (Name, Kind))
-> Either ((Name, Kind), (Name, Kind)) (Name, Kind)
forall a b. a -> Either a b
Left
            ( (String -> Name
mkName (String
"b" String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
i), (Kind -> Kind -> Kind) -> Kind -> Cxt -> Kind
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Kind -> Kind -> Kind
AppT Kind
binderT [Name -> Kind
VarT Name
n, Name -> Kind
VarT Name
l])
            , (String -> Name
mkName (String
"x" String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
i), Name -> Cxt -> Kind
PeelConT ''AST [Kind
binderT, Kind
signatureType, Name -> Kind
VarT Name
l]))
      | Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
term =
          (Name, Kind) -> Either ((Name, Kind), (Name, Kind)) (Name, Kind)
forall a b. b -> Either a b
Right (String -> Name
mkName (String
"x" String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
i), Name -> Cxt -> Kind
PeelConT ''AST [Kind
binderT, Kind
signatureType, Name -> Kind
VarT Name
n])
      where
        l :: Name
l = String -> Name
mkName (String
"l" String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
i)
    toPatternArgType a
i (a
_bang, Kind
type_)
      = (Name, Kind) -> Either ((Name, Kind), (Name, Kind)) (Name, Kind)
forall a b. b -> Either a b
Right (String -> Name
mkName (String
"z" String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
i), Kind
type_)

    mkPatternName :: Name -> Name
mkPatternName Name
conName = String -> Name
mkName (Int -> String -> String
forall {a}. Int -> [a] -> [a]
dropEnd (String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
"Sig") (Name -> String
nameBase Name
conName))
    dropEnd :: Int -> [a] -> [a]
dropEnd Int
k = [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> [a]
forall {a}. Int -> [a] -> [a]
drop Int
k ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse

    collapse :: Either (a, a) a -> [a]
collapse = \case
      Left (a
x, a
y) -> [a
x, a
y]
      Right a
x -> [a
x]