{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}

-- |
-- Module      :   Grisette.Internal.TH.Derivation.UnaryOpCommon
-- Copyright   :   (c) Sirui Lu 2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.TH.Derivation.UnaryOpCommon
  ( UnaryOpClassConfig (..),
    UnaryOpFieldConfig (..),
    UnaryOpConfig (..),
    UnaryOpFunConfig (..),
    FieldFunExp,
    defaultFieldResFun,
    defaultFieldFunExp,
    genUnaryOpClass,
    defaultUnaryOpInstanceTypeFromConfig,
  )
where

import Control.Monad (replicateM, zipWithM)
import qualified Data.List as List
import qualified Data.Map as M
import Data.Maybe (catMaybes, fromMaybe, mapMaybe)
import qualified Data.Set as S
import Grisette.Internal.TH.Derivation.Common
  ( CheckArgsResult
      ( CheckArgsResult,
        argVars,
        constructors,
        keptVars
      ),
    DeriveConfig (unconstrainedPositions),
    checkArgs,
    ctxForVar,
    evalModeSpecializeList,
    extraConstraint,
    freshenCheckArgsResult,
    isVarUsedInFields,
    specializeResult,
  )
import Grisette.Internal.TH.Util (allUsedNames)
import Language.Haskell.TH
  ( Body (NormalB),
    Clause (Clause),
    Dec (FunD, InstanceD),
    Exp (VarE),
    Kind,
    Name,
    Pat (VarP, WildP),
    Q,
    Type (AppT, ConT, VarT),
    appE,
    clause,
    conP,
    conT,
    funD,
    nameBase,
    newName,
    normalB,
    varE,
    varP,
  )
import Language.Haskell.TH.Datatype
  ( ConstructorInfo (constructorFields, constructorName, constructorVariant),
    ConstructorVariant,
    TypeSubstitution (freeVariables),
    resolveTypeSynonyms,
  )

-- | Type of field function expression generator.
type FieldFunExp = M.Map Name Name -> M.Map Name [Name] -> Type -> Q Exp

-- | Default field function expression generator.
defaultFieldFunExp :: [Name] -> FieldFunExp
defaultFieldFunExp :: [Name] -> FieldFunExp
defaultFieldFunExp [Name]
unaryOpFunNames Map Name Name
argToFunPat Map Name [Name]
_ = Type -> Q Exp
go
  where
    go :: Type -> Q Exp
go Type
ty = do
      let allArgNames :: Set Name
allArgNames = Map Name Name -> Set Name
forall k a. Map k a -> Set k
M.keysSet Map Name Name
argToFunPat
      let typeHasNoArg :: Type -> Bool
typeHasNoArg Type
ty =
            [Name] -> Set Name
forall a. Ord a => [a] -> Set a
S.fromList ([Type] -> [Name]
forall a. TypeSubstitution a => a -> [Name]
freeVariables [Type
ty])
              Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
`S.intersection` Set Name
allArgNames
              Set Name -> Set Name -> Bool
forall a. Eq a => a -> a -> Bool
== Set Name
forall a. Set a
S.empty
      let fun0 :: Q Exp
fun0 = Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Name -> Q Exp) -> Name -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Name] -> Name
forall a. HasCallStack => [a] -> a
head [Name]
unaryOpFunNames
          fun1 :: Type -> Q Exp
fun1 Type
b = [|$(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Name -> Q Exp) -> Name -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Name]
unaryOpFunNames [Name] -> Int -> Name
forall a. HasCallStack => [a] -> Int -> a
!! Int
1) $(Type -> Q Exp
go Type
b)|]
          fun2 :: Type -> Type -> Q Exp
fun2 Type
b Type
c = [|$(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Name -> Q Exp) -> Name -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Name]
unaryOpFunNames [Name] -> Int -> Name
forall a. HasCallStack => [a] -> Int -> a
!! Int
2) $(Type -> Q Exp
go Type
b) $(Type -> Q Exp
go Type
c)|]
          fun3 :: Type -> Type -> Type -> Q Exp
fun3 Type
b Type
c Type
d =
            [|$(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Name -> Q Exp) -> Name -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Name]
unaryOpFunNames [Name] -> Int -> Name
forall a. HasCallStack => [a] -> Int -> a
!! Int
3) $(Type -> Q Exp
go Type
b) $(Type -> Q Exp
go Type
c) $(Type -> Q Exp
go Type
d)|]
      case Type
ty of
        AppT (AppT (AppT (VarT Name
_) Type
b) Type
c) Type
d -> Type -> Type -> Type -> Q Exp
fun3 Type
b Type
c Type
d
        AppT (AppT (VarT Name
_) Type
b) Type
c -> Type -> Type -> Q Exp
fun2 Type
b Type
c
        AppT (VarT Name
_) Type
b -> Type -> Q Exp
fun1 Type
b
        Type
_ | Type -> Bool
typeHasNoArg Type
ty -> Q Exp
fun0
        AppT Type
a Type
b | Type -> Bool
typeHasNoArg Type
a -> Type -> Q Exp
fun1 Type
b
        AppT (AppT Type
a Type
b) Type
c | Type -> Bool
typeHasNoArg Type
a -> Type -> Type -> Q Exp
fun2 Type
b Type
c
        AppT (AppT (AppT Type
a Type
b) Type
c) Type
d | Type -> Bool
typeHasNoArg Type
a -> Type -> Type -> Type -> Q Exp
fun3 Type
b Type
c Type
d
        VarT Name
nm -> case Name -> Map Name Name -> Maybe Name
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
nm Map Name Name
argToFunPat of
          Just Name
pname -> Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
pname
          Maybe Name
_ -> String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Exp) -> String -> Q Exp
forall a b. (a -> b) -> a -> b
$ String
"defaultFieldFunExp: unsupported type: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Type -> String
forall a. Show a => a -> String
show Type
ty
        Type
_ -> String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Exp) -> String -> Q Exp
forall a b. (a -> b) -> a -> b
$ String
"defaultFieldFunExp: unsupported type: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Type -> String
forall a. Show a => a -> String
show Type
ty

-- | Configuration for a unary function field expression generation on a GADT.
data UnaryOpConfig where
  UnaryOpConfig ::
    (UnaryOpFunConfig config) => config -> [Name] -> UnaryOpConfig

-- | Default field result function.
defaultFieldResFun ::
  ConstructorVariant -> Name -> [Exp] -> Int -> Exp -> Exp -> Q (Exp, [Bool])
defaultFieldResFun :: ConstructorVariant
-> Name -> [Exp] -> Int -> Exp -> Exp -> Q (Exp, [Bool])
defaultFieldResFun ConstructorVariant
_ Name
_ [Exp]
extraPatExps Int
_ Exp
fieldPatExp Exp
defaultFieldFunExp = do
  res <-
    Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE
      ( (Q Exp -> Exp -> Q Exp) -> Q Exp -> [Exp] -> Q Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
          (\Q Exp
exp Exp
name -> Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE Q Exp
exp (Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
name))
          (Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
defaultFieldFunExp)
          [Exp]
extraPatExps
      )
      (Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
fieldPatExp)
  return (res, (True <$ extraPatExps))

funPatAndExps ::
  FieldFunExp ->
  (Int -> [String]) ->
  [(Type, Kind)] ->
  [Type] ->
  Q ([Pat], [[Pat]], [Exp])
funPatAndExps :: FieldFunExp
-> (Int -> [String])
-> [(Type, Type)]
-> [Type]
-> Q ([Pat], [[Pat]], [Exp])
funPatAndExps FieldFunExp
fieldFunExpGen Int -> [String]
extraLiftedPatNames [(Type, Type)]
argTypes [Type]
fields = do
  let usedArgs :: Set Name
usedArgs = [Name] -> Set Name
forall a. Ord a => [a] -> Set a
S.fromList ([Name] -> Set Name) -> [Name] -> Set Name
forall a b. (a -> b) -> a -> b
$ [Type] -> [Name]
forall a. TypeSubstitution a => a -> [Name]
freeVariables [Type]
fields
  let liftedNames :: [String]
liftedNames = Int -> [String]
extraLiftedPatNames ([(Type, Type)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Type, Type)]
argTypes)
  args <-
    ((Type, Type) -> Q (Name, Maybe (Name, [Name])))
-> [(Type, Type)] -> Q [(Name, Maybe (Name, [Name]))]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse
      ( \(Type
ty, Type
_) -> do
          case Type
ty of
            VarT Name
nm ->
              if Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Name
nm Set Name
usedArgs
                then do
                  pname <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"p"
                  epname <- traverse newName liftedNames
                  return (nm, Just (pname, epname))
                else (Name, Maybe (Name, [Name])) -> Q (Name, Maybe (Name, [Name]))
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ('undefined, Maybe (Name, [Name])
forall a. Maybe a
Nothing)
            Type
_ -> (Name, Maybe (Name, [Name])) -> Q (Name, Maybe (Name, [Name]))
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ('undefined, Maybe (Name, [Name])
forall a. Maybe a
Nothing)
      )
      [(Type, Type)]
argTypes
  let argToFunPat =
        [(Name, Name)] -> Map Name Name
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, Name)] -> Map Name Name)
-> [(Name, Name)] -> Map Name Name
forall a b. (a -> b) -> a -> b
$ ((Name, Maybe (Name, [Name])) -> Maybe (Name, Name))
-> [(Name, Maybe (Name, [Name]))] -> [(Name, Name)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(Name
nm, Maybe (Name, [Name])
mpat) -> ((Name, [Name]) -> (Name, Name))
-> Maybe (Name, [Name]) -> Maybe (Name, Name)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Name
nm,) (Name -> (Name, Name))
-> ((Name, [Name]) -> Name) -> (Name, [Name]) -> (Name, Name)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [Name]) -> Name
forall a b. (a, b) -> a
fst) Maybe (Name, [Name])
mpat) [(Name, Maybe (Name, [Name]))]
args
  let argToLiftedPat =
        [(Name, [Name])] -> Map Name [Name]
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, [Name])] -> Map Name [Name])
-> [(Name, [Name])] -> Map Name [Name]
forall a b. (a -> b) -> a -> b
$ ((Name, Maybe (Name, [Name])) -> Maybe (Name, [Name]))
-> [(Name, Maybe (Name, [Name]))] -> [(Name, [Name])]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(Name
nm, Maybe (Name, [Name])
mpat) -> ((Name, [Name]) -> (Name, [Name]))
-> Maybe (Name, [Name]) -> Maybe (Name, [Name])
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Name
nm,) ([Name] -> (Name, [Name]))
-> ((Name, [Name]) -> [Name]) -> (Name, [Name]) -> (Name, [Name])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [Name]) -> [Name]
forall a b. (a, b) -> b
snd) Maybe (Name, [Name])
mpat) [(Name, Maybe (Name, [Name]))]
args
  let funPats = ((Name, Maybe (Name, [Name])) -> Pat)
-> [(Name, Maybe (Name, [Name]))] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Pat -> ((Name, [Name]) -> Pat) -> Maybe (Name, [Name]) -> Pat
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Pat
WildP (Name -> Pat
VarP (Name -> Pat) -> ((Name, [Name]) -> Name) -> (Name, [Name]) -> Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [Name]) -> Name
forall a b. (a, b) -> a
fst) (Maybe (Name, [Name]) -> Pat)
-> ((Name, Maybe (Name, [Name])) -> Maybe (Name, [Name]))
-> (Name, Maybe (Name, [Name]))
-> Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Maybe (Name, [Name])) -> Maybe (Name, [Name])
forall a b. (a, b) -> b
snd) [(Name, Maybe (Name, [Name]))]
args
  let extraLiftedPats =
        ((Name, Maybe (Name, [Name])) -> [Pat])
-> [(Name, Maybe (Name, [Name]))] -> [[Pat]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
          ( [Pat] -> ((Name, [Name]) -> [Pat]) -> Maybe (Name, [Name]) -> [Pat]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
              (Int -> Pat -> [Pat]
forall a. Int -> a -> [a]
replicate ([String] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
liftedNames) Pat
WildP)
              ((Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> Pat
VarP ([Name] -> [Pat])
-> ((Name, [Name]) -> [Name]) -> (Name, [Name]) -> [Pat]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [Name]) -> [Name]
forall a b. (a, b) -> b
snd)
              (Maybe (Name, [Name]) -> [Pat])
-> ((Name, Maybe (Name, [Name])) -> Maybe (Name, [Name]))
-> (Name, Maybe (Name, [Name]))
-> [Pat]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Maybe (Name, [Name])) -> Maybe (Name, [Name])
forall a b. (a, b) -> b
snd
          )
          [(Name, Maybe (Name, [Name]))]
args
  defaultFieldFunExps <-
    traverse
      (fieldFunExpGen argToFunPat argToLiftedPat)
      fields
  return (funPats, extraLiftedPats, defaultFieldFunExps)

-- | Generate a clause for a unary function on a GADT.
genUnaryOpFieldClause ::
  UnaryOpFieldConfig ->
  [(Type, Kind)] ->
  Int ->
  Int ->
  ConstructorInfo ->
  Q Clause
genUnaryOpFieldClause :: UnaryOpFieldConfig
-> [(Type, Type)] -> Int -> Int -> ConstructorInfo -> Q Clause
genUnaryOpFieldClause
  (UnaryOpFieldConfig {[String]
Int -> [String]
Int
-> Int
-> ConstructorVariant
-> Name
-> [Exp]
-> [Exp]
-> Q (Exp, [Bool])
FieldFunExp
ConstructorVariant
-> Name -> [Exp] -> Int -> Exp -> Exp -> Q (Exp, [Bool])
extraPatNames :: [String]
extraLiftedPatNames :: Int -> [String]
fieldResFun :: ConstructorVariant
-> Name -> [Exp] -> Int -> Exp -> Exp -> Q (Exp, [Bool])
fieldCombineFun :: Int
-> Int
-> ConstructorVariant
-> Name
-> [Exp]
-> [Exp]
-> Q (Exp, [Bool])
fieldFunExp :: FieldFunExp
fieldFunExp :: UnaryOpFieldConfig -> FieldFunExp
fieldCombineFun :: UnaryOpFieldConfig
-> Int
-> Int
-> ConstructorVariant
-> Name
-> [Exp]
-> [Exp]
-> Q (Exp, [Bool])
fieldResFun :: UnaryOpFieldConfig
-> ConstructorVariant
-> Name
-> [Exp]
-> Int
-> Exp
-> Exp
-> Q (Exp, [Bool])
extraLiftedPatNames :: UnaryOpFieldConfig -> Int -> [String]
extraPatNames :: UnaryOpFieldConfig -> [String]
..})
  [(Type, Type)]
argTypes
  Int
totalConNumber
  Int
conIdx
  ConstructorInfo
conInfo = do
    fields <- (Type -> Q Type) -> [Type] -> Q [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Type -> Q Type
resolveTypeSynonyms ([Type] -> Q [Type]) -> [Type] -> Q [Type]
forall a b. (a -> b) -> a -> b
$ ConstructorInfo -> [Type]
constructorFields ConstructorInfo
conInfo
    (funPats, funLiftedPats, defaultFieldFunExps) <-
      funPatAndExps fieldFunExp extraLiftedPatNames argTypes fields
    extraPatNames <- traverse newName extraPatNames
    let extraPatExps = (Name -> Exp) -> [Name] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> Exp
VarE [Name]
extraPatNames
    fieldsPatNames <- replicateM (length fields) $ newName "field"
    let extraPats = (Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> Pat
VarP [Name]
extraPatNames
    fieldPats <- conP (constructorName conInfo) (fmap varP fieldsPatNames)
    let fieldPatExps = (Name -> Exp) -> [Name] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Name -> Exp
VarE [Name]
fieldsPatNames

    fieldResExpsAndArgsUsed <-
      sequence $
        zipWith3
          ( fieldResFun
              (constructorVariant conInfo)
              (constructorName conInfo)
              extraPatExps
          )
          [0 ..]
          fieldPatExps
          defaultFieldFunExps
    let fieldResExps = (Exp, [Bool]) -> Exp
forall a b. (a, b) -> a
fst ((Exp, [Bool]) -> Exp) -> [(Exp, [Bool])] -> [Exp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Exp, [Bool])]
fieldResExpsAndArgsUsed
    let extraArgsUsedByFields = (Exp, [Bool]) -> [Bool]
forall a b. (a, b) -> b
snd ((Exp, [Bool]) -> [Bool]) -> [(Exp, [Bool])] -> [[Bool]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Exp, [Bool])]
fieldResExpsAndArgsUsed

    (resExp, extraArgsUsedByResult) <-
      fieldCombineFun
        totalConNumber
        conIdx
        (constructorVariant conInfo)
        (constructorName conInfo)
        extraPatExps
        fieldResExps
    let resUsedNames = Exp -> Set Name
allUsedNames Exp
resExp
    let extraArgsUsed =
          ([Bool] -> Bool) -> [[Bool]] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ([[Bool]] -> [Bool]) -> [[Bool]] -> [Bool]
forall a b. (a -> b) -> a -> b
$
            [[Bool]] -> [[Bool]]
forall a. [[a]] -> [[a]]
List.transpose ([[Bool]] -> [[Bool]]) -> [[Bool]] -> [[Bool]]
forall a b. (a -> b) -> a -> b
$
              [Bool]
extraArgsUsedByResult [Bool] -> [[Bool]] -> [[Bool]]
forall a. a -> [a] -> [a]
: [[Bool]]
extraArgsUsedByFields
    let extraArgsPats =
          (Pat -> Bool -> Pat) -> [Pat] -> [Bool] -> [Pat]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
            (\Pat
pat Bool
used -> if Bool
used then Pat
pat else Pat
WildP)
            [Pat]
extraPats
            [Bool]
extraArgsUsed
    let transformPat (VarP Name
nm) =
          if Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Name
nm Set Name
resUsedNames then Name -> Pat
VarP Name
nm else Pat
WildP
        transformPat Pat
p = Pat
p
    return $
      Clause
        ( fmap transformPat $
            concat (zipWith (:) funPats funLiftedPats)
              ++ extraArgsPats
              ++ [fieldPats]
        )
        (NormalB resExp)
        []

-- | Configuration for a unary operation type class generation on a GADT.
data UnaryOpClassConfig = UnaryOpClassConfig
  { UnaryOpClassConfig -> [UnaryOpConfig]
unaryOpConfigs :: [UnaryOpConfig],
    UnaryOpClassConfig -> [Name]
unaryOpInstanceNames :: [Name],
    UnaryOpClassConfig -> Maybe [Name]
unaryOpContextNames :: Maybe [Name],
    UnaryOpClassConfig -> DeriveConfig -> Q [(Type, Type)]
unaryOpExtraVars :: DeriveConfig -> Q [(Type, Kind)],
    UnaryOpClassConfig
-> DeriveConfig
-> [(Type, Type)]
-> [(Type, Type)]
-> Name
-> Q Type
unaryOpInstanceTypeFromConfig ::
      DeriveConfig ->
      [(Type, Kind)] ->
      [(Type, Kind)] ->
      Name ->
      Q Type,
    UnaryOpClassConfig -> Bool
unaryOpAllowExistential :: Bool
  }

-- | Default unary operation instance type generator.
defaultUnaryOpInstanceTypeFromConfig ::
  DeriveConfig -> [(Type, Kind)] -> [(Type, Kind)] -> Name -> Q Type
defaultUnaryOpInstanceTypeFromConfig :: DeriveConfig -> [(Type, Type)] -> [(Type, Type)] -> Name -> Q Type
defaultUnaryOpInstanceTypeFromConfig DeriveConfig
_ [(Type, Type)]
_ [(Type, Type)]
_ = Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT

-- | Configuration for the derivation rules for a unary operation that can be
-- derived by transforming each field and then combining the results.
data UnaryOpFieldConfig = UnaryOpFieldConfig
  { UnaryOpFieldConfig -> [String]
extraPatNames :: [String],
    UnaryOpFieldConfig -> Int -> [String]
extraLiftedPatNames :: Int -> [String],
    UnaryOpFieldConfig
-> ConstructorVariant
-> Name
-> [Exp]
-> Int
-> Exp
-> Exp
-> Q (Exp, [Bool])
fieldResFun ::
      ConstructorVariant ->
      Name ->
      [Exp] ->
      Int ->
      Exp ->
      Exp ->
      Q (Exp, [Bool]),
    UnaryOpFieldConfig
-> Int
-> Int
-> ConstructorVariant
-> Name
-> [Exp]
-> [Exp]
-> Q (Exp, [Bool])
fieldCombineFun ::
      -- \| Total number of constructors
      Int ->
      -- \| Constructor index
      Int ->
      -- \| Constructor variant
      ConstructorVariant ->
      -- \| Constructor name
      Name ->
      -- \| Extra pattern expressions
      [Exp] ->
      -- \| Field result expressions
      [Exp] ->
      Q (Exp, [Bool]),
    UnaryOpFieldConfig -> FieldFunExp
fieldFunExp :: FieldFunExp
  }

-- | Configuration for the derivation rules for a unary operation.
class UnaryOpFunConfig config where
  genUnaryOpFun ::
    -- | Derive configuration
    DeriveConfig ->
    -- | Configuration
    config ->
    -- | Function names
    [Name] ->
    -- | Number of functor arguments to the class
    Int ->
    -- | Extra variables
    [(Type, Kind)] ->
    -- | Kept variables
    [(Type, Kind)] ->
    -- | Argument variables
    [(Type, Kind)] ->
    -- | Whether the variable is used in fields
    (Name -> Bool) ->
    -- | Constructor infos
    [ConstructorInfo] ->
    Q Dec

instance UnaryOpFunConfig UnaryOpFieldConfig where
  genUnaryOpFun :: DeriveConfig
-> UnaryOpFieldConfig
-> [Name]
-> Int
-> [(Type, Type)]
-> [(Type, Type)]
-> [(Type, Type)]
-> (Name -> Bool)
-> [ConstructorInfo]
-> Q Dec
genUnaryOpFun DeriveConfig
_ UnaryOpFieldConfig
_ [Name]
funNames Int
n [(Type, Type)]
_ [(Type, Type)]
_ [(Type, Type)]
_ Name -> Bool
_ [] =
    Name -> [Q Clause] -> Q Dec
forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
funD ([Name]
funNames [Name] -> Int -> Name
forall a. HasCallStack => [a] -> Int -> a
!! Int
n) [[Q Pat] -> Q Body -> [Q Dec] -> Q Clause
forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
clause [] (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|error "impossible"|]) []]
  genUnaryOpFun DeriveConfig
_ UnaryOpFieldConfig
config [Name]
funNames Int
n [(Type, Type)]
_ [(Type, Type)]
_ [(Type, Type)]
argTypes Name -> Bool
_ [ConstructorInfo]
constructors = do
    clauses <-
      (Int -> ConstructorInfo -> Q Clause)
-> [Int] -> [ConstructorInfo] -> Q [Clause]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM
        ( UnaryOpFieldConfig
-> [(Type, Type)] -> Int -> Int -> ConstructorInfo -> Q Clause
genUnaryOpFieldClause
            UnaryOpFieldConfig
config
            [(Type, Type)]
argTypes
            ([ConstructorInfo] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ConstructorInfo]
constructors)
        )
        [Int
0 ..]
        [ConstructorInfo]
constructors
    let instanceFunName = [Name]
funNames [Name] -> Int -> Name
forall a. HasCallStack => [a] -> Int -> a
!! Int
n
    return $ FunD instanceFunName clauses

-- | Generate a unary operation type class instance for a data type.
genUnaryOpClass ::
  DeriveConfig ->
  UnaryOpClassConfig ->
  Int ->
  Name ->
  Q [Dec]
genUnaryOpClass :: DeriveConfig -> UnaryOpClassConfig -> Int -> Name -> Q [Dec]
genUnaryOpClass DeriveConfig
deriveConfig (UnaryOpClassConfig {Bool
[Name]
[UnaryOpConfig]
Maybe [Name]
DeriveConfig -> Q [(Type, Type)]
DeriveConfig -> [(Type, Type)] -> [(Type, Type)] -> Name -> Q Type
unaryOpConfigs :: UnaryOpClassConfig -> [UnaryOpConfig]
unaryOpInstanceNames :: UnaryOpClassConfig -> [Name]
unaryOpContextNames :: UnaryOpClassConfig -> Maybe [Name]
unaryOpExtraVars :: UnaryOpClassConfig -> DeriveConfig -> Q [(Type, Type)]
unaryOpInstanceTypeFromConfig :: UnaryOpClassConfig
-> DeriveConfig
-> [(Type, Type)]
-> [(Type, Type)]
-> Name
-> Q Type
unaryOpAllowExistential :: UnaryOpClassConfig -> Bool
unaryOpConfigs :: [UnaryOpConfig]
unaryOpInstanceNames :: [Name]
unaryOpContextNames :: Maybe [Name]
unaryOpExtraVars :: DeriveConfig -> Q [(Type, Type)]
unaryOpInstanceTypeFromConfig :: DeriveConfig -> [(Type, Type)] -> [(Type, Type)] -> Name -> Q Type
unaryOpAllowExistential :: Bool
..}) Int
n Name
typName = do
  result@CheckArgsResult {..} <-
    [(Int, EvalModeTag)] -> CheckArgsResult -> Q CheckArgsResult
specializeResult (DeriveConfig -> [(Int, EvalModeTag)]
evalModeSpecializeList DeriveConfig
deriveConfig)
      (CheckArgsResult -> Q CheckArgsResult)
-> Q CheckArgsResult -> Q CheckArgsResult
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Bool -> CheckArgsResult -> Q CheckArgsResult
freshenCheckArgsResult Bool
True
      (CheckArgsResult -> Q CheckArgsResult)
-> Q CheckArgsResult -> Q CheckArgsResult
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> Int -> Name -> Bool -> Int -> Q CheckArgsResult
checkArgs
        (Name -> String
nameBase (Name -> String) -> Name -> String
forall a b. (a -> b) -> a -> b
$ [Name] -> Name
forall a. HasCallStack => [a] -> a
head [Name]
unaryOpInstanceNames)
        ([Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
unaryOpInstanceNames Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        Name
typName
        Bool
unaryOpAllowExistential
        Int
n
  extraVars <- unaryOpExtraVars deriveConfig

  let isTypeUsedInFields (VarT Name
nm) = CheckArgsResult -> Name -> Bool
isVarUsedInFields CheckArgsResult
result Name
nm
      isTypeUsedInFields Type
_ = Bool
False
  contextInstanceTypes <-
    traverse
      (unaryOpInstanceTypeFromConfig deriveConfig extraVars keptVars)
      (fromMaybe unaryOpInstanceNames unaryOpContextNames)
  ctxs <-
    traverse (uncurry $ ctxForVar contextInstanceTypes) $
      filter (isTypeUsedInFields . fst) $
        fmap snd $
          filter (not . (`elem` unconstrainedPositions deriveConfig) . fst) $
            zip [0 ..] keptVars
  let keptType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT (Name -> Type
ConT Name
typName) ([Type] -> Type) -> [Type] -> Type
forall a b. (a -> b) -> a -> b
$ ((Type, Type) -> Type) -> [(Type, Type)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Type, Type) -> Type
forall a b. (a, b) -> a
fst [(Type, Type)]
keptVars
  instanceFuns <-
    traverse
      ( \(UnaryOpConfig config
config [Name]
funNames) ->
          DeriveConfig
-> config
-> [Name]
-> Int
-> [(Type, Type)]
-> [(Type, Type)]
-> [(Type, Type)]
-> (Name -> Bool)
-> [ConstructorInfo]
-> Q Dec
forall config.
UnaryOpFunConfig config =>
DeriveConfig
-> config
-> [Name]
-> Int
-> [(Type, Type)]
-> [(Type, Type)]
-> [(Type, Type)]
-> (Name -> Bool)
-> [ConstructorInfo]
-> Q Dec
genUnaryOpFun
            DeriveConfig
deriveConfig
            config
config
            [Name]
funNames
            Int
n
            [(Type, Type)]
extraVars
            [(Type, Type)]
keptVars
            [(Type, Type)]
argVars
            (CheckArgsResult -> Name -> Bool
isVarUsedInFields CheckArgsResult
result)
            [ConstructorInfo]
constructors
      )
      unaryOpConfigs
  let instanceName = [Name]
unaryOpInstanceNames [Name] -> Int -> Name
forall a. HasCallStack => [a] -> Int -> a
!! Int
n
  instanceTypes <-
    traverse
      (unaryOpInstanceTypeFromConfig deriveConfig extraVars keptVars)
      unaryOpInstanceNames
  let instanceType = Type -> Type -> Type
AppT ([Type]
instanceTypes [Type] -> Int -> Type
forall a. HasCallStack => [a] -> Int -> a
!! Int
n) Type
keptType
  extraPreds <-
    extraConstraint
      deriveConfig
      typName
      instanceName
      extraVars
      keptVars
      constructors
  return
    [ InstanceD
        Nothing
        ( extraPreds
            ++ if null constructors
              then []
              else catMaybes ctxs
        )
        instanceType
        instanceFuns
    ]