{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
{-# HLINT ignore "Use <=<" #-}
module Data.Effect.HFunctor.TH.Internal where
import Control.Monad (replicateM, zipWithM)
import Data.Effect.HFunctor (HFunctor, hfmap)
import Data.Effect.TH.Internal (
ConInfo (ConInfo),
DataInfo (DataInfo),
conArgs,
conGadtReturnType,
conName,
occurs,
tyVarName,
tyVarType,
unkindType,
)
import Data.Foldable (foldl')
import Data.Functor ((<&>))
import Data.List.Infinite (Infinite, prependList)
import Data.Text qualified as T
import Formatting (int, sformat, shown, stext, (%))
import Language.Haskell.TH (
Body (NormalB),
Clause (Clause),
Dec (FunD, InstanceD, PragmaD),
Exp (AppE, CaseE, ConE, LamE, TupE, VarE),
Inline (Inline),
Match (Match),
Name,
Pat (ConP, TupP, VarP),
Phases (AllPhases),
Pragma (InlineP),
Q,
Quote (..),
RuleMatch (FunLike),
TyVarBndr (PlainTV),
Type (AppT, ArrowT, ConT, ForallT, SigT, TupleT, VarT),
appE,
nameBase,
pprint,
)
import Language.Haskell.TH qualified as TH
deriveHFunctor :: (Infinite (Q TH.Type) -> Q TH.Type) -> DataInfo -> Q [Dec]
deriveHFunctor :: (Infinite (Q Type) -> Q Type) -> DataInfo -> Q [Dec]
deriveHFunctor Infinite (Q Type) -> Q Type
manualCxt (DataInfo Cxt
_ Name
name [TyVarBndr ()]
args [ConInfo]
cons) = do
Name
mapFnName <- forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"_f"
let mapFn :: Exp
mapFn = Name -> Exp
VarE Name
mapFnName
initArgs :: [TyVarBndr ()]
initArgs = forall a. [a] -> [a]
init [TyVarBndr ()]
args
hfArgs :: [TyVarBndr ()]
hfArgs = forall a. [a] -> [a]
init [TyVarBndr ()]
initArgs
hfArgNames :: Cxt
hfArgNames = forall a b. (a -> b) -> [a] -> [b]
map (Name -> Type
VarT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. TyVarBndr a -> Name
tyVarName) [TyVarBndr ()]
hfArgs
hfmapClause :: ConInfo -> Q Clause
hfmapClause :: ConInfo -> Q Clause
hfmapClause ConInfo{[BangType]
Maybe Type
Name
conGadtReturnType :: Maybe Type
conArgs :: [BangType]
conName :: Name
conName :: ConInfo -> Name
conGadtReturnType :: ConInfo -> Maybe Type
conArgs :: ConInfo -> [BangType]
..} = do
let f :: TyVarBndr ()
f = case Maybe Type
conGadtReturnType of
Maybe Type
Nothing -> forall a. [a] -> a
last [TyVarBndr ()]
initArgs
Just Type
t -> case Type
t of
Type
_ `AppT` VarT Name
n `AppT` Type
_ -> forall flag. Name -> flag -> TyVarBndr flag
PlainTV Name
n ()
Type
_ `AppT` (VarT Name
n `SigT` Type
_) `AppT` Type
_ -> forall flag. Name -> flag -> TyVarBndr flag
PlainTV Name
n ()
Type
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Encounted unknown structure: " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> [Char]
pprint Type
t
hfmapE :: TH.Type -> Exp -> Q Exp
hfmapE :: Type -> Exp -> Q Exp
hfmapE Type
tk
| Type -> Bool
fNotOccurs Type
t = forall (f :: * -> *) a. Applicative f => a -> f a
pure
| Bool
otherwise = \Exp
x -> case Type
t of
VarT Name
n `AppT` Type
a | Name
n forall a. Eq a => a -> a -> Bool
== forall a. TyVarBndr a -> Name
tyVarName TyVarBndr ()
f Bool -> Bool -> Bool
&& Type -> Bool
fNotOccurs Type
a -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Exp
mapFn Exp -> Exp -> Exp
`AppE` Exp
x
Type
ArrowT `AppT` Type
c `AppT` Type
d ->
(Exp -> Q Exp) -> Q Exp
wrapLam \Exp
y -> Type -> Exp -> Q Exp
hfmapE Type
d forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp
x `AppE`) forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> Exp -> Q Exp
cohfmapE Type
c Exp
y
Type
g `AppT` Type
a
| Type -> Bool
fNotOccurs Type
g ->
((Name -> Exp
VarE 'fmap `AppE`) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> Q Exp) -> Q Exp
wrapLam (Type -> Exp -> Q Exp
hfmapE Type
a)) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> (Exp -> Exp -> Exp
`AppE` Exp
x)
Type
ff `AppT` Type
g `AppT` Type
a
| Type -> Bool
fNotOccurs Type
ff Bool -> Bool -> Bool
&& Type -> Bool
fNotOccurs Type
a ->
((Name -> Exp
VarE 'hfmap `AppE`) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> Q Exp) -> Q Exp
wrapLam (Type -> Exp -> Q Exp
hfmapE forall a b. (a -> b) -> a -> b
$ Type
g Type -> Type -> Type
`AppT` Type
a)) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> (Exp -> Exp -> Exp
`AppE` Exp
x)
ForallT [TyVarBndr Specificity]
_ Cxt
_ Type
a -> Type -> Exp -> Q Exp
hfmapE Type
a Exp
x
Type
_ ->
case (Type -> Exp -> Q Exp) -> Type -> Exp -> Maybe (Q Exp)
mapTupleE Type -> Exp -> Q Exp
hfmapE Type
t Exp
x of
Just Q Exp
e -> Q Exp
e
Maybe (Q Exp)
Nothing -> forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail forall a b. (a -> b) -> a -> b
$ [Char]
"Encounted unsupported structure: " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> [Char]
pprint Type
t
where
t :: Type
t = Type -> Type
unkindType Type
tk
cohfmapE :: TH.Type -> Exp -> Q Exp
cohfmapE :: Type -> Exp -> Q Exp
cohfmapE Type
tk
| Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall a. TyVarBndr a -> Name
tyVarName TyVarBndr ()
f Name -> Type -> Bool
`occurs` Type
t = forall (f :: * -> *) a. Applicative f => a -> f a
pure
| Bool
otherwise = \Exp
x -> case Type
t of
VarT Name
n `AppT` Type
a
| Name
n forall a. Eq a => a -> a -> Bool
== forall a. TyVarBndr a -> Name
tyVarName TyVarBndr ()
f Bool -> Bool -> Bool
&& Type -> Bool
fNotOccurs Type
a ->
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail forall a b. (a -> b) -> a -> b
$ [Char]
"Functor type variable occurs in contravariant position: " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> [Char]
pprint Type
t
Type
ArrowT `AppT` Type
c `AppT` Type
d ->
(Exp -> Q Exp) -> Q Exp
wrapLam \Exp
y -> Type -> Exp -> Q Exp
cohfmapE Type
d forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp
x `AppE`) forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Type -> Exp -> Q Exp
hfmapE Type
c Exp
y
Type
g `AppT` Type
a
| Type -> Bool
fNotOccurs Type
g ->
((Name -> Exp
VarE 'fmap `AppE`) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Exp -> Q Exp) -> Q Exp
wrapLam (Type -> Exp -> Q Exp
cohfmapE Type
a)) forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> (Exp -> Exp -> Exp
`AppE` Exp
x)
Type
ff `AppT` Type
_ `AppT` Type
a
| Type -> Bool
fNotOccurs Type
ff Bool -> Bool -> Bool
&& Type -> Bool
fNotOccurs Type
a ->
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail forall a b. (a -> b) -> a -> b
$ [Char]
"Functor type variable occurs in contravariant position: " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> [Char]
pprint Type
t
ForallT [TyVarBndr Specificity]
_ Cxt
_ Type
b' -> Type -> Exp -> Q Exp
cohfmapE Type
b' Exp
x
Type
_ ->
case (Type -> Exp -> Q Exp) -> Type -> Exp -> Maybe (Q Exp)
mapTupleE Type -> Exp -> Q Exp
cohfmapE Type
t Exp
x of
Just Q Exp
e -> Q Exp
e
Maybe (Q Exp)
Nothing -> forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail forall a b. (a -> b) -> a -> b
$ [Char]
"Encounted unsupported structure: " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> [Char]
pprint Type
t
where
t :: Type
t = Type -> Type
unkindType Type
tk
fNotOccurs :: Type -> Bool
fNotOccurs = Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. TyVarBndr a -> Name
tyVarName TyVarBndr ()
f `occurs`)
[Name]
vars <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall (t :: * -> *) a. Foldable t => t a -> Int
length [BangType]
conArgs) (forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"x")
[Exp]
mappedArgs <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Type -> Exp -> Q Exp
hfmapE (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [BangType]
conArgs) (forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE [Name]
vars)
let body :: Exp
body = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Exp -> Exp -> Exp
AppE (Name -> Exp
ConE Name
conName) [Exp]
mappedArgs
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
mapFnName, Name -> Cxt -> [Pat] -> Pat
ConP Name
conName [] (forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
vars)] (Exp -> Body
NormalB Exp
body) []
Type
cxt <-
Infinite (Q Type) -> Q Type
manualCxt forall a b. (a -> b) -> a -> b
$
forall a b. (a -> b) -> [a] -> [b]
map (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. TyVarBndr a -> Type
tyVarType) [TyVarBndr ()]
hfArgs
forall a. [a] -> Infinite a -> Infinite a
`prependList` forall a. HasCallStack => [Char] -> a
error
( Text -> [Char]
T.unpack forall a b. (a -> b) -> a -> b
$
forall a. Format Text a -> a
sformat
( Format (Name -> Int -> Text -> Text) (Name -> Int -> Text -> Text)
"Too many data type arguments in use. The number of usable type arguments in the data type ‘"
forall r a r'. Format r a -> Format r' r -> Format r' a
% forall a r. Show a => Format r (a -> r)
shown
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format (Int -> Text -> Text) (Int -> Text -> Text)
"’ to be derived is "
forall r a r'. Format r a -> Format r' r -> Format r' a
% forall a r. Integral a => Format r (a -> r)
int
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format (Text -> Text) (Text -> Text)
". ("
forall r a r'. Format r a -> Format r' r -> Format r' a
% forall r. Format r (Text -> r)
stext
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format Text Text
")"
)
Name
name
(forall (t :: * -> *) a. Foldable t => t a -> Int
length [TyVarBndr ()]
hfArgs)
(Text -> [Text] -> Text
T.intercalate Text
", " forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map ((\Text
t -> Text
"‘" forall a. Semigroup a => a -> a -> a
<> Text
t forall a. Semigroup a => a -> a -> a
<> Text
"’") forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Text
T.pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> [Char]
nameBase forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. TyVarBndr a -> Name
tyVarName) [TyVarBndr ()]
hfArgs)
)
Dec
hfmapDecls <- Name -> [Clause] -> Dec
FunD 'hfmap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ConInfo -> Q Clause
hfmapClause [ConInfo]
cons
let fnInline :: Dec
fnInline = Pragma -> Dec
PragmaD (Name -> Inline -> RuleMatch -> Phases -> Pragma
InlineP 'hfmap Inline
Inline RuleMatch
FunLike Phases
AllPhases)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
[ Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD
forall a. Maybe a
Nothing
[Type
cxt]
(Name -> Type
ConT ''HFunctor Type -> Type -> Type
`AppT` forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Type -> Type -> Type
AppT (Name -> Type
ConT Name
name) Cxt
hfArgNames)
[Dec
hfmapDecls, Dec
fnInline]
]
wrapLam :: (Exp -> Q Exp) -> Q Exp
wrapLam :: (Exp -> Q Exp) -> Q Exp
wrapLam Exp -> Q Exp
f = do
Name
x <- forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"x"
[Pat] -> Exp -> Exp
LamE [Name -> Pat
VarP Name
x] forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Exp -> Q Exp
f (Name -> Exp
VarE Name
x)
mapTupleE :: (TH.Type -> Exp -> Q Exp) -> TH.Type -> Exp -> Maybe (Q Exp)
mapTupleE :: (Type -> Exp -> Q Exp) -> Type -> Exp -> Maybe (Q Exp)
mapTupleE Type -> Exp -> Q Exp
f Type
t Exp
e = do
Cxt
es <- Type -> Maybe Cxt
decomposeTupleT Type
t
let n :: Int
n = forall (t :: * -> *) a. Foldable t => t a -> Int
length Cxt
es
forall a. a -> Maybe a
Just do
[Name]
xs <- Int -> [Char] -> Q [Name]
newNames Int
n [Char]
"x"
[Exp]
ys <- forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Type -> Exp -> Q Exp
f Cxt
es forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE [Name]
xs
forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Exp -> [Match] -> Exp
CaseE Exp
e [Pat -> Body -> [Dec] -> Match
Match ([Pat] -> Pat
TupP forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
xs) (Exp -> Body
NormalB forall a b. (a -> b) -> a -> b
$ [Maybe Exp] -> Exp
TupE forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a. a -> Maybe a
Just [Exp]
ys) []]
decomposeTupleT :: TH.Type -> Maybe [TH.Type]
decomposeTupleT :: Type -> Maybe Cxt
decomposeTupleT = Cxt -> Int -> Type -> Maybe Cxt
go [] Int
0
where
go :: [TH.Type] -> Int -> TH.Type -> Maybe [TH.Type]
go :: Cxt -> Int -> Type -> Maybe Cxt
go Cxt
acc !Int
n = \case
TupleT Int
m | Int
m forall a. Eq a => a -> a -> Bool
== Int
n -> forall a. a -> Maybe a
Just Cxt
acc
Type
f `AppT` Type
a -> Cxt -> Int -> Type -> Maybe Cxt
go (Type
a forall a. a -> [a] -> [a]
: Cxt
acc) (Int
n forall a. Num a => a -> a -> a
+ Int
1) Type
f
Type
_ -> forall a. Maybe a
Nothing
{-# INLINE decomposeTupleT #-}
newNames :: Int -> String -> Q [Name]
newNames :: Int -> [Char] -> Q [Name]
newNames Int
n [Char]
name = forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
name)
iter :: (Eq t, Num t, Quote m) => t -> m Exp -> m Exp -> m Exp
iter :: forall t (m :: * -> *).
(Eq t, Num t, Quote m) =>
t -> m Exp -> m Exp -> m Exp
iter t
0 m Exp
_ m Exp
e = m Exp
e
iter t
n m Exp
f m Exp
e = forall t (m :: * -> *).
(Eq t, Num t, Quote m) =>
t -> m Exp -> m Exp -> m Exp
iter (t
n forall a. Num a => a -> a -> a
- t
1) m Exp
f (m Exp
f forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` m Exp
e)