{-# LANGUAGE Trustworthy, TemplateHaskell, LambdaCase, ViewPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE PatternSynonyms #-}
------------------------------------------------------------------------
-- |
-- Module      :  Data.Extensible.Effect.TH
-- Copyright   :  (c) Fumiaki Kinoshita 2021
-- License     :  BSD3
--
-- Maintainer  :  Fumiaki Kinoshita <fumiexcel@gmail.com>
--
------------------------------------------------------------------------
module Data.Extensible.Effect.TH (decEffects
  , decEffectSet
  , decEffectSuite
  , customDecEffects) where

import Data.Extensible.Effect
import Data.List (nub)
import Language.Haskell.TH
import Data.Char
import Control.Monad
import Type.Membership

-- | Generate named effects from a GADT declaration.
--
-- @
-- decEffects [d|
--  data Blah a b x where
--    Blah :: Int -> a -> Blah a b b
--  |]
-- @
--
-- generates
--
-- @
-- type Blah a b = \"Blah\" >: Action '[Int, a] b
-- blah :: forall xs a b
--   . Associate \"Blah\" (Action '[Int, a] b) xs
--   => Int -> a -> Eff xs b
-- blah a0 a1
--   = liftEff
--     (Data.Proxy.Proxy :: Data.Proxy.Proxy \"Blah\")
--     (AArgument a0 (AArgument a1 AResult))
-- @
decEffects :: DecsQ -> DecsQ
decEffects :: DecsQ -> DecsQ
decEffects = Bool -> Bool -> DecsQ -> DecsQ
customDecEffects Bool
False Bool
True

-- | Instead of making a type synonym for individual actions, it defines a list
-- of actions.
decEffectSet :: DecsQ -> DecsQ
decEffectSet :: DecsQ -> DecsQ
decEffectSet = Bool -> Bool -> DecsQ -> DecsQ
customDecEffects Bool
True Bool
False

-- | Generates type synonyms for the set of actions and also individual actions.
decEffectSuite :: DecsQ -> DecsQ
decEffectSuite :: DecsQ -> DecsQ
decEffectSuite = Bool -> Bool -> DecsQ -> DecsQ
customDecEffects Bool
True Bool
True

-- | Generate effect suite with custom settings.
customDecEffects :: Bool -- ^ generate a synonym of the set of actions
    -> Bool -- ^ generate synonyms for individual actions
    -> DecsQ -> DecsQ
customDecEffects :: Bool -> Bool -> DecsQ -> DecsQ
customDecEffects Bool
synSet Bool
synActions DecsQ
decs = DecsQ
decs DecsQ -> ([Dec] -> DecsQ) -> DecsQ
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \[Dec]
ds -> ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Dec]] -> DecsQ) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> a -> b
$ [Dec] -> (Dec -> DecsQ) -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Dec]
ds ((Dec -> DecsQ) -> Q [[Dec]]) -> (Dec -> DecsQ) -> Q [[Dec]]
forall a b. (a -> b) -> a -> b
$ \case
  DataD Cxt
_ Name
dataName [TyVarBndr]
tparams Maybe Kind
_ [Con]
cs [DerivClause]
_
    -> do
      ([(Name, Kind)]
cxts, [[Dec]]
dcs) <- ([((Name, Kind), [Dec])] -> ([(Name, Kind)], [[Dec]]))
-> Q [((Name, Kind), [Dec])] -> Q ([(Name, Kind)], [[Dec]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [((Name, Kind), [Dec])] -> ([(Name, Kind)], [[Dec]])
forall a b. [(a, b)] -> ([a], [b])
unzip (Q [((Name, Kind), [Dec])] -> Q ([(Name, Kind)], [[Dec]]))
-> Q [((Name, Kind), [Dec])] -> Q ([(Name, Kind)], [[Dec]])
forall a b. (a -> b) -> a -> b
$ (Con -> Q ((Name, Kind), [Dec]))
-> [Con] -> Q [((Name, Kind), [Dec])]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ([TyVarBndr] -> Con -> Q ((Name, Kind), [Dec])
con2Eff [TyVarBndr]
tparams) [Con]
cs

      let vars :: [TyVarBndr]
vars = (Name -> TyVarBndr) -> [Name] -> [TyVarBndr]
forall a b. (a -> b) -> [a] -> [b]
map Name -> TyVarBndr
mkPlainTV ([Name] -> [TyVarBndr]) -> [Name] -> [TyVarBndr]
forall a b. (a -> b) -> a -> b
$ [Name] -> [Name]
forall a. Eq a => [a] -> [a]
nub ([Name] -> [Name]) -> [Name] -> [Name]
forall a b. (a -> b) -> a -> b
$ ((Name, Kind) -> [Name]) -> [(Name, Kind)] -> [Name]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Kind -> [Name]
varsT (Kind -> [Name])
-> ((Name, Kind) -> Kind) -> (Name, Kind) -> [Name]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Kind) -> Kind
forall a b. (a, b) -> b
snd) [(Name, Kind)]
cxts
      [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> DecsQ) -> [Dec] -> DecsQ
forall a b. (a -> b) -> a -> b
$ [Name -> [TyVarBndr] -> Kind -> Dec
TySynD Name
dataName [TyVarBndr]
vars (Cxt -> Kind
typeListT (Cxt -> Kind) -> Cxt -> Kind
forall a b. (a -> b) -> a -> b
$ ((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)]
cxts) | Bool
synSet]
          [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [ Name -> [TyVarBndr] -> Kind -> Dec
TySynD Name
k ((Name -> TyVarBndr) -> [Name] -> [TyVarBndr]
forall a b. (a -> b) -> [a] -> [b]
map Name -> TyVarBndr
mkPlainTV ([Name] -> [TyVarBndr]) -> [Name] -> [TyVarBndr]
forall a b. (a -> b) -> a -> b
$ [Name] -> [Name]
forall a. Eq a => [a] -> [a]
nub ([Name] -> [Name]) -> [Name] -> [Name]
forall a b. (a -> b) -> a -> b
$ Kind -> [Name]
varsT Kind
t) Kind
t | Bool
synActions, (Name
k, Kind
t) <- [(Name, Kind)]
cxts]
          [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Dec]]
dcs
  Dec
_ -> String -> DecsQ
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"mkEffects accepts GADT declaration"
  where
#if MIN_VERSION_template_haskell(2,17,0)
    mkPlainTV n = PlainTV n ()
#else
    mkPlainTV :: Name -> TyVarBndr
mkPlainTV = Name -> TyVarBndr
PlainTV
#endif

#if MIN_VERSION_template_haskell(2,17,0)
con2Eff :: [TyVarBndr ()] -> Con -> Q ((Name, Type), [Dec])
#else
con2Eff :: [TyVarBndr] -> Con -> Q ((Name, Type), [Dec])
#endif
con2Eff :: [TyVarBndr] -> Con -> Q ((Name, Kind), [Dec])
con2Eff [TyVarBndr]
_ (GadtC [Name
name] [BangType]
st (AppT Kind
_ Kind
resultT))
  = ((Name, Kind), [Dec]) -> Q ((Name, Kind), [Dec])
forall (m :: * -> *) a. Monad m => a -> m a
return (((Name, Kind), [Dec]) -> Q ((Name, Kind), [Dec]))
-> ((Name, Kind), [Dec]) -> Q ((Name, Kind), [Dec])
forall a b. (a -> b) -> a -> b
$ Name -> Cxt -> Kind -> ((Name, Kind), [Dec])
effectFunD Name
name ((BangType -> Kind) -> [BangType] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map BangType -> Kind
forall a b. (a, b) -> b
snd [BangType]
st) Kind
resultT
con2Eff [TyVarBndr]
tparams (ForallC [TyVarBndr]
_ Cxt
eqs (NormalC Name
name [BangType]
st))
  = ((Name, Kind), [Dec]) -> Q ((Name, Kind), [Dec])
forall (m :: * -> *) a. Monad m => a -> m a
return (((Name, Kind), [Dec]) -> Q ((Name, Kind), [Dec]))
-> ((Name, Kind), [Dec]) -> Q ((Name, Kind), [Dec])
forall a b. (a -> b) -> a -> b
$ [TyVarBndr] -> Cxt -> Name -> [BangType] -> ((Name, Kind), [Dec])
fromMangledGADT [TyVarBndr]
tparams Cxt
eqs Name
name [BangType]
st
con2Eff [TyVarBndr]
tparams (ForallC [TyVarBndr]
_ Cxt
_ Con
c) = [TyVarBndr] -> Con -> Q ((Name, Kind), [Dec])
con2Eff [TyVarBndr]
tparams Con
c
con2Eff [TyVarBndr]
_ Con
p = do
  IO () -> Q ()
forall a. IO a -> Q a
runIO (Con -> IO ()
forall a. Show a => a -> IO ()
print Con
p)
  String -> Q ((Name, Kind), [Dec])
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Unsupported constructor"

#if MIN_VERSION_template_haskell(2,17,0)
fromMangledGADT :: [TyVarBndr ()] -> [Type] -> Name -> [(Strict, Type)] -> ((Name, Type), [Dec])
#else
fromMangledGADT :: [TyVarBndr] -> [Type] -> Name -> [(Strict, Type)] -> ((Name, Type), [Dec])
#endif
fromMangledGADT :: [TyVarBndr] -> Cxt -> Name -> [BangType] -> ((Name, Kind), [Dec])
fromMangledGADT [TyVarBndr]
tyvars_ Cxt
eqs Name
con [BangType]
fieldTypes
  = Name -> Cxt -> Kind -> ((Name, Kind), [Dec])
effectFunD Name
con Cxt
argumentsT Kind
result
  where
#if MIN_VERSION_template_haskell(2,17,0)
    getTV (PlainTV n _) = n
    getTV (KindedTV n _ _) = n
#else
    getTV :: TyVarBndr -> Name
getTV (PlainTV Name
n) = Name
n
    getTV (KindedTV Name
n Kind
_) = Name
n
#endif

    tyvars :: [Name]
tyvars = (TyVarBndr -> Name) -> [TyVarBndr] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr -> Name
getTV [TyVarBndr]
tyvars_

    dic_ :: [(Name, Kind)]
dic_ = [(Name
v, Kind
t) | AppT (AppT Kind
EqualityT (VarT Name
v)) Kind
t <- Cxt
eqs]
    dic :: [(Name, Kind)]
dic = [(Name, Kind)]
dic_ [(Name, Kind)] -> [(Name, Kind)] -> [(Name, Kind)]
forall a. [a] -> [a] -> [a]
++ [(Name
t, Name -> Kind
VarT Name
v) | (Name
v, VarT Name
t) <- [(Name, Kind)]
dic_]

    params' :: [(Name, Name)]
params' = do
      (Name
t, Name
v) <- [Name] -> [Name] -> [(Name, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
tyvars [Name]
uniqueNames
      case Name -> [(Name, Kind)] -> Maybe Kind
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
t [(Name, Kind)]
dic of
        Just (VarT Name
p) -> (Name, Name) -> [(Name, Name)]
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
t, Name
p)
        Maybe Kind
_ -> (Name, Name) -> [(Name, Name)]
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
t, Name
v)

    argumentsT :: Cxt
argumentsT = (BangType -> Kind) -> [BangType] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (\case
      (Strict
_, VarT Name
n) -> Kind -> (Name -> Kind) -> Maybe Name -> Kind
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Name -> Kind
VarT Name
n) Name -> Kind
VarT (Maybe Name -> Kind) -> Maybe Name -> Kind
forall a b. (a -> b) -> a -> b
$ Name -> [(Name, Name)] -> Maybe Name
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
n [(Name, Name)]
params'
      (Strict
_, Kind
x) -> Kind
x) [BangType]
fieldTypes

    result :: Kind
result = case Name -> [(Name, Kind)] -> Maybe Kind
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup ([Name] -> Name
forall a. [a] -> a
last [Name]
tyvars) [(Name, Kind)]
dic of
      Just (VarT Name
v) -> case Name -> [(Name, Name)] -> Maybe Name
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Name
v [(Name, Name)]
params' of
        Just Name
p -> Name -> Kind
VarT Name
p
        Maybe Name
Nothing -> Name -> Kind
VarT Name
v
      Just Kind
t -> Kind
t
      Maybe Kind
Nothing -> Name -> Kind
VarT (Name -> Kind) -> Name -> Kind
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"x"

varsT :: Type -> [Name]
varsT :: Kind -> [Name]
varsT (VarT Name
v) = [Name
v]
varsT (AppT Kind
s Kind
t) = Kind -> [Name]
varsT Kind
s [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
++ Kind -> [Name]
varsT Kind
t
varsT Kind
_ = []

effectFunD :: Name
  -> [Type]
  -> Type
  -> ((Name, Type), [Dec])
effectFunD :: Name -> Cxt -> Kind -> ((Name, Kind), [Dec])
effectFunD Name
key Cxt
argumentsT Kind
resultT = ((Name
key, Name -> Kind
PromotedT '(:>) Kind -> Kind -> Kind
`AppT` Kind
nameT Kind -> Kind -> Kind
`AppT` Kind
actionT)
  , [Name -> Kind -> Dec
SigD Name
fName Kind
typ, Name -> [Clause] -> Dec
FunD Name
fName [Kind -> Int -> Clause
effClause Kind
nameT (Cxt -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Cxt
argumentsT)]]) where

    varList :: Name
varList = String -> Name
mkName String
"xs"

    fName :: Name
fName = let (Char
ch : String
rest) = Name -> String
nameBase Name
key in String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ Char -> Char
toLower Char
ch Char -> String -> String
forall a. a -> [a] -> [a]
: String
rest

#if MIN_VERSION_template_haskell(2,17,0)
    mkPlainTV n = PlainTV n SpecifiedSpec
#else
    mkPlainTV :: Name -> TyVarBndr
mkPlainTV = Name -> TyVarBndr
PlainTV
#endif

    typ :: Kind
typ = [TyVarBndr] -> Cxt -> Kind -> Kind
ForallT ((Name -> TyVarBndr) -> [Name] -> [TyVarBndr]
forall a b. (a -> b) -> [a] -> [b]
map Name -> TyVarBndr
mkPlainTV ([Name] -> [TyVarBndr]) -> [Name] -> [TyVarBndr]
forall a b. (a -> b) -> a -> b
$ Name
varList Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: Kind -> [Name]
varsT Kind
resultT [Name] -> [Name] -> [Name]
forall a. [a] -> [a] -> [a]
++ (Kind -> [Name]) -> Cxt -> [Name]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Kind -> [Name]
varsT Cxt
argumentsT)
        [Kind -> Kind -> Name -> Kind
associateT Kind
nameT Kind
actionT Name
varList]
        (Kind -> Kind) -> Kind -> Kind
forall a b. (a -> b) -> a -> b
$ Name -> Cxt -> Kind -> Kind
effectFunT Name
varList Cxt
argumentsT Kind
resultT

    -- Action [a, B, C] R
    actionT :: Kind
actionT = Name -> Kind
ConT ''Action Kind -> Kind -> Kind
`AppT` Cxt -> Kind
typeListT Cxt
argumentsT Kind -> Kind -> Kind
`AppT` Kind
resultT

    nameT :: Kind
nameT = TyLit -> Kind
LitT (TyLit -> Kind) -> TyLit -> Kind
forall a b. (a -> b) -> a -> b
$ String -> TyLit
StrTyLit (String -> TyLit) -> String -> TyLit
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase Name
key

effectFunT :: Name
  -> [Type]
  -> Type
  -> Type
effectFunT :: Name -> Cxt -> Kind -> Kind
effectFunT Name
varList Cxt
argumentsT Kind
resultT
  = (Kind -> Kind -> Kind) -> Kind -> Cxt -> Kind
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Kind
x Kind
y -> Kind
ArrowT Kind -> Kind -> Kind
`AppT` Kind
x Kind -> Kind -> Kind
`AppT` Kind
y) Kind
rt Cxt
argumentsT where
    rt :: Kind
rt = Name -> Kind
ConT ''Eff Kind -> Kind -> Kind
`AppT` Name -> Kind
VarT Name
varList Kind -> Kind -> Kind
`AppT` Kind
resultT

uniqueNames :: [Name]
uniqueNames :: [Name]
uniqueNames = (String -> Name) -> [String] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map String -> Name
mkName ([String] -> [Name]) -> [String] -> [Name]
forall a b. (a -> b) -> a -> b
$ (Int -> [String]) -> [Int] -> [String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Int -> String -> [String]) -> String -> Int -> [String]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> String -> [String]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM [Char
'a'..Char
'z']) [Int
1..]

typeListT :: [Type] -> Type
typeListT :: Cxt -> Kind
typeListT = (Kind -> Kind -> Kind) -> Kind -> Cxt -> Kind
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Kind
x Kind
y -> Kind
PromotedConsT Kind -> Kind -> Kind
`AppT` Kind
x Kind -> Kind -> Kind
`AppT` Kind
y) Kind
PromotedNilT

associateT :: Type -- key
  -> Type -- type
  -> Name -- variable
  -> Type
associateT :: Kind -> Kind -> Name -> Kind
associateT Kind
nameT Kind
t Name
xs = Name -> Kind
ConT ''Lookup Kind -> Kind -> Kind
`AppT` Name -> Kind
VarT Name
xs Kind -> Kind -> Kind
`AppT` Kind
nameT Kind -> Kind -> Kind
`AppT` Kind
t

effClause :: Type -- effect key
  -> Int -- number of arguments
  -> Clause
effClause :: Kind -> Int -> Clause
effClause Kind
nameT Int
n = [Pat] -> Body -> [Dec] -> Clause
Clause ((Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
argNames) (Exp -> Body
NormalB Exp
rhs) [] where
  -- liftEff (Proxy :: Proxy "Foo")
  lifter :: Exp
lifter = Name -> Exp
VarE 'liftEff Exp -> Exp -> Exp
`AppE` (Name -> Exp
ConE 'Proxy Exp -> Kind -> Exp
`SigE` Kind -> Kind -> Kind
AppT (Name -> Kind
ConT ''Proxy) Kind
nameT)

  argNames :: [Name]
argNames = (Int -> Name) -> [Int] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map (String -> Name
mkName (String -> Name) -> (Int -> String) -> Int -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"a" String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> String) -> (Int -> String) -> Int -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String
forall a. Show a => a -> String
show) [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]

  rhs :: Exp
rhs = Exp
lifter Exp -> Exp -> Exp
`AppE` (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Exp
x Exp
y -> Name -> Exp
ConE 'AArgument Exp -> Exp -> Exp
`AppE` Exp
x Exp -> Exp -> Exp
`AppE` Exp
y)
    (Name -> Exp
ConE 'AResult)
    ((Name -> Exp) -> [Name] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE [Name]
argNames)