{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Use <=<" #-}

-- 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 licensed under the BSD3 License as
    shown in [1].  The modified code, in its entirety, is licensed under
    MPL 2.0. When redistributing, please ensure that you do not remove
    the BSD3 License text as indicated in [1].
    <https://github.com/pa-ba/compdata/blob/master/src/Data/Comp/Multi/Derive/HFunctor.hs>

    [1] Copyright (c) 2010--2011 Patrick Bahr, Tom Hvitved

        All rights reserved.

        Redistribution and use in source and binary forms, with or without
        modification, are permitted provided that the following conditions
        are met:

        1. Redistributions of source code must retain the above copyright
        notice, this list of conditions and the following disclaimer.

        2. Redistributions in binary form must reproduce the above copyright
        notice, this list of conditions and the following disclaimer in the
        documentation and/or other materials provided with the distribution.

        3. Neither the name of the author nor the names of his contributors
        may be used to endorse or promote products derived from this software
        without specific prior written permission.

        THIS SOFTWARE IS PROVIDED BY THE AUTHORS ``AS IS'' AND ANY EXPRESS OR
        IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
        WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
        DISCLAIMED.  IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE FOR
        ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
        DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
        OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
        HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
        STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
        ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
        POSSIBILITY OF SUCH DAMAGE.
-}

{- |
Copyright   :  (c) 2010-2011 Patrick Bahr, Tom Hvitved
               (c) 2023 Yamada Ryo
License     :  MPL-2.0 (see the file LICENSE)
Maintainer  :  ymdfield@outlook.jp
Stability   :  experimental
Portability :  portable
-}
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

{- |
Derive an instance of t'Data.Effect.HFunctor.HFunctor' for a type constructor of any higher-order
kind taking at least two arguments.
-}
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

        -- The algorithm is based on: https://gitlab.haskell.org/ghc/ghc/-/wikis/commentary/compiler/derive-functor
        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)
                        -- todo: tuple support
                        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 #-}

-- * Utility functions

{- |
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)