{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE TemplateHaskell #-}
module Control.Monad.Free.Foil.TH.Signature where

import           Language.Haskell.TH
import           Language.Haskell.TH.Syntax

import           Control.Monad              (forM_)
import           Control.Monad.Foil.TH.Util
import           Data.Maybe                 (catMaybes)

-- | Generate a signature for the free foil (or free scoped monads)
-- based on a naïve recursive abstract syntax representation,
-- with clearly separated types for terms, variable identifiers, scoped terms,
-- and patterns (binders).
mkSignature
  :: Name -- ^ Type name for raw terms.
  -> Name -- ^ Type name for raw variable identifiers.
  -> Name -- ^ Type name for raw scoped terms.
  -> Name -- ^ Type name for raw patterns.
  -> Q [Dec]
mkSignature :: Name -> Name -> Name -> Name -> Q [Dec]
mkSignature Name
termT Name
nameT Name
scopeT Name
patternT = do
  Name
scope <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"scope"
  Name
term <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"term"
  TyConI (DataD Cxt
_ctx Name
_name [TyVarBndr BndrVis]
termTVars Maybe Kind
_kind [Con]
termCons [DerivClause]
_deriv) <- Name -> Q Info
reify Name
termT

  [Con]
signatureCons <- [Maybe Con] -> [Con]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Con] -> [Con]) -> Q [Maybe Con] -> Q [Con]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Con -> Q (Maybe Con)) -> [Con] -> Q [Maybe Con]
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 (Maybe Con)
toSignatureCons Name
scope Name
term) [Con]
termCons

  Q () -> Q ()
addModFinalizer (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ DocLoc -> String -> Q ()
putDoc (Name -> DocLoc
DeclDoc Name
signatureT)
    (String
"/Generated/ with '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show 'mkSignature String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'. A signature bifunctor, specifying the nodes of a syntax tree corresponding to '" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
termT String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"'.")
  [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return
    [ Cxt
-> Name
-> [TyVarBndr BndrVis]
-> Maybe Kind
-> [Con]
-> [DerivClause]
-> Dec
DataD [] Name
signatureT ([TyVarBndr BndrVis]
termTVars [TyVarBndr BndrVis] -> [TyVarBndr BndrVis] -> [TyVarBndr BndrVis]
forall a. [a] -> [a] -> [a]
++ [Name -> BndrVis -> TyVarBndr BndrVis
forall flag. Name -> flag -> TyVarBndr flag
PlainTV Name
scope BndrVis
BndrReq, Name -> BndrVis -> TyVarBndr BndrVis
forall flag. Name -> flag -> TyVarBndr flag
PlainTV Name
term BndrVis
BndrReq]) Maybe Kind
forall a. Maybe a
Nothing [Con]
signatureCons
      [Maybe DerivStrategy -> Cxt -> DerivClause
DerivClause Maybe DerivStrategy
forall a. Maybe a
Nothing [Name -> Kind
ConT ''Functor, Name -> Kind
ConT ''Foldable, Name -> Kind
ConT ''Traversable]]
    ]
  where
    signatureT :: Name
signatureT = String -> Name
mkName (Name -> String
nameBase Name
termT String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"Sig")

    toSignatureCons :: Name -> Name -> Con -> Q (Maybe Con)
    toSignatureCons :: Name -> Name -> Con -> Q (Maybe Con)
toSignatureCons Name
scope Name
term Con
con' = case Con
con' of
      -- treat constructors with a single variable field as variable constructor and ignore
      NormalC Name
_conName [BangType]
types | [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or [ Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
nameT | (Bang
_bang, PeelConT Name
typeName Cxt
_typeParams) <- [BangType]
types ]
        -> Maybe Con -> Q (Maybe Con)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Con
forall a. Maybe a
Nothing
      RecC Name
_conName [VarBangType]
types | [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or [ Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
nameT | (Name
_name, Bang
_bang, PeelConT Name
typeName Cxt
_typeParams) <- [VarBangType]
types ]
        -> Maybe Con -> Q (Maybe Con)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Con
forall a. Maybe a
Nothing

      NormalC Name
conName [BangType]
params -> do
        Q () -> Q ()
addModFinalizer (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ DocLoc -> String -> Q ()
putDoc (Name -> DocLoc
DeclDoc Name
conName') (String
"Corresponds to '" 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
"'.")
        Con -> Maybe Con
forall a. a -> Maybe a
Just (Con -> Maybe Con)
-> ([Maybe BangType] -> Con) -> [Maybe BangType] -> Maybe Con
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [BangType] -> Con
NormalC Name
conName' ([BangType] -> Con)
-> ([Maybe BangType] -> [BangType]) -> [Maybe BangType] -> Con
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe BangType] -> [BangType]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe BangType] -> Maybe Con)
-> Q [Maybe BangType] -> Q (Maybe Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BangType -> Q (Maybe BangType))
-> [BangType] -> Q [Maybe BangType]
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 BangType -> Q (Maybe BangType)
forall {m :: * -> *} {a}.
MonadFail m =>
(a, Kind) -> m (Maybe (a, Kind))
toSignatureParam [BangType]
params
        where
          conName' :: Name
conName' = String -> Name
mkName (Name -> String
nameBase Name
conName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"Sig")
      RecC Name
conName [VarBangType]
params -> do
        Q () -> Q ()
addModFinalizer (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ DocLoc -> String -> Q ()
putDoc (Name -> DocLoc
DeclDoc Name
conName') (String
"Corresponds to '" 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
"'.")
        Con -> Maybe Con
forall a. a -> Maybe a
Just (Con -> Maybe Con)
-> ([Maybe VarBangType] -> Con) -> [Maybe VarBangType] -> Maybe Con
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [VarBangType] -> Con
RecC Name
conName' ([VarBangType] -> Con)
-> ([Maybe VarBangType] -> [VarBangType])
-> [Maybe VarBangType]
-> Con
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe VarBangType] -> [VarBangType]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe VarBangType] -> Maybe Con)
-> Q [Maybe VarBangType] -> Q (Maybe Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VarBangType -> Q (Maybe VarBangType))
-> [VarBangType] -> Q [Maybe VarBangType]
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 VarBangType -> Q (Maybe VarBangType)
forall {f :: * -> *} {a} {b}.
MonadFail f =>
(a, b, Kind) -> f (Maybe (a, b, Kind))
toSignatureParam' [VarBangType]
params
        where
          conName' :: Name
conName' = String -> Name
mkName (Name -> String
nameBase Name
conName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"Sig")
      InfixC BangType
l Name
conName BangType
r -> do
        Q () -> Q ()
addModFinalizer (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ DocLoc -> String -> Q ()
putDoc (Name -> DocLoc
DeclDoc Name
conName') (String
"Corresponds to '" 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
"'.")
        Con -> Maybe Con
forall a. a -> Maybe a
Just (Con -> Maybe Con) -> Q Con -> Q (Maybe Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((BangType -> Name -> BangType -> Con)
-> Name -> BangType -> BangType -> Con
forall a b c. (a -> b -> c) -> b -> a -> c
flip BangType -> Name -> BangType -> Con
InfixC Name
conName' (BangType -> BangType -> Con) -> Q BangType -> Q (BangType -> Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BangType -> Q BangType
forall {m :: * -> *} {a}. MonadFail m => (a, Kind) -> m (a, Kind)
toInfixParam BangType
l Q (BangType -> Con) -> Q BangType -> Q Con
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BangType -> Q BangType
forall {m :: * -> *} {a}. MonadFail m => (a, Kind) -> m (a, Kind)
toInfixParam BangType
r)
        where
          conName' :: Name
conName' = String -> Name
mkName (Name -> String
nameBase Name
conName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"---")
      ForallC [TyVarBndr Specificity]
params Cxt
ctx Con
con -> (Con -> Con) -> Maybe Con -> Maybe Con
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([TyVarBndr Specificity] -> Cxt -> Con -> Con
ForallC [TyVarBndr Specificity]
params Cxt
ctx) (Maybe Con -> Maybe Con) -> Q (Maybe Con) -> Q (Maybe Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Name -> Con -> Q (Maybe Con)
toSignatureCons Name
scope Name
term Con
con
      GadtC [Name]
conNames [BangType]
argTypes Kind
retType -> do
        let conNames' :: [Name]
conNames' = (Name -> Name) -> [Name] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (\Name
conName -> String -> Name
mkName (Name -> String
nameBase Name
conName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"---")) [Name]
conNames
        [(Name, Name)] -> ((Name, Name) -> Q ()) -> Q ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Name] -> [Name] -> [(Name, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
conNames [Name]
conNames') (((Name, Name) -> Q ()) -> Q ()) -> ((Name, Name) -> Q ()) -> Q ()
forall a b. (a -> b) -> a -> b
$ \(Name
conName, Name
conName') ->
          Q () -> Q ()
addModFinalizer (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ DocLoc -> String -> Q ()
putDoc (Name -> DocLoc
DeclDoc Name
conName') (String
"Corresponds to '" 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
"'.")
        Con -> Maybe Con
forall a. a -> Maybe a
Just (Con -> Maybe Con) -> Q Con -> Q (Maybe Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Name] -> [BangType] -> Kind -> Con
GadtC [Name]
conNames' ([BangType] -> Kind -> Con) -> Q [BangType] -> Q (Kind -> Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Maybe BangType] -> [BangType]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe BangType] -> [BangType])
-> Q [Maybe BangType] -> Q [BangType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BangType -> Q (Maybe BangType))
-> [BangType] -> Q [Maybe BangType]
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 BangType -> Q (Maybe BangType)
forall {m :: * -> *} {a}.
MonadFail m =>
(a, Kind) -> m (Maybe (a, Kind))
toSignatureParam [BangType]
argTypes) Q (Kind -> Con) -> Q Kind -> Q Con
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Q Kind
retType')
        where
          retType' :: Q Kind
retType' = case Kind
retType of
            PeelConT Name
typeName Cxt
typeParams
              | Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
termT -> Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Cxt -> Kind
PeelConT Name
signatureT (Cxt
typeParams Cxt -> Cxt -> Cxt
forall a. [a] -> [a] -> [a]
++ [Name -> Kind
VarT Name
scope, Name -> Kind
VarT Name
term]))
            Kind
_ -> String -> Q Kind
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"unexpected return type in a GADT constructor"
      RecGadtC [Name]
conNames [VarBangType]
argTypes Kind
retType -> do
        let conNames' :: [Name]
conNames' = (Name -> Name) -> [Name] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (\Name
conName -> String -> Name
mkName (Name -> String
nameBase Name
conName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"---")) [Name]
conNames
        [(Name, Name)] -> ((Name, Name) -> Q ()) -> Q ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Name] -> [Name] -> [(Name, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
conNames [Name]
conNames') (((Name, Name) -> Q ()) -> Q ()) -> ((Name, Name) -> Q ()) -> Q ()
forall a b. (a -> b) -> a -> b
$ \(Name
conName, Name
conName') ->
          Q () -> Q ()
addModFinalizer (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ DocLoc -> String -> Q ()
putDoc (Name -> DocLoc
DeclDoc Name
conName') (String
"Corresponds to '" 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
"'.")
        Con -> Maybe Con
forall a. a -> Maybe a
Just (Con -> Maybe Con) -> Q Con -> Q (Maybe Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Name] -> [VarBangType] -> Kind -> Con
RecGadtC [Name]
conNames' ([VarBangType] -> Kind -> Con)
-> Q [VarBangType] -> Q (Kind -> Con)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Maybe VarBangType] -> [VarBangType]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe VarBangType] -> [VarBangType])
-> Q [Maybe VarBangType] -> Q [VarBangType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (VarBangType -> Q (Maybe VarBangType))
-> [VarBangType] -> Q [Maybe VarBangType]
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 VarBangType -> Q (Maybe VarBangType)
forall {f :: * -> *} {a} {b}.
MonadFail f =>
(a, b, Kind) -> f (Maybe (a, b, Kind))
toSignatureParam' [VarBangType]
argTypes) Q (Kind -> Con) -> Q Kind -> Q Con
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Q Kind
retType')
        where
          retType' :: Q Kind
retType' = case Kind
retType of
            PeelConT Name
typeName Cxt
typeParams
              | Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
termT -> Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Cxt -> Kind
PeelConT Name
signatureT (Cxt
typeParams Cxt -> Cxt -> Cxt
forall a. [a] -> [a] -> [a]
++ [Name -> Kind
VarT Name
scope, Name -> Kind
VarT Name
term]))
            Kind
_ -> String -> Q Kind
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"unexpected return type in a GADT constructor"

      where
        toInfixParam :: (a, Kind) -> m (a, Kind)
toInfixParam (a
bang_, Kind
type_) = (a, Kind) -> m (Maybe (a, Kind))
forall {m :: * -> *} {a}.
MonadFail m =>
(a, Kind) -> m (Maybe (a, Kind))
toSignatureParam (a
bang_, Kind
type_) m (Maybe (a, Kind))
-> (Maybe (a, Kind) -> m (a, Kind)) -> m (a, Kind)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe (a, Kind)
Nothing -> (a, Kind) -> m (a, Kind)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
bang_, Name -> Kind
VarT ''())
          Just (a, Kind)
bt -> (a, Kind) -> m (a, Kind)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a, Kind)
bt

        toSignatureParam' :: (a, b, Kind) -> f (Maybe (a, b, Kind))
toSignatureParam' (a
name, b
bang_, Kind
type_) = ((b, Kind) -> (a, b, Kind))
-> Maybe (b, Kind) -> Maybe (a, b, Kind)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (b, Kind) -> (a, b, Kind)
forall {b} {c}. (b, c) -> (a, b, c)
k (Maybe (b, Kind) -> Maybe (a, b, Kind))
-> f (Maybe (b, Kind)) -> f (Maybe (a, b, Kind))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (b, Kind) -> f (Maybe (b, Kind))
forall {m :: * -> *} {a}.
MonadFail m =>
(a, Kind) -> m (Maybe (a, Kind))
toSignatureParam (b
bang_, Kind
type_)
          where
            k :: (b, c) -> (a, b, c)
k (b
x, c
y) = (a
name, b
x, c
y)

        toSignatureParam :: (a, Kind) -> m (Maybe (a, Kind))
toSignatureParam (a
bang_, PeelConT Name
typeName Cxt
_typeParams)
          | Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
nameT = String -> m (Maybe (a, Kind))
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"variable with other stuff in constructor: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Con -> String
forall a. Show a => a -> String
show Con
con')
          | Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
patternT = Maybe (a, Kind) -> m (Maybe (a, Kind))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (a, Kind)
forall a. Maybe a
Nothing -- skip binders, they will be inserted automatically with each scoped term
          | Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
scopeT = Maybe (a, Kind) -> m (Maybe (a, Kind))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a, Kind) -> Maybe (a, Kind)
forall a. a -> Maybe a
Just (a
bang_, Name -> Kind
VarT Name
scope))
          | Name
typeName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
termT = Maybe (a, Kind) -> m (Maybe (a, Kind))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a, Kind) -> Maybe (a, Kind)
forall a. a -> Maybe a
Just (a
bang_, Name -> Kind
VarT Name
term))
        toSignatureParam (a, Kind)
bt = Maybe (a, Kind) -> m (Maybe (a, Kind))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a, Kind) -> Maybe (a, Kind)
forall a. a -> Maybe a
Just (a, Kind)
bt)  -- everything else remains as is