{-# LANGUAGE QualifiedDo #-}

-- | 'Builder' for type signatures
--
-- In Template Haskell, a signature ('SigD') refers to a symbol's type signature with a 'Name' (e.g. _length :: [a] -> Int_). Here, by _signature_, we mean just the type signature, without a name.
module Language.Haskell.TH.Natural.Syntax.Signature (
    -- * Builder
    newSignature,
    SignatureBuilder,

    -- * State
    SignatureState (..),

    -- * Functions
    addToForall,
    addConstraint,
    addParam,
    setResultType,

    -- * Lenses
    tyVarBndr,
    constraints,
    params,
    result,

    -- * Re-export
    newTypeVar,
    module Language.Haskell.TH.Natural.Syntax.Builder.Monad,
) where

import Control.Lens ((?=), (^.), (|>=))
import Control.Lens.TH
import Data.Constructor.Extract
import Language.Haskell.TH (Q, Type (AppT, ArrowT))
import qualified Language.Haskell.TH as TH
import Language.Haskell.TH.Gen
import Language.Haskell.TH.Natural.Syntax.Builder
import qualified Language.Haskell.TH.Natural.Syntax.Builder as B
import Language.Haskell.TH.Natural.Syntax.Builder.Monad
import Language.Haskell.TH.Natural.Syntax.Name
import Language.Haskell.TH.Syntax.ExtractedCons hiding (inline, tyVarBndr)

type SignatureBuilder prev next a = Builder SignatureState prev next a

data SignatureState = MkSBS
    { SignatureState -> [TyVarBndr Specificity]
_tyVarBndr :: [TH.TyVarBndr TH.Specificity]
    , SignatureState -> [Type]
_constraints :: [TH.Type]
    , SignatureState -> [Type]
_params :: [TH.Type]
    , SignatureState -> Maybe Type
_result :: Maybe TH.Type
    }

makeLenses ''SignatureState

instance GenType (SignatureBuilder step Ready ()) where
    genTy :: SignatureBuilder step 'Ready () -> Q Type
genTy = (ForallT -> Type) -> Q ForallT -> Q Type
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ForallT -> Type
forall con ty. ExtractedConstructor con ty => con -> ty
fromEC (Q ForallT -> Q Type)
-> (SignatureBuilder step 'Ready () -> Q ForallT)
-> SignatureBuilder step 'Ready ()
-> Q Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SignatureBuilder step 'Ready () -> Q ForallT
forall (step :: BuilderStep).
SignatureBuilder step 'Ready () -> Q ForallT
newSignature

-- | Builds a type signature
newSignature :: SignatureBuilder step Ready () -> Q ForallT
newSignature :: forall (step :: BuilderStep).
SignatureBuilder step 'Ready () -> Q ForallT
newSignature SignatureBuilder step 'Ready ()
builder = do
    MkSBS{..} <- SignatureBuilder step 'Ready ()
-> SignatureState -> Q SignatureState
forall {k} (m :: * -> *) s (step :: k) (end :: k).
Monad m =>
BaseBuilder m s step end () -> s -> m s
runBaseBuilder SignatureBuilder step 'Ready ()
builder ([TyVarBndr Specificity]
-> [Type] -> [Type] -> Maybe Type -> SignatureState
MkSBS [] [] [] Maybe Type
forall a. Maybe a
Nothing)
    resTy <- case _result of
        Maybe Type
Nothing -> String -> Q Type
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
Prelude.fail String
"The signature does not contain a return type."
        Just Type
r -> Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
r
    let funcType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Type
param -> ((Type
ArrowT Type -> Type -> Type
`AppT` Type
param) Type -> Type -> Type
`AppT`)) Type
resTy [Type]
_params
    return $ MkForallT _tyVarBndr _constraints funcType

-- | Adds the given type variable to the _forall_ list.
--
-- Using this function should comply with the 'forall-or-nothing' rule (https://ghc.gitlab.haskell.org/ghc/doc/users_guide/exts/explicit_forall.html#the-forall-or-nothing-rule)
addToForall :: TypeVarName -> SignatureBuilder step step ()
addToForall :: forall {k} (step :: k).
TypeVarName -> SignatureBuilder step step ()
addToForall TypeVarName
tyVar = StateT SignatureState Q ()
-> BaseBuilder Q SignatureState step step ()
forall {k} s (m :: * -> *) a (prev :: k) (curr :: k).
StateT s m a -> BaseBuilder m s prev curr a
unsafeWithState (StateT SignatureState Q ()
 -> BaseBuilder Q SignatureState step step ())
-> StateT SignatureState Q ()
-> BaseBuilder Q SignatureState step step ()
forall a b. (a -> b) -> a -> b
$ ([TyVarBndr Specificity] -> Identity [TyVarBndr Specificity])
-> SignatureState -> Identity SignatureState
Lens' SignatureState [TyVarBndr Specificity]
tyVarBndr (([TyVarBndr Specificity] -> Identity [TyVarBndr Specificity])
 -> SignatureState -> Identity SignatureState)
-> TyVarBndr Specificity -> StateT SignatureState Q ()
forall s (m :: * -> *) b a.
(MonadState s m, Snoc b b a a) =>
ASetter s s b b -> a -> m ()
|>= Name -> Specificity -> TyVarBndr Specificity
forall flag. Name -> flag -> TyVarBndr flag
TH.PlainTV (TypeVarName
tyVar TypeVarName -> Getting Name TypeVarName Name -> Name
forall s a. s -> Getting a s a -> a
^. Getting Name TypeVarName Name
forall a b. HasName a b => Lens' a b
Lens' TypeVarName Name
name) Specificity
TH.SpecifiedSpec

-- | Add the given type to the set of constraints
addConstraint :: (GenType a) => a -> SignatureBuilder step step ()
addConstraint :: forall {k} a (step :: k).
GenType a =>
a -> SignatureBuilder step step ()
addConstraint a
tyBuilder = do
    constr <- Q Type -> BaseBuilder Q SignatureState step step Type
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Type -> BaseBuilder Q SignatureState step step Type)
-> Q Type -> BaseBuilder Q SignatureState step step Type
forall a b. (a -> b) -> a -> b
$ a -> Q Type
forall a. GenType a => a -> Q Type
genTy a
tyBuilder
    unsafeWithState $
        constraints |>= constr

-- | Set the type as the nth parameter of the function's signature
--
-- (n being the number of time 'addParam' was called)
addParam :: (GenType a) => a -> SignatureBuilder step step ()
addParam :: forall {k} a (step :: k).
GenType a =>
a -> SignatureBuilder step step ()
addParam a
tyBuilder = do
    param <- Q Type -> BaseBuilder Q SignatureState step step Type
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Type -> BaseBuilder Q SignatureState step step Type)
-> Q Type -> BaseBuilder Q SignatureState step step Type
forall a b. (a -> b) -> a -> b
$ a -> Q Type
forall a. GenType a => a -> Q Type
genTy a
tyBuilder
    unsafeWithState $
        params |>= param

-- | Set the result type in the function's signature
setResultType :: (GenType a) => a -> SignatureBuilder step Ready ()
setResultType :: forall a (step :: BuilderStep).
GenType a =>
a -> SignatureBuilder step 'Ready ()
setResultType a
tyBuilder = B.do
    resTy <- Q Type -> BaseBuilder Q SignatureState step step Type
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Type -> BaseBuilder Q SignatureState step step Type)
-> Q Type -> BaseBuilder Q SignatureState step step Type
forall a b. (a -> b) -> a -> b
$ a -> Q Type
forall a. GenType a => a -> Q Type
genTy a
tyBuilder
    unsafeWithState $ result ?= resTy