{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module      :   Grisette.Internal.TH.Ctor.UnifiedConstructor
-- 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.Ctor.UnifiedConstructor
  ( makeUnifiedCtorWith,
    makePrefixedUnifiedCtor,
    makeNamedUnifiedCtor,
    makeUnifiedCtor,
  )
where

import Control.Monad (join, replicateM, when, zipWithM)
import Data.Maybe (catMaybes)
import Grisette.Internal.Core.Data.Class.Mergeable (Mergeable, Mergeable1, Mergeable2)
import Grisette.Internal.TH.Ctor.Common
  ( decapitalizeTransformer,
    prefixTransformer,
    withNameTransformer,
  )
import Grisette.Internal.TH.Derivation.Common (ctxForVar)
import Grisette.Internal.TH.Util (constructorInfoToType, putHaddock, tvIsMode)
import Grisette.Internal.Unified.EvalModeTag (EvalModeTag)
import Grisette.Internal.Unified.UnifiedData
  ( GetData,
    UnifiedData,
    wrapData,
  )
import Language.Haskell.TH (conT, pprint, varT)
import Language.Haskell.TH.Datatype
  ( ConstructorInfo (constructorFields, constructorName),
    DatatypeInfo (datatypeCons, datatypeVars),
    reifyDatatype,
    tvKind,
    tvName,
  )
import Language.Haskell.TH.Datatype.TyVarBndr (TyVarBndrSpec, kindedTVSpecified)
import Language.Haskell.TH.Lib (appE, appTypeE, lamE, varE, varP)
import Language.Haskell.TH.Syntax
  ( Body (NormalB),
    Clause (Clause),
    Dec (FunD, SigD),
    Exp (ConE),
    Name,
    Pred,
    Q,
    Type (AppT, ArrowT, ConT, ForallT, VarT),
    mkName,
    newName,
  )

-- | Generate smart constructors to create unified values with provided name
-- transformer.
--
-- For a type @T mode a b c@ with constructors @T1@, @T2@, etc., this function
-- will generate smart constructors with the name transformed, e.g., given the
-- name transformer @(\name -> "mk" ++ name)@, it will generate @mkT1@, @mkT2@,
-- @mkT2@, etc.
--
-- The generated smart constructors will contruct values of type
-- @GetData mode (T mode a b c)@.
makeUnifiedCtorWith :: [Name] -> (String -> String) -> Name -> Q [Dec]
makeUnifiedCtorWith :: [Name] -> (String -> String) -> Name -> Q [Dec]
makeUnifiedCtorWith = ([String] -> Name -> Q [Dec])
-> (String -> String) -> Name -> Q [Dec]
withNameTransformer (([String] -> Name -> Q [Dec])
 -> (String -> String) -> Name -> Q [Dec])
-> ([Name] -> [String] -> Name -> Q [Dec])
-> [Name]
-> (String -> String)
-> Name
-> Q [Dec]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Name] -> [String] -> Name -> Q [Dec]
makeNamedUnifiedCtor

-- | Generate smart constructors to create unified values.
--
-- For a type @T mode a b c@ with constructors @T1@, @T2@, etc., this function
-- will generate smart constructors with the given prefix, e.g., @mkT1@, @mkT2@,
-- etc.
--
-- The generated smart constructors will contruct values of type
-- @GetData mode (T mode a b c)@.
makePrefixedUnifiedCtor ::
  [Name] ->
  -- | Prefix for generated wrappers
  String ->
  -- | The type to generate the wrappers for
  Name ->
  Q [Dec]
makePrefixedUnifiedCtor :: [Name] -> String -> Name -> Q [Dec]
makePrefixedUnifiedCtor [Name]
modeCtx =
  [Name] -> (String -> String) -> Name -> Q [Dec]
makeUnifiedCtorWith [Name]
modeCtx ((String -> String) -> Name -> Q [Dec])
-> (String -> String -> String) -> String -> Name -> Q [Dec]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String -> String
prefixTransformer

-- | Generate smart constructors to create unified values.
--
-- For a type @T mode a b c@ with constructors @T1@, @T2@, etc., this function
-- will generate smart constructors with the names decapitalized, e.g.,
-- @t1@, @t2@, etc.
--
-- The generated smart constructors will contruct values of type
-- @GetData mode (T mode a b c)@.
makeUnifiedCtor ::
  [Name] ->
  -- | The type to generate the wrappers for
  Name ->
  Q [Dec]
makeUnifiedCtor :: [Name] -> Name -> Q [Dec]
makeUnifiedCtor [Name]
modeCtx = [Name] -> (String -> String) -> Name -> Q [Dec]
makeUnifiedCtorWith [Name]
modeCtx String -> String
decapitalizeTransformer

-- | Generate smart constructors to create unified values.
--
-- For a type @T mode a b c@ with constructors @T1@, @T2@, etc., this function
-- will generate smart constructors with the given names.
--
-- The generated smart constructors will contruct values of type
-- @GetData mode (T mode a b c)@.
makeNamedUnifiedCtor ::
  [Name] ->
  -- | Names for generated wrappers
  [String] ->
  -- | The type to generate the wrappers for
  Name ->
  Q [Dec]
makeNamedUnifiedCtor :: [Name] -> [String] -> Name -> Q [Dec]
makeNamedUnifiedCtor [Name]
modeCtx [String]
names Name
typName = do
  d <- Name -> Q DatatypeInfo
reifyDatatype Name
typName
  let constructors = DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
d
  when (length names /= length constructors) $
    fail "Number of names does not match the number of constructors"
  let modeVars = (TyVarBndr_ () -> Bool) -> [TyVarBndr_ ()] -> [TyVarBndr_ ()]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Name -> Type
ConT ''EvalModeTag) (Type -> Bool) -> (TyVarBndr_ () -> Type) -> TyVarBndr_ () -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVarBndr_ () -> Type
forall flag. TyVarBndr_ flag -> Type
tvKind) (DatatypeInfo -> [TyVarBndr_ ()]
datatypeVars DatatypeInfo
d)
  -- when (length modeVars /= 1) $
  --  fail "Expected exactly one EvalModeTag variable in the datatype."
  case modeVars of
    [TyVarBndr_ ()
mode] -> do
      ds <-
        (String -> ConstructorInfo -> Q [Dec])
-> [String] -> [ConstructorInfo] -> Q [[Dec]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM
          ([Name]
-> DatatypeInfo
-> Maybe TyVarBndrSpec
-> Type
-> String
-> ConstructorInfo
-> Q [Dec]
mkSingleWrapper [Name]
modeCtx DatatypeInfo
d Maybe TyVarBndrSpec
forall a. Maybe a
Nothing (Type -> String -> ConstructorInfo -> Q [Dec])
-> Type -> String -> ConstructorInfo -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ Name -> Type
VarT (Name -> Type) -> Name -> Type
forall a b. (a -> b) -> a -> b
$ TyVarBndr_ () -> Name
forall flag. TyVarBndr_ flag -> Name
tvName TyVarBndr_ ()
mode)
          [String]
names
          [ConstructorInfo]
constructors
      return $ join ds
    [] -> do
      n <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"mode"
      let newBndr = Name -> Type -> TyVarBndrSpec
kindedTVSpecified Name
n (Name -> Type
ConT ''EvalModeTag)
      ds <-
        zipWithM
          (mkSingleWrapper modeCtx d (Just newBndr) (VarT n))
          names
          constructors
      return $ join ds
    [TyVarBndr_ ()]
_ -> String -> Q [Dec]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Expected one or zero EvalModeTag variable in the datatype."

augmentFinalType :: Type -> Type -> Q ([Pred], Type)
augmentFinalType :: Type -> Type -> Q ([Type], Type)
augmentFinalType Type
mode (AppT a :: Type
a@(AppT Type
ArrowT Type
_) Type
t) = do
  (pred, ret) <- Type -> Type -> Q ([Type], Type)
augmentFinalType Type
mode Type
t
  return (pred, AppT a ret)
augmentFinalType Type
mode Type
t = do
  r <- [t|GetData $(Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
mode) $(Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t)|]
  predu <- [t|UnifiedData $(return mode) $(return t)|]
  return ([predu], r)

augmentConstructorType ::
  [Name] -> Maybe TyVarBndrSpec -> Type -> Type -> Q Type
augmentConstructorType :: [Name] -> Maybe TyVarBndrSpec -> Type -> Type -> Q Type
augmentConstructorType
  [Name]
modeCtx
  Maybe TyVarBndrSpec
freshModeBndr
  Type
mode
  (ForallT [TyVarBndrSpec]
tybinders [Type]
ctx Type
ty1) = do
    (preds, augmentedTyp) <- Type -> Type -> Q ([Type], Type)
augmentFinalType Type
mode Type
ty1
    let modeBndrsInForall = (TyVarBndrSpec -> Bool) -> [TyVarBndrSpec] -> [TyVarBndrSpec]
forall a. (a -> Bool) -> [a] -> [a]
filter TyVarBndrSpec -> Bool
forall flag. TyVarBndr_ flag -> Bool
tvIsMode [TyVarBndrSpec]
tybinders
    mergeablePreds <-
      catMaybes
        <$> traverse
          ( \TyVarBndrSpec
bndr ->
              [Type] -> Type -> Type -> Q (Maybe Type)
ctxForVar
                (Name -> Type
ConT (Name -> Type) -> [Name] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [''Mergeable, ''Mergeable1, ''Mergeable2])
                (Name -> Type
VarT (Name -> Type) -> Name -> Type
forall a b. (a -> b) -> a -> b
$ TyVarBndrSpec -> Name
forall flag. TyVarBndr_ flag -> Name
tvName TyVarBndrSpec
bndr)
                (TyVarBndrSpec -> Type
forall flag. TyVarBndr_ flag -> Type
tvKind TyVarBndrSpec
bndr)
          )
          tybinders
    modePred <-
      case (modeBndrsInForall, freshModeBndr) of
        ([TyVarBndrSpec
bndr], Maybe TyVarBndrSpec
Nothing) ->
          (Name -> Q Type) -> [Name] -> Q [Type]
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 (\Name
nm -> [t|$(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
nm) $(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT (Name -> Q Type) -> Name -> Q Type
forall a b. (a -> b) -> a -> b
$ TyVarBndrSpec -> Name
forall flag. TyVarBndr_ flag -> Name
tvName TyVarBndrSpec
bndr)|]) [Name]
modeCtx
        ([], Just TyVarBndrSpec
bndr) ->
          (Name -> Q Type) -> [Name] -> Q [Type]
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 (\Name
nm -> [t|$(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
nm) $(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT (Name -> Q Type) -> Name -> Q Type
forall a b. (a -> b) -> a -> b
$ TyVarBndrSpec -> Name
forall flag. TyVarBndr_ flag -> Name
tvName TyVarBndrSpec
bndr)|]) [Name]
modeCtx
        ([TyVarBndrSpec], Maybe TyVarBndrSpec)
_ -> String -> Q [Type]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Unsupported constructor type."
    case freshModeBndr of
      Just TyVarBndrSpec
bndr -> do
        Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Q Type) -> Type -> Q Type
forall a b. (a -> b) -> a -> b
$
          [TyVarBndrSpec] -> [Type] -> Type -> Type
ForallT
            (TyVarBndrSpec
bndr TyVarBndrSpec -> [TyVarBndrSpec] -> [TyVarBndrSpec]
forall a. a -> [a] -> [a]
: [TyVarBndrSpec]
tybinders)
            ([Type]
modePred [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
mergeablePreds [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
preds [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
ctx)
            Type
augmentedTyp
      Maybe TyVarBndrSpec
Nothing ->
        Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Q Type) -> Type -> Q Type
forall a b. (a -> b) -> a -> b
$
          [TyVarBndrSpec] -> [Type] -> Type -> Type
ForallT
            [TyVarBndrSpec]
tybinders
            ([Type]
modePred [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
mergeablePreds [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
preds [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
ctx)
            Type
augmentedTyp
augmentConstructorType [Name]
_ Maybe TyVarBndrSpec
freshModeBndr Type
mode Type
ty = do
  (preds, augmentedTyp) <- Type -> Type -> Q ([Type], Type)
augmentFinalType Type
mode Type
ty
  case freshModeBndr of
    Just TyVarBndrSpec
bndr -> Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Q Type) -> Type -> Q Type
forall a b. (a -> b) -> a -> b
$ [TyVarBndrSpec] -> [Type] -> Type -> Type
ForallT [TyVarBndrSpec
bndr] ([Type]
preds) Type
augmentedTyp
    Maybe TyVarBndrSpec
Nothing ->
      String -> Q Type
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Type) -> String -> Q Type
forall a b. (a -> b) -> a -> b
$
        String
"augmentConstructorType: unsupported constructor type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Ppr a => a -> String
pprint Type
ty

augmentExpr :: Type -> Int -> Exp -> Q Exp
augmentExpr :: Type -> Int -> Exp -> Q Exp
augmentExpr Type
mode Int
n Exp
f = do
  xs <- Int -> Q Name -> Q [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"x")
  let args = (Name -> Q Pat) -> [Name] -> [Q Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP [Name]
xs
  lamE
    args
    ( ( appE
          (appTypeE [|wrapData|] (return mode))
          (foldl appE (return f) (map varE xs))
      )
    )

mkSingleWrapper ::
  [Name] ->
  DatatypeInfo ->
  Maybe TyVarBndrSpec ->
  Type ->
  String ->
  ConstructorInfo ->
  Q [Dec]
mkSingleWrapper :: [Name]
-> DatatypeInfo
-> Maybe TyVarBndrSpec
-> Type
-> String
-> ConstructorInfo
-> Q [Dec]
mkSingleWrapper [Name]
modeCtx DatatypeInfo
dataType Maybe TyVarBndrSpec
freshModeBndr Type
mode String
name ConstructorInfo
info = do
  constructorTyp <- DatatypeInfo -> ConstructorInfo -> Q Type
constructorInfoToType DatatypeInfo
dataType ConstructorInfo
info
  augmentedTyp <-
    augmentConstructorType modeCtx freshModeBndr mode constructorTyp
  let oriName = ConstructorInfo -> Name
constructorName ConstructorInfo
info
  let retName = String -> Name
mkName String
name
  expr <- augmentExpr mode (length $ constructorFields info) (ConE oriName)
  putHaddock retName $
    "Smart constructor for v'"
      <> show oriName
      <> "' to construct unified value."
  return
    [ SigD retName augmentedTyp,
      FunD retName [Clause [] (NormalB expr) []]
    ]