{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

-- Originally ported from code written by Sandy Maguire (@isovector), available
-- at https://github.com/IxpertaSolutions/freer-effects/pull/28.

{-|
This module provides Template Haskell functions for automatically generating
effect operation functions (that is, functions that use 'send') from a given
effect algebra. For example, using the @FileSystem@ effect from the example in
the module documentation for "Control.Monad.Freer", we can write the following:

@
data FileSystem r where
  ReadFile :: 'FilePath' -> FileSystem 'String'
  WriteFile :: 'FilePath' -> 'String' -> FileSystem ()
'makeEffect' ''FileSystem
@

This will automatically generate the following functions:

@
readFile :: 'Member' FileSystem effs => 'FilePath' -> 'Eff' effs 'String'
readFile a = 'send' (ReadFile a)

writeFile :: 'Member' FileSystem effs => 'FilePath' -> 'String' -> 'Eff' effs ()
writeFile a b = 'send' (WriteFile a b)
@
-}
module Control.Monad.Freer.TH
  ( makeEffect
  , makeEffect_
  )
where

import Control.Monad (forM, unless)
import Control.Monad.Freer (send, Member, Eff)
import Data.Char (toLower)
import Language.Haskell.TH
import Prelude


-- | If @T@ is a GADT representing an effect algebra, as described in the module
-- documentation for "Control.Monad.Freer", @$('makeEffect' ''T)@ automatically
-- generates a function that uses 'send' with each operation. For more
-- information, see the module documentation for "Control.Monad.Freer.TH".
makeEffect :: Name -> Q [Dec]
makeEffect :: Name -> Q [Dec]
makeEffect = Bool -> Name -> Q [Dec]
genFreer Bool
True

-- | Like 'makeEffect', but does not provide type signatures. This can be used
-- to attach Haddock comments to individual arguments for each generated
-- function.
--
-- @
-- data Lang x where
--   Output :: String -> Lang ()
--
-- makeEffect_ ''Lang
--
-- -- | Output a string.
-- output :: Member Lang effs
--        => String    -- ^ String to output.
--        -> Eff effs ()  -- ^ No result.
-- @
--
-- Note that 'makeEffect_' must be used /before/ the explicit type signatures.
makeEffect_ :: Name -> Q [Dec]
makeEffect_ :: Name -> Q [Dec]
makeEffect_ = Bool -> Name -> Q [Dec]
genFreer Bool
False

-- | Generates declarations and possibly signatures for functions to lift GADT
-- constructors into 'Eff' actions.
genFreer :: Bool -> Name -> Q [Dec]
genFreer :: Bool -> Name -> Q [Dec]
genFreer Bool
makeSigs Name
tcName = do
  -- The signatures for the generated definitions require FlexibleContexts.
  Extension -> Q Bool
isExtEnabled Extension
FlexibleContexts
    Q Bool -> (Bool -> Q ()) -> Q ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Bool -> Q () -> Q ()) -> Q () -> Bool -> Q ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"makeEffect requires FlexibleContexts to be enabled")

  Name -> Q Info
reify Name
tcName Q Info -> (Info -> Q [Dec]) -> Q [Dec]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    TyConI (DataD Cxt
_ Name
_ [TyVarBndr]
_ Maybe Kind
_ [Con]
cons [DerivClause]
_) -> do
      [Dec]
sigs <- (Dec -> Bool) -> [Dec] -> [Dec]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Dec -> Bool
forall a b. a -> b -> a
const Bool
makeSigs) ([Dec] -> [Dec]) -> Q [Dec] -> Q [Dec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Con -> Q Dec) -> [Con] -> Q [Dec]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Con -> Q Dec
genSig [Con]
cons
      [Dec]
decs <- (Con -> Q Dec) -> [Con] -> Q [Dec]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Con -> Q Dec
genDecl [Con]
cons
      [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> Q [Dec]) -> [Dec] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ [Dec]
sigs [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [Dec]
decs

    Info
_ -> String -> Q [Dec]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"makeEffect expects a type constructor"

-- | Given the name of a GADT constructor, return the name of the corresponding
-- lifted function.
getDeclName :: Name -> Name
getDeclName :: Name -> Name
getDeclName = String -> Name
mkName (String -> Name) -> (Name -> String) -> Name -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Char) -> String -> String
forall a. (a -> a) -> [a] -> [a]
overFirst Char -> Char
toLower (String -> String) -> (Name -> String) -> Name -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase
 where
  overFirst :: (a -> a) -> [a] -> [a]
overFirst a -> a
f (a
a : [a]
as) = a -> a
f a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
as
  overFirst a -> a
_ [a]
as       = [a]
as

-- | Builds a function definition of the form @x a b c = send $ X a b c@.
genDecl :: Con -> Q Dec
genDecl :: Con -> Q Dec
genDecl (ForallC [TyVarBndr]
_       Cxt
_     Con
con) = Con -> Q Dec
genDecl Con
con
genDecl (GadtC   [Name
cName] [BangType]
tArgs Kind
_  ) = do
  let fnName :: Name
fnName = Name -> Name
getDeclName Name
cName
  let arity :: Int
arity  = [BangType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [BangType]
tArgs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  [Name]
dTypeVars <- [Int] -> (Int -> Q Name) -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
0 .. Int
arity] ((Int -> Q Name) -> Q [Name]) -> (Int -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ Q Name -> Int -> Q Name
forall a b. a -> b -> a
const (Q Name -> Int -> Q Name) -> Q Name -> Int -> Q Name
forall a b. (a -> b) -> a -> b
$ String -> Q Name
newName String
"a"
  Dec -> Q Dec
forall (m :: * -> *) a. Monad m => a -> m a
return (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$ Name -> [Clause] -> Dec
FunD Name
fnName ([Clause] -> Dec) -> (Clause -> [Clause]) -> Clause -> Dec
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Clause -> [Clause]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Clause -> Dec) -> Clause -> Dec
forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause
    (Name -> Pat
VarP (Name -> Pat) -> [Name] -> [Pat]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
dTypeVars)
    (Exp -> Body
NormalB (Exp -> Body) -> (Exp -> Exp) -> Exp -> Body
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'send) (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
      (\Exp
b -> Exp -> Exp -> Exp
AppE Exp
b (Exp -> Exp) -> (Name -> Exp) -> Name -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp
VarE)
      (Name -> Exp
ConE Name
cName)
      [Name]
dTypeVars
    )
    []
genDecl Con
_ = String -> Q Dec
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"genDecl expects a GADT constructor"

-- | Generates a function type from the corresponding GADT type constructor
-- @x :: Member (Effect e) effs => a -> b -> c -> Eff effs r@.
genType :: Con -> Q Type
genType :: Con -> Q Kind
genType (ForallC [TyVarBndr]
tyVarBindings Cxt
conCtx Con
con)
  = [TyVarBndr] -> Cxt -> Kind -> Kind
ForallT [TyVarBndr]
tyVarBindings Cxt
conCtx (Kind -> Kind) -> Q Kind -> Q Kind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Con -> Q Kind
genType Con
con
genType (GadtC   [Name]
_ [BangType]
tArgs' (AppT Kind
eff Kind
tRet)) = do
  Name
effs <- String -> Q Name
newName String
"effs"
  let
    tArgs :: Cxt
tArgs            = (BangType -> Kind) -> [BangType] -> Cxt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap BangType -> Kind
forall a b. (a, b) -> b
snd [BangType]
tArgs'
    memberConstraint :: Kind
memberConstraint = Name -> Kind
ConT ''Member Kind -> Kind -> Kind
`AppT` Kind
eff Kind -> Kind -> Kind
`AppT` Name -> Kind
VarT Name
effs
    resultType :: Kind
resultType       = Name -> Kind
ConT ''Eff Kind -> Kind -> Kind
`AppT` Name -> Kind
VarT Name
effs Kind -> Kind -> Kind
`AppT` Kind
tRet

  Kind -> Q Kind
forall (m :: * -> *) a. Monad m => a -> m a
return
#if MIN_VERSION_template_haskell(2,17,0)
    .  ForallT [PlainTV effs SpecifiedSpec] [memberConstraint]
#else
    (Kind -> Q Kind) -> (Cxt -> Kind) -> Cxt -> Q Kind
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  [TyVarBndr] -> Cxt -> Kind -> Kind
ForallT [Name -> TyVarBndr
PlainTV Name
effs] [Kind
memberConstraint]
#endif
    (Kind -> Kind) -> (Cxt -> Kind) -> Cxt -> Kind
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  Cxt -> Kind
foldArrows
    (Cxt -> Q Kind) -> Cxt -> Q Kind
forall a b. (a -> b) -> a -> b
$  Cxt
tArgs
    Cxt -> Cxt -> Cxt
forall a. [a] -> [a] -> [a]
++ [Kind
resultType]
-- TODO: Although this should never happen, we obviously need a better error message below.
genType Con
_       = String -> Q Kind
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"genSig expects a GADT constructor"

-- | Turn all (KindedTV tv StarT) into (PlainTV tv) in the given type
-- This can prevent the need for KindSignatures
simplifyBndrs :: Type -> Type
simplifyBndrs :: Kind -> Kind
simplifyBndrs (ForallT [TyVarBndr]
bndrs Cxt
tcxt Kind
t) = [TyVarBndr] -> Cxt -> Kind -> Kind
ForallT ((TyVarBndr -> TyVarBndr) -> [TyVarBndr] -> [TyVarBndr]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr -> TyVarBndr
simplifyBndr [TyVarBndr]
bndrs) Cxt
tcxt (Kind -> Kind
simplifyBndrs Kind
t)
simplifyBndrs (AppT Kind
t1 Kind
t2) = Kind -> Kind -> Kind
AppT (Kind -> Kind
simplifyBndrs Kind
t1) (Kind -> Kind
simplifyBndrs Kind
t2)
simplifyBndrs (SigT Kind
t Kind
k) = Kind -> Kind -> Kind
SigT (Kind -> Kind
simplifyBndrs Kind
t) Kind
k
simplifyBndrs (InfixT Kind
t1 Name
n Kind
t2) = Kind -> Name -> Kind -> Kind
InfixT (Kind -> Kind
simplifyBndrs Kind
t1) Name
n (Kind -> Kind
simplifyBndrs Kind
t2)
simplifyBndrs (UInfixT Kind
t1 Name
n Kind
t2) = Kind -> Name -> Kind -> Kind
InfixT (Kind -> Kind
simplifyBndrs Kind
t1) Name
n (Kind -> Kind
simplifyBndrs Kind
t2)
simplifyBndrs (ParensT Kind
t) = Kind -> Kind
ParensT (Kind -> Kind
simplifyBndrs Kind
t)
simplifyBndrs Kind
t = Kind
t

-- | Turn TvVarBndrs of the form (KindedTV tv StarT) into (PlainTV tv)
-- This can prevent the need for KindSignatures
#if MIN_VERSION_template_haskell(2,17,0)
simplifyBndr :: TyVarBndrSpec -> TyVarBndrSpec
simplifyBndr (KindedTV tv f StarT) = PlainTV tv f
#else
simplifyBndr :: TyVarBndr -> TyVarBndr
simplifyBndr :: TyVarBndr -> TyVarBndr
simplifyBndr (KindedTV Name
tv Kind
StarT) = Name -> TyVarBndr
PlainTV Name
tv
#endif
simplifyBndr TyVarBndr
bndr = TyVarBndr
bndr

-- | Generates a type signature of the form
-- @x :: Member (Effect e) effs => a -> b -> c -> Eff effs r@.
genSig :: Con -> Q Dec
genSig :: Con -> Q Dec
genSig Con
con = do
  let
    getConName :: Con -> f Name
getConName (ForallC [TyVarBndr]
_ Cxt
_ Con
c) = Con -> f Name
getConName Con
c
    getConName (GadtC [Name
n] [BangType]
_ Kind
_) = Name -> f Name
forall (f :: * -> *) a. Applicative f => a -> f a
pure Name
n
    getConName Con
c = String -> f Name
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> f Name) -> String -> f Name
forall a b. (a -> b) -> a -> b
$ String
"failed to get GADT name from " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Con -> String
forall a. Show a => a -> String
show Con
c
  Name
conName <- Con -> Q Name
forall (f :: * -> *). MonadFail f => Con -> f Name
getConName Con
con
  Name -> Kind -> Dec
SigD (Name -> Name
getDeclName Name
conName) (Kind -> Dec) -> (Kind -> Kind) -> Kind -> Dec
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Kind -> Kind
simplifyBndrs (Kind -> Dec) -> Q Kind -> Q Dec
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Con -> Q Kind
genType Con
con

-- | Folds a list of 'Type's into a right-associative arrow 'Type'.
foldArrows :: [Type] -> Type
foldArrows :: Cxt -> Kind
foldArrows = (Kind -> Kind -> Kind) -> Cxt -> Kind
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 (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)