{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module      :   Grisette.Internal.TH.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.UnifiedConstructor
  ( mkUnifiedConstructor,
    mkUnifiedConstructor',
  )
where

import Control.Monad (join, replicateM, when, zipWithM)
import Grisette.Internal.TH.Util (constructorInfoToType, occName, putHaddock)
import Grisette.Unified.Internal.EvalModeTag (EvalModeTag)
import Grisette.Unified.Internal.UnifiedData
  ( GetData,
    UnifiedData,
    wrapData,
  )
import Language.Haskell.TH (pprint)
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.
--
-- 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)@.
mkUnifiedConstructor ::
  -- | Prefix for generated wrappers
  String ->
  -- | The type to generate the wrappers for
  Name ->
  Q [Dec]
mkUnifiedConstructor :: String -> Name -> Q [Dec]
mkUnifiedConstructor String
prefix Name
typName = do
  DatatypeInfo
d <- Name -> Q DatatypeInfo
reifyDatatype Name
typName
  let constructorNames :: [String]
constructorNames = Name -> String
occName (Name -> String)
-> (ConstructorInfo -> Name) -> ConstructorInfo -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConstructorInfo -> Name
constructorName (ConstructorInfo -> String) -> [ConstructorInfo] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
d
  [String] -> Name -> Q [Dec]
mkUnifiedConstructor' ((String
prefix String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> String) -> [String] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
constructorNames) Name
typName

-- | 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)@.
mkUnifiedConstructor' ::
  -- | Names for generated wrappers
  [String] ->
  -- | The type to generate the wrappers for
  Name ->
  Q [Dec]
mkUnifiedConstructor' :: [String] -> Name -> Q [Dec]
mkUnifiedConstructor' [String]
names Name
typName = do
  DatatypeInfo
d <- Name -> Q DatatypeInfo
reifyDatatype Name
typName
  let constructors :: [ConstructorInfo]
constructors = DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
d
  Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([String] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
names Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [ConstructorInfo] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ConstructorInfo]
constructors) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
    String -> Q ()
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Number of names does not match the number of constructors"
  let modeVars :: [TyVarBndr_ ()]
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 [TyVarBndr_ ()]
modeVars of
    [TyVarBndr_ ()
mode] -> do
      [[Dec]]
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
          (DatatypeInfo
-> Maybe TyVarBndrSpec
-> Type
-> String
-> ConstructorInfo
-> Q [Dec]
mkSingleWrapper 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
      [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> Q [Dec]) -> [Dec] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ [[Dec]] -> [Dec]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join [[Dec]]
ds
    [] -> do
      Name
n <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"mode"
      let newBndr :: TyVarBndrSpec
newBndr = Name -> Type -> TyVarBndrSpec
kindedTVSpecified Name
n (Name -> Type
ConT ''EvalModeTag)
      [[Dec]]
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
          (DatatypeInfo
-> Maybe TyVarBndrSpec
-> Type
-> String
-> ConstructorInfo
-> Q [Dec]
mkSingleWrapper DatatypeInfo
d (TyVarBndrSpec -> Maybe TyVarBndrSpec
forall a. a -> Maybe a
Just TyVarBndrSpec
newBndr) (Name -> Type
VarT Name
n))
          [String]
names
          [ConstructorInfo]
constructors
      [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec] -> Q [Dec]) -> [Dec] -> Q [Dec]
forall a b. (a -> b) -> a -> b
$ [[Dec]] -> [Dec]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join [[Dec]]
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
  ([Type]
pred, Type
ret) <- Type -> Type -> Q ([Type], Type)
augmentFinalType Type
mode Type
t
  ([Type], Type) -> Q ([Type], Type)
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Type]
pred, Type -> Type -> Type
AppT Type
a Type
ret)
augmentFinalType Type
mode Type
t = do
  Type
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)|]
  Type
predu <- [t|UnifiedData $(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)|]
  ([Type], Type) -> Q ([Type], Type)
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Type
predu], Type
r)

augmentConstructorType :: Maybe TyVarBndrSpec -> Type -> Type -> Q Type
augmentConstructorType :: Maybe TyVarBndrSpec -> Type -> Type -> Q Type
augmentConstructorType Maybe TyVarBndrSpec
modeBndr Type
mode (ForallT [TyVarBndrSpec]
tybinders [Type]
ctx Type
ty1) = do
  ([Type]
preds, Type
augmentedTyp) <- Type -> Type -> Q ([Type], Type)
augmentFinalType Type
mode Type
ty1
  case Maybe TyVarBndrSpec
modeBndr 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 TyVarBndrSpec -> [TyVarBndrSpec] -> [TyVarBndrSpec]
forall a. a -> [a] -> [a]
: [TyVarBndrSpec]
tybinders) ([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]
preds [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
ctx) Type
augmentedTyp
augmentConstructorType Maybe TyVarBndrSpec
modeBndr Type
mode Type
ty = do
  ([Type]
preds, Type
augmentedTyp) <- Type -> Type -> Q ([Type], Type)
augmentFinalType Type
mode Type
ty
  case Maybe TyVarBndrSpec
modeBndr 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
  [Name]
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 :: [Q Pat]
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
  [Q Pat] -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => [m Pat] -> m Exp -> m Exp
lamE
    [Q Pat]
args
    ( ( Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE
          (Q Exp -> Q Type -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Type -> m Exp
appTypeE [|wrapData|] (Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
mode))
          ((Q Exp -> Q Exp -> Q Exp) -> Q Exp -> [Q 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 -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
f) ((Name -> Q Exp) -> [Name] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
xs))
      )
    )

mkSingleWrapper :: DatatypeInfo -> Maybe TyVarBndrSpec -> Type -> String -> ConstructorInfo -> Q [Dec]
mkSingleWrapper :: DatatypeInfo
-> Maybe TyVarBndrSpec
-> Type
-> String
-> ConstructorInfo
-> Q [Dec]
mkSingleWrapper DatatypeInfo
dataType Maybe TyVarBndrSpec
modeBndr Type
mode String
name ConstructorInfo
info = do
  Type
constructorTyp <- DatatypeInfo -> ConstructorInfo -> Q Type
constructorInfoToType DatatypeInfo
dataType ConstructorInfo
info
  Type
augmentedTyp <- Maybe TyVarBndrSpec -> Type -> Type -> Q Type
augmentConstructorType Maybe TyVarBndrSpec
modeBndr Type
mode Type
constructorTyp
  let oriName :: Name
oriName = ConstructorInfo -> Name
constructorName ConstructorInfo
info
  let retName :: Name
retName = String -> Name
mkName String
name
  Exp
expr <- Type -> Int -> Exp -> Q Exp
augmentExpr Type
mode ([Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Type] -> Int) -> [Type] -> Int
forall a b. (a -> b) -> a -> b
$ ConstructorInfo -> [Type]
constructorFields ConstructorInfo
info) (Name -> Exp
ConE Name
oriName)
  Name -> String -> Q ()
putHaddock Name
retName (String -> Q ()) -> String -> Q ()
forall a b. (a -> b) -> a -> b
$
    String
"Smart constructor for v'"
      String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Name -> String
forall a. Show a => a -> String
show Name
oriName
      String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"' to construct unified value."
  [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return
    [ Name -> Type -> Dec
SigD Name
retName Type
augmentedTyp,
      Name -> [Clause] -> Dec
FunD Name
retName [[Pat] -> Body -> [Dec] -> Clause
Clause [] (Exp -> Body
NormalB Exp
expr) []]
    ]