-- | 'Builder' for top-level function declaration
module Language.Haskell.TH.Natural.Syntax.Func (
    -- * Builder
    newFunc,
    FuncBuilder,

    -- * State
    FuncBuilderState (..),

    -- * Functions

    -- ** Clause
    addClause,
    bodyFromExp,

    -- ** Signature
    setSignature,

    -- ** Pragmas
    inline,
    setInline,
    addPragma,

    -- * Lenses
    pragmas,
    dec,
    signature,

    -- * Reexports
    module Language.Haskell.TH.Natural.Syntax.Builder.Monad,
) where

import Control.Lens
import Control.Monad.State
import Data.Constructor.Extract
import Data.Maybe (maybeToList)
import Language.Haskell.TH (Q, mkName)
import qualified Language.Haskell.TH as TH
import Language.Haskell.TH.Gen
import Language.Haskell.TH.Natural.Syntax.Builder
import Language.Haskell.TH.Natural.Syntax.Builder.Monad
import Language.Haskell.TH.Syntax.ExtractedCons hiding (fName, inline)

data FuncBuilderState = MkFBS
    { FuncBuilderState -> [Pragma]
_pragmas :: [TH.Pragma]
    , FuncBuilderState -> FunD
_dec :: FunD
    , FuncBuilderState -> Maybe SigD
_signature :: Maybe SigD
    }

makeLenses ''FuncBuilderState

type FuncBuilder = Builder FuncBuilderState

-- TODO Should not be ready if 0 clause

-- | Builds a function declaration. The string argument is the name of the function
newFunc :: String -> FuncBuilder step Ready () -> Q [TH.Dec]
newFunc :: forall (step :: BuilderStep).
String -> FuncBuilder step 'Ready () -> Q [Dec]
newFunc String
fName FuncBuilder step 'Ready ()
builder = do
    MkFBS{..} <- FuncBuilder step 'Ready ()
-> FuncBuilderState -> Q FuncBuilderState
forall {k} (m :: * -> *) s (step :: k) (end :: k).
Monad m =>
BaseBuilder m s step end () -> s -> m s
runBaseBuilder FuncBuilder step 'Ready ()
builder ([Pragma] -> FunD -> Maybe SigD -> FuncBuilderState
MkFBS [] (Name -> [Clause] -> FunD
MkFunD (String -> Name
mkName String
fName) []) Maybe SigD
forall a. Maybe a
Nothing)
    return ((TH.PragmaD <$> _pragmas) ++ (fromEC <$> maybeToList _signature) ++ [fromEC _dec])

-- | Set the signature of the function
setSignature :: (GenType a) => a -> FuncBuilder step step ()
setSignature :: forall {k} a (step :: k).
GenType a =>
a -> FuncBuilder step step ()
setSignature a
sigBuilder = do
    sig <- Q Type -> BaseBuilder Q FuncBuilderState 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 FuncBuilderState step step Type)
-> Q Type -> BaseBuilder Q FuncBuilderState step step Type
forall a b. (a -> b) -> a -> b
$ a -> Q Type
forall a. GenType a => a -> Q Type
genTy a
sigBuilder
    fName <- view (dec . name)
    signature ?= MkSigD fName sig
    return ()

-- | Add a clause to the function
addClause :: Clause -> FuncBuilder step Ready ()
addClause :: forall (step :: BuilderStep). Clause -> FuncBuilder step 'Ready ()
addClause Clause
c = BaseBuilder Q FuncBuilderState step step ()
-> BaseBuilder Q FuncBuilderState step 'Ready ()
forall {k} (m :: * -> *) s (step :: k) a (next :: k).
BaseBuilder m s step step a -> BaseBuilder m s step next a
impure (BaseBuilder Q FuncBuilderState step step ()
 -> BaseBuilder Q FuncBuilderState step 'Ready ())
-> BaseBuilder Q FuncBuilderState step step ()
-> BaseBuilder Q FuncBuilderState step 'Ready ()
forall a b. (a -> b) -> a -> b
$ ((FunD -> Identity FunD)
-> FuncBuilderState -> Identity FuncBuilderState
Lens' FuncBuilderState FunD
dec ((FunD -> Identity FunD)
 -> FuncBuilderState -> Identity FuncBuilderState)
-> (([Clause] -> Identity [Clause]) -> FunD -> Identity FunD)
-> ([Clause] -> Identity [Clause])
-> FuncBuilderState
-> Identity FuncBuilderState
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Clause] -> Identity [Clause]) -> FunD -> Identity FunD
forall a b. HasClauses a b => Lens' a b
Lens' FunD [Clause]
clauses) (([Clause] -> Identity [Clause])
 -> FuncBuilderState -> Identity FuncBuilderState)
-> Clause -> BaseBuilder Q FuncBuilderState step step ()
forall s (m :: * -> *) b a.
(MonadState s m, Snoc b b a a) =>
ASetter s s b b -> a -> m ()
|>= Clause
c

-- | Uses an Exp as the body of a function
--
-- Warning: This operation is destructive, and replaces all previous clauses set using 'addClause'
bodyFromExp :: (GenExpr b) => b -> FuncBuilder step Ready ()
bodyFromExp :: forall b (step :: BuilderStep).
GenExpr b =>
b -> FuncBuilder step 'Ready ()
bodyFromExp b
qe = BaseBuilder Q FuncBuilderState step step ()
-> BaseBuilder Q FuncBuilderState step 'Ready ()
forall {k} (m :: * -> *) s (step :: k) a (next :: k).
BaseBuilder m s step step a -> BaseBuilder m s step next a
impure (BaseBuilder Q FuncBuilderState step step ()
 -> BaseBuilder Q FuncBuilderState step 'Ready ())
-> BaseBuilder Q FuncBuilderState step step ()
-> BaseBuilder Q FuncBuilderState step 'Ready ()
forall a b. (a -> b) -> a -> b
$ do
    e <- Q Exp -> BaseBuilder Q FuncBuilderState step step Exp
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Exp -> BaseBuilder Q FuncBuilderState step step Exp)
-> Q Exp -> BaseBuilder Q FuncBuilderState step step Exp
forall a b. (a -> b) -> a -> b
$ b -> Q Exp
forall a. GenExpr a => a -> Q Exp
genExpr b
qe
    (dec . clauses) .= [TH.Clause [] (TH.NormalB e) []]

-- | Add an inline pragma to the function
inline :: FuncBuilder step step ()
inline :: forall {k} (step :: k). FuncBuilder step step ()
inline = Inline -> RuleMatch -> Phases -> FuncBuilder step step ()
forall {k} (step :: k).
Inline -> RuleMatch -> Phases -> FuncBuilder step step ()
setInline Inline
TH.Inline RuleMatch
TH.FunLike Phases
TH.AllPhases

-- | Sets an inline pragma to the function
setInline :: TH.Inline -> TH.RuleMatch -> TH.Phases -> FuncBuilder step step ()
setInline :: forall {k} (step :: k).
Inline -> RuleMatch -> Phases -> FuncBuilder step step ()
setInline Inline
i RuleMatch
rm Phases
phs = do
    fName <- Getting Name FuncBuilderState Name
-> BaseBuilder Q FuncBuilderState step step Name
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view ((FunD -> Const Name FunD)
-> FuncBuilderState -> Const Name FuncBuilderState
Lens' FuncBuilderState FunD
dec ((FunD -> Const Name FunD)
 -> FuncBuilderState -> Const Name FuncBuilderState)
-> ((Name -> Const Name Name) -> FunD -> Const Name FunD)
-> Getting Name FuncBuilderState Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name -> Const Name Name) -> FunD -> Const Name FunD
forall a b. HasName a b => Lens' a b
Lens' FunD Name
name)
    modify $ over pragmas $ filter $ \case
        TH.InlineP{} -> Bool
False
        Pragma
_ -> Bool
True
    let newInlineP = Name -> Inline -> RuleMatch -> Phases -> Pragma
TH.InlineP Name
fName Inline
i RuleMatch
rm Phases
phs
    pragmas <|= newInlineP

-- | Add a 'Pragma' alongside the function declaration
addPragma :: TH.Pragma -> FuncBuilder step step ()
addPragma :: forall {k} (step :: k). Pragma -> FuncBuilder step step ()
addPragma Pragma
p = ASetter FuncBuilderState FuncBuilderState [Pragma] [Pragma]
Lens' FuncBuilderState [Pragma]
pragmas ASetter FuncBuilderState FuncBuilderState [Pragma] [Pragma]
-> Pragma -> BaseBuilder Q FuncBuilderState step step ()
forall s (m :: * -> *) b a.
(MonadState s m, Cons b b a a) =>
ASetter s s b b -> a -> m ()
<|= Pragma
p