{-# LANGUAGE TemplateHaskellQuotes #-}

-- This Source Code Form is subject to the terms of the Mozilla Public
-- License, v. 2.0. If a copy of the MPL was not distributed with this
-- file, You can obtain one at https://mozilla.org/MPL/2.0/.

{-
The code before modification is BSD3 licensed, (c) 2010-2011 Patrick Bahr.
<https://github.com/pa-ba/compdata/blob/master/src/Data/Comp/Multi/Derive/HFunctor.hs>

This fork was made to work around the problem that the
'Control.Effect.Class.Machinery.TH.makeEffect' function that generates multiple
definitions at once for convenience is not possible with only the original
'Data.Comp.Multi.Derive.makeHFunctor' function due to TH limitations,
because the original function takes the name of the data type as an argument,
but there is no version that takes 'DataInfo' as an argument (the data type
reification and the HFunctor derivation process are not separated as functions).
-}

{- |
Copyright   :  (c) 2010-2011 Patrick Bahr
               (c) 2023 Yamada Ryo
License     :  MPL-2.0 (see the file LICENSE)
Maintainer  :  ymdfield@outlook.jp
Stability   :  experimental
Portability :  portable
-}
module Data.Effect.Class.TH.HFunctor.Internal where

import Control.Effect.Class.Machinery.HFunctor (HFunctor, hfmap)
import Control.Monad (replicateM, (<=<))
import Data.Maybe (catMaybes)
import Language.Haskell.TH (
    Body (NormalB),
    Clause (Clause),
    Con (ForallC, GadtC, InfixC, NormalC, RecC),
    Cxt,
    Dec (DataD, InstanceD, NewtypeD),
    DerivClause,
    Exp,
    Info (TyConI),
    Name,
    Pat (ConP, VarP, WildP),
    Q,
    Quote (..),
    TyVarBndr (..),
    Type (AppT, ConT, ForallT, SigT, VarT),
    appE,
    conE,
    funD,
    varE,
 )
import Language.Haskell.TH.Syntax (StrictType)

{- |
Derive an instance of 'HFunctor' for a type constructor of any higher-order kind taking at least two
arguments, from 'DataInfo'.
-}
deriveHFunctor :: DataInfo flag -> Q [Dec]
deriveHFunctor :: forall flag. DataInfo flag -> Q [Dec]
deriveHFunctor (DataInfo Cxt
_cxt Name
name [TyVarBndr flag]
args [Con]
constrs [DerivClause]
_deriving) = do
    let args' :: [TyVarBndr flag]
args' = forall a. [a] -> [a]
init [TyVarBndr flag]
args
        fArg :: Type
fArg = Name -> Type
VarT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. TyVarBndr a -> Name
tyVarName forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
last [TyVarBndr flag]
args'
        argNames :: Cxt
argNames = 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) (forall a. [a] -> [a]
init [TyVarBndr flag]
args')
        complType :: Type
complType = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT (Name -> Type
ConT Name
name) Cxt
argNames
        classType :: Type
classType = Type -> Type -> Type
AppT (Name -> Type
ConT ''HFunctor) Type
complType
    [(Q Exp, Pat,
  (Int -> Q Exp -> Q Exp) -> (Q Exp -> Q Exp) -> [Q Exp], Bool,
  [Q Exp], [(Int, Name)])]
constrs' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall t c.
(Name, [[t]])
-> Q (Q Exp, Pat, (t -> Q Exp -> c) -> (Q Exp -> c) -> [c], Bool,
      [Q Exp], [(t, Name)])
mkPatAndVars forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. Type -> (a, Cxt, Maybe Type) -> (a, [[Int]])
isFarg Type
fArg forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Con -> Q (Name, Cxt, Maybe Type)
normalConExp) [Con]
constrs
    Dec
hfmapDecl <- forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
funD 'hfmap (forall a b. (a -> b) -> [a] -> [b]
map forall {m :: * -> *} {m :: * -> *} {t} {t :: * -> *} {a} {e} {f}.
(Quote m, Quote m, Eq t, Num t, Foldable t) =>
(m Exp, Pat, (t -> m Exp -> m Exp) -> (a -> a) -> t (m Exp), Bool,
 e, f)
-> m Clause
hfmapClause [(Q Exp, Pat,
  (Int -> Q Exp -> Q Exp) -> (Q Exp -> Q Exp) -> [Q Exp], Bool,
  [Q Exp], [(Int, Name)])]
constrs')
    forall (m :: * -> *) a. Monad m => a -> m a
return [Cxt -> Type -> [Dec] -> Dec
mkInstanceD [] Type
classType [Dec
hfmapDecl]]
  where
    isFarg :: Type -> (a, Cxt, Maybe Type) -> (a, [[Int]])
isFarg Type
fArg (a
constr, Cxt
args_, Maybe Type
ty) = (a
constr, forall a b. (a -> b) -> [a] -> [b]
map (Type -> Type -> [Int]
`containsType'` Type -> Maybe Type -> Type
getBinaryFArg Type
fArg Maybe Type
ty) Cxt
args_)
    filterVar :: (t -> t -> t) -> (t -> t) -> [t] -> t -> t
filterVar t -> t -> t
_ t -> t
nonFarg [] t
x = t -> t
nonFarg t
x
    filterVar t -> t -> t
farg t -> t
_ [t
depth] t
x = t -> t -> t
farg t
depth t
x
    filterVar t -> t -> t
_ t -> t
_ [t]
_ t
_ = forall a. HasCallStack => [Char] -> a
error [Char]
"functor variable occurring twice in argument type"
    filterVars :: [[t]] -> [b] -> (t -> b -> c) -> (b -> c) -> [c]
filterVars [[t]]
args_ [b]
varNs t -> b -> c
farg b -> c
nonFarg = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall {t} {t} {t}. (t -> t -> t) -> (t -> t) -> [t] -> t -> t
filterVar t -> b -> c
farg b -> c
nonFarg) [[t]]
args_ [b]
varNs
    mkCPat :: Name -> [Name] -> Pat
mkCPat Name
constr [Name]
varNs = Name -> Cxt -> [Pat] -> Pat
ConP Name
constr [] forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
mkPat [Name]
varNs
    mkPat :: Name -> Pat
mkPat = Name -> Pat
VarP
    mkPatAndVars ::
        (Name, [[t]]) ->
        Q (Q Exp, Pat, (t -> Q Exp -> c) -> (Q Exp -> c) -> [c], Bool, [Q Exp], [(t, Name)])
    mkPatAndVars :: forall t c.
(Name, [[t]])
-> Q (Q Exp, Pat, (t -> Q Exp -> c) -> (Q Exp -> c) -> [c], Bool,
      [Q Exp], [(t, Name)])
mkPatAndVars (Name
constr, [[t]]
args_) =
        do
            [Name]
varNs <- Int -> [Char] -> Q [Name]
newNames (forall (t :: * -> *) a. Foldable t => t a -> Int
length [[t]]
args_) [Char]
"x"
            forall (m :: * -> *) a. Monad m => a -> m a
return
                ( forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
constr
                , Name -> [Name] -> Pat
mkCPat Name
constr [Name]
varNs
                , \t -> Q Exp -> c
f Q Exp -> c
g -> forall {t} {b} {c}.
[[t]] -> [b] -> (t -> b -> c) -> (b -> c) -> [c]
filterVars [[t]]
args_ [Name]
varNs (\t
d Name
x -> t -> Q Exp -> c
f t
d (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
x)) (Q Exp -> c
g forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). Quote m => Name -> m Exp
varE)
                , Bool -> Bool
not (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[t]]
args_)
                , forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
varNs
                , forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$ forall {t} {b} {c}.
[[t]] -> [b] -> (t -> b -> c) -> (b -> c) -> [c]
filterVars [[t]]
args_ [Name]
varNs (forall a b c. ((a, b) -> c) -> a -> b -> c
curry forall a. a -> Maybe a
Just) (forall a b. a -> b -> a
const forall a. Maybe a
Nothing)
                )
    hfmapClause :: (m Exp, Pat, (t -> m Exp -> m Exp) -> (a -> a) -> t (m Exp), Bool,
 e, f)
-> m Clause
hfmapClause (m Exp
con, Pat
pat, (t -> m Exp -> m Exp) -> (a -> a) -> t (m Exp)
vars', Bool
hasFargs, e
_, f
_) =
        do
            Name
fn <- forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"f"
            let f :: m Exp
f = forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
fn
                fp :: Pat
fp = if Bool
hasFargs then Name -> Pat
VarP Name
fn else Pat
WildP
                vars :: t (m Exp)
vars = (t -> m Exp -> m Exp) -> (a -> a) -> t (m Exp)
vars' (\t
d m Exp
x -> forall t (m :: * -> *).
(Eq t, Num t, Quote m) =>
t -> m Exp -> m Exp -> m Exp
iter t
d [|fmap|] m Exp
f forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` m Exp
x) forall a. a -> a
id
            Exp
body <- forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE m Exp
con t (m Exp)
vars
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
fp, Pat
pat] (Exp -> Body
NormalB Exp
body) []

-- | A reified information of a datatype.
data DataInfo flag = DataInfo
    { forall flag. DataInfo flag -> Cxt
dataCxt :: Cxt
    , forall flag. DataInfo flag -> Name
dataName :: Name
    , forall flag. DataInfo flag -> [TyVarBndr flag]
dataTyVars :: [TyVarBndr flag]
    , forall flag. DataInfo flag -> [Con]
dataCons :: [Con]
    , forall flag. DataInfo flag -> [DerivClause]
dataDerivings :: [DerivClause]
    }

{- |
This function abstracts away @newtype@ declaration, it turns them into
@data@ declarations.
-}
abstractNewtype :: Info -> Maybe (DataInfo ())
abstractNewtype :: Info -> Maybe (DataInfo ())
abstractNewtype = \case
    TyConI (NewtypeD Cxt
cxt Name
name [TyVarBndr ()]
args Maybe Type
_ Con
constr [DerivClause]
derive) -> forall a. a -> Maybe a
Just (forall flag.
Cxt
-> Name
-> [TyVarBndr flag]
-> [Con]
-> [DerivClause]
-> DataInfo flag
DataInfo Cxt
cxt Name
name [TyVarBndr ()]
args [Con
constr] [DerivClause]
derive)
    TyConI (DataD Cxt
cxt Name
name [TyVarBndr ()]
args Maybe Type
_ [Con]
constrs [DerivClause]
derive) -> forall a. a -> Maybe a
Just (forall flag.
Cxt
-> Name
-> [TyVarBndr flag]
-> [Con]
-> [DerivClause]
-> DataInfo flag
DataInfo Cxt
cxt Name
name [TyVarBndr ()]
args [Con]
constrs [DerivClause]
derive)
    Info
_ -> forall a. Maybe a
Nothing

-- | Convert the reified information of the datatype to a definition.
infoToDataD :: DataInfo () -> Dec
infoToDataD :: DataInfo () -> Dec
infoToDataD (DataInfo Cxt
cxt Name
name [TyVarBndr ()]
args [Con]
cons [DerivClause]
deriv) = Cxt
-> Name
-> [TyVarBndr ()]
-> Maybe Type
-> [Con]
-> [DerivClause]
-> Dec
DataD Cxt
cxt Name
name [TyVarBndr ()]
args forall a. Maybe a
Nothing [Con]
cons [DerivClause]
deriv

{- |
This function provides the name and the arity of the given data
constructor, and if it is a GADT also its type.
-}
normalCon :: Con -> (Name, [StrictType], Maybe Type)
normalCon :: Con -> (Name, [StrictType], Maybe Type)
normalCon (NormalC Name
constr [StrictType]
args) = (Name
constr, [StrictType]
args, forall a. Maybe a
Nothing)
normalCon (RecC Name
constr [VarBangType]
args) = (Name
constr, forall a b. (a -> b) -> [a] -> [b]
map (\(Name
_, Bang
s, Type
t) -> (Bang
s, Type
t)) [VarBangType]
args, forall a. Maybe a
Nothing)
normalCon (InfixC StrictType
a Name
constr StrictType
b) = (Name
constr, [StrictType
a, StrictType
b], forall a. Maybe a
Nothing)
normalCon (ForallC [TyVarBndr Specificity]
_ Cxt
_ Con
constr) = Con -> (Name, [StrictType], Maybe Type)
normalCon Con
constr
normalCon (GadtC (Name
constr : [Name]
_) [StrictType]
args Type
typ) = (Name
constr, [StrictType]
args, forall a. a -> Maybe a
Just Type
typ)
normalCon Con
_ = forall a. HasCallStack => [Char] -> a
error [Char]
"missing case for 'normalCon'"

normalConExp :: Con -> Q (Name, [Type], Maybe Type)
normalConExp :: Con -> Q (Name, Cxt, Maybe Type)
normalConExp Con
con = forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name
n, forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [StrictType]
ts, Maybe Type
t)
  where
    (Name
n, [StrictType]
ts, Maybe Type
t) = Con -> (Name, [StrictType], Maybe Type)
normalCon Con
con

containsType' :: Type -> Type -> [Int]
containsType' :: Type -> Type -> [Int]
containsType' = forall {t}. Num t => t -> Type -> Type -> [t]
run Int
0
  where
    run :: t -> Type -> Type -> [t]
run t
n Type
s Type
t
        | Type
s forall a. Eq a => a -> a -> Bool
== Type
t = [t
n]
        | Bool
otherwise = case Type
s of
            ForallT [TyVarBndr Specificity]
_ Cxt
_ Type
s' -> t -> Type -> Type -> [t]
run t
n Type
s' Type
t
            -- only going through the right-hand side counts!
            AppT Type
s1 Type
s2 -> t -> Type -> Type -> [t]
run t
n Type
s1 Type
t forall a. [a] -> [a] -> [a]
++ t -> Type -> Type -> [t]
run (t
n forall a. Num a => a -> a -> a
+ t
1) Type
s2 Type
t
            SigT Type
s' Type
_ -> t -> Type -> Type -> [t]
run t
n Type
s' Type
t
            Type
_ -> []

{- |
Auxiliary function to extract the first argument of a binary type
application (the second argument of this function). If the second
argument is @Nothing@ or not of the right shape, the first argument
is returned as a default.
-}
getBinaryFArg :: Type -> Maybe Type -> Type
getBinaryFArg :: Type -> Maybe Type -> Type
getBinaryFArg Type
_ (Just (AppT (AppT Type
_ Type
t) Type
_)) = Type
t
getBinaryFArg Type
def Maybe Type
_ = Type
def

mkInstanceD :: Cxt -> Type -> [Dec] -> Dec
mkInstanceD :: Cxt -> Type -> [Dec] -> Dec
mkInstanceD = Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD forall a. Maybe a
Nothing

{- |
This function provides a list (of the given length) of new names based
on the given string.
-}
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)

-- | pures the name of a type variable.
tyVarName :: TyVarBndr a -> Name
tyVarName :: forall a. TyVarBndr a -> Name
tyVarName (PlainTV Name
n a
_) = Name
n
tyVarName (KindedTV Name
n a
_ Type
_) = Name
n