{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE Unsafe #-}
module Data.Eliminator.TH (
deriveElim
, deriveElimNamed
) where
import Control.Applicative
import Control.Monad
import Data.Char (isUpper)
import Data.Foldable
import qualified Data.Kind as Kind (Type)
import Data.Maybe
import Data.Singletons.Prelude
import Language.Haskell.TH
import Language.Haskell.TH.Datatype
import Language.Haskell.TH.Desugar hiding (NewOrData(..))
deriveElim :: Name -> Q [Dec]
deriveElim :: Name -> Q [Dec]
deriveElim dataName :: Name
dataName = String -> Name -> Q [Dec]
deriveElimNamed (Name -> String
eliminatorName Name
dataName) Name
dataName
deriveElimNamed :: String -> Name -> Q [Dec]
deriveElimNamed :: String -> Name -> Q [Dec]
deriveElimNamed funName :: String
funName dataName :: Name
dataName = do
info :: DatatypeInfo
info@(DatatypeInfo { datatypeVars :: DatatypeInfo -> [TyVarBndr]
datatypeVars = [TyVarBndr]
dataVarBndrs
, datatypeVariant :: DatatypeInfo -> DatatypeVariant
datatypeVariant = DatatypeVariant
variant
, datatypeCons :: DatatypeInfo -> [ConstructorInfo]
datatypeCons = [ConstructorInfo]
cons
}) <- Name -> Q DatatypeInfo
reifyDatatype Name
dataName
let noDataFamilies :: Q a
noDataFamilies =
String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Eliminators for data family instances are currently not supported"
case DatatypeVariant
variant of
DataInstance -> Q ()
forall a. Q a
noDataFamilies
NewtypeInstance -> Q ()
forall a. Q a
noDataFamilies
Datatype -> () -> Q ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
Newtype -> () -> Q ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
Name
predVar <- String -> Q Name
newName "p"
Name
singVar <- String -> Q Name
newName "s"
let elimN :: Name
elimN = String -> Name
mkName String
funName
promDataKind :: Type
promDataKind = DatatypeInfo -> Type
datatypeType DatatypeInfo
info
predVarBndr :: TyVarBndr
predVarBndr = Name -> Type -> TyVarBndr
KindedTV Name
predVar (Type -> Name -> Type -> Type
InfixT Type
promDataKind ''(~>) (Name -> Type
ConT ''Kind.Type))
singVarBndr :: TyVarBndr
singVarBndr = Name -> Type -> TyVarBndr
KindedTV Name
singVar Type
promDataKind
[Type]
caseTypes <- (ConstructorInfo -> Q Type) -> [ConstructorInfo] -> Q [Type]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Name -> Name -> ConstructorInfo -> Q Type
caseType Name
dataName Name
predVar) [ConstructorInfo]
cons
let returnType :: Type
returnType = Name -> Type -> Type
predType Name
predVar (Name -> Type
VarT Name
singVar)
bndrsPrefix :: [TyVarBndr]
bndrsPrefix = [TyVarBndr]
dataVarBndrs [TyVarBndr] -> [TyVarBndr] -> [TyVarBndr]
forall a. [a] -> [a] -> [a]
++ [TyVarBndr
predVarBndr]
allBndrs :: [TyVarBndr]
allBndrs = [TyVarBndr]
bndrsPrefix [TyVarBndr] -> [TyVarBndr] -> [TyVarBndr]
forall a. [a] -> [a] -> [a]
++ [TyVarBndr
singVarBndr]
elimType :: Type
elimType = [TyVarBndr] -> [Type] -> Type -> Type
ForallT [TyVarBndr]
allBndrs []
([Type] -> Type -> Type
ravel (Name -> Type
singType Name
singVarType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
caseTypes) Type
returnType)
qelimDef :: Q Dec
qelimDef
| [ConstructorInfo] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ConstructorInfo]
cons
= do Name
singVal <- String -> Q Name
newName "singVal"
Dec -> Q Dec
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$ Name -> [Clause] -> Dec
FunD Name
elimN [[Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
singVal] (Exp -> Body
NormalB (Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE Name
singVal) [])) []]
| Bool
otherwise
= do [Clause]
caseClauses
<- (Int -> ConstructorInfo -> Q Clause)
-> [ConstructorInfo] -> Q [Clause]
forall (f :: * -> *) a b.
Applicative f =>
(Int -> a -> f b) -> [a] -> f [b]
itraverse (\i :: Int
i -> Name -> Name -> [Name] -> Int -> Int -> ConstructorInfo -> Q Clause
caseClause Name
dataName Name
elimN
((TyVarBndr -> Name) -> [TyVarBndr] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr -> Name
tyVarBndrName [TyVarBndr]
bndrsPrefix)
Int
i ([ConstructorInfo] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [ConstructorInfo]
cons)) [ConstructorInfo]
cons
Dec -> Q Dec
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$ Name -> [Clause] -> Dec
FunD Name
elimN [Clause]
caseClauses
Dec
elimDef <- Q Dec
qelimDef
[Dec] -> Q [Dec]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name -> Type -> Dec
SigD Name
elimN Type
elimType, Dec
elimDef]
caseType :: Name -> Name -> ConstructorInfo -> Q Type
caseType :: Name -> Name -> ConstructorInfo -> Q Type
caseType dataName :: Name
dataName predVar :: Name
predVar
(ConstructorInfo { constructorName :: ConstructorInfo -> Name
constructorName = Name
conName
, constructorVars :: ConstructorInfo -> [TyVarBndr]
constructorVars = [TyVarBndr]
conVars
, constructorContext :: ConstructorInfo -> [Type]
constructorContext = [Type]
conContext
, constructorFields :: ConstructorInfo -> [Type]
constructorFields = [Type]
fieldTypes })
= do Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([TyVarBndr] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TyVarBndr]
conVars Bool -> Bool -> Bool
&& [Type] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Type]
conContext) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q ()) -> String -> Q ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
[ "Eliminators for GADTs or datatypes with existentially quantified"
, "data constructors currently not supported"
]
[Name]
vars <- String -> Int -> Q [Name]
newNameList "f" (Int -> Q [Name]) -> Int -> Q [Name]
forall a b. (a -> b) -> a -> b
$ [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
fieldTypes
let returnType :: Type
returnType = Name -> Type -> Type
predType Name
predVar
((Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Type -> Type -> Type
AppT (Name -> Type
ConT Name
conName) ((Name -> Type) -> [Name] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Type
VarT [Name]
vars))
mbInductiveType :: Name -> Type -> Maybe Type
mbInductiveType var :: Name
var varType :: Type
varType =
let inductiveArg :: Type
inductiveArg = Name -> Type -> Type
predType Name
predVar (Name -> Type
VarT Name
var)
in Name -> Type -> Type -> Maybe Type
forall a. Name -> Type -> a -> Maybe a
mbInductiveCase Name
dataName Type
varType Type
inductiveArg
Type -> Q Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> Q Type) -> Type -> Q Type
forall a b. (a -> b) -> a -> b
$ ((Name, Type) -> Type -> Type) -> Type -> [(Name, Type)] -> Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' (\(var :: Name
var, varType :: Type
varType) t :: Type
t ->
[TyVarBndr] -> [Type] -> Type -> Type
ForallT [Name -> Type -> TyVarBndr
KindedTV Name
var Type
varType]
[]
([Type] -> Type -> Type
ravel (Name -> Type
singType Name
varType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:Maybe Type -> [Type]
forall a. Maybe a -> [a]
maybeToList (Name -> Type -> Maybe Type
mbInductiveType Name
var Type
varType)) Type
t))
Type
returnType
([Name] -> [Type] -> [(Name, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
vars [Type]
fieldTypes)
caseClause :: Name -> Name -> [Name] -> Int -> Int
-> ConstructorInfo -> Q Clause
caseClause :: Name -> Name -> [Name] -> Int -> Int -> ConstructorInfo -> Q Clause
caseClause dataName :: Name
dataName elimN :: Name
elimN bndrNamesPrefix :: [Name]
bndrNamesPrefix conIndex :: Int
conIndex numCons :: Int
numCons
(ConstructorInfo { constructorName :: ConstructorInfo -> Name
constructorName = Name
conName
, constructorFields :: ConstructorInfo -> [Type]
constructorFields = [Type]
fieldTypes })
= do let numFields :: Int
numFields = [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
fieldTypes
[Name]
singVars <- String -> Int -> Q [Name]
newNameList "s" Int
numFields
[Name]
singVarSigs <- String -> Int -> Q [Name]
newNameList "sTy" Int
numFields
Name
usedCaseVar <- String -> Q Name
newName "useThis"
[Name]
caseVars <- Int -> (Int -> Q Name) -> Q [Name]
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> f a) -> f [a]
ireplicateA Int
numCons ((Int -> Q Name) -> Q [Name]) -> (Int -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ \i :: Int
i ->
if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
conIndex
then Name -> Q Name
forall (f :: * -> *) a. Applicative f => a -> f a
pure Name
usedCaseVar
else String -> Q Name
newName ("_p" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i)
let singConName :: Name
singConName = Name -> Name
singDataConName Name
conName
mkSingVarPat :: Name -> Name -> Pat
mkSingVarPat var :: Name
var varSig :: Name
varSig = Pat -> Type -> Pat
SigP (Name -> Pat
VarP Name
var) (Name -> Type
singType Name
varSig)
singVarPats :: [Pat]
singVarPats = (Name -> Name -> Pat) -> [Name] -> [Name] -> [Pat]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> Name -> Pat
mkSingVarPat [Name]
singVars [Name]
singVarSigs
mbInductiveArg :: Name -> Name -> Type -> Maybe Exp
mbInductiveArg singVar :: Name
singVar singVarSig :: Name
singVarSig varType :: Type
varType =
let prefix :: Exp
prefix = Exp -> [Type] -> Exp
foldAppType (Name -> Exp
VarE Name
elimN)
([Type] -> Exp) -> [Type] -> Exp
forall a b. (a -> b) -> a -> b
$ (Name -> Type) -> [Name] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Type
VarT [Name]
bndrNamesPrefix
[Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Name -> Type
VarT Name
singVarSig]
inductiveArg :: Exp
inductiveArg = Exp -> [Exp] -> Exp
foldExp Exp
prefix
([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE Name
singVarExp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
:(Name -> Exp) -> [Name] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE [Name]
caseVars
in Name -> Type -> Exp -> Maybe Exp
forall a. Name -> Type -> a -> Maybe a
mbInductiveCase Name
dataName Type
varType Exp
inductiveArg
mkArg :: Exp -> (Name, Name, Type) -> Exp
mkArg f :: Exp
f (singVar :: Name
singVar, singVarSig :: Name
singVarSig, varType :: Type
varType) =
Exp -> [Exp] -> Exp
foldExp Exp
f ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE Name
singVar
Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: Maybe Exp -> [Exp]
forall a. Maybe a -> [a]
maybeToList (Name -> Name -> Type -> Maybe Exp
mbInductiveArg Name
singVar Name
singVarSig Type
varType)
rhs :: Exp
rhs = (Exp -> (Name, Name, Type) -> Exp)
-> Exp -> [(Name, Name, Type)] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Exp -> (Name, Name, Type) -> Exp
mkArg (Name -> Exp
VarE Name
usedCaseVar) ([(Name, Name, Type)] -> Exp) -> [(Name, Name, Type)] -> Exp
forall a b. (a -> b) -> a -> b
$
[Name] -> [Name] -> [Type] -> [(Name, Name, Type)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 [Name]
singVars [Name]
singVarSigs [Type]
fieldTypes
Clause -> Q Clause
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Clause -> Q Clause) -> Clause -> Q Clause
forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause (Name -> [Pat] -> Pat
ConP Name
singConName [Pat]
singVarPats Pat -> [Pat] -> [Pat]
forall a. a -> [a] -> [a]
: (Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
caseVars)
(Exp -> Body
NormalB Exp
rhs)
[]
mbInductiveCase :: Name -> Type -> a -> Maybe a
mbInductiveCase :: Name -> Type -> a -> Maybe a
mbInductiveCase dataName :: Name
dataName varType :: Type
varType inductiveArg :: a
inductiveArg
= case Type -> (Type, [TypeArg])
unfoldType Type
varType of
(headTy :: Type
headTy, _)
| Type
ListT <- Type
headTy
, Name
dataName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== ''[]
-> a -> Maybe a
forall a. a -> Maybe a
Just a
inductiveArg
| ConT n :: Name
n <- Type
headTy
, Name
dataName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n
-> a -> Maybe a
forall a. a -> Maybe a
Just a
inductiveArg
| Bool
otherwise
-> Maybe a
forall a. Maybe a
Nothing
singType :: Name -> Type
singType :: Name -> Type
singType x :: Name
x = Name -> Type
ConT ''Sing Type -> Type -> Type
`AppT` Name -> Type
VarT Name
x
predType :: Name -> Type -> Type
predType :: Name -> Type -> Type
predType p :: Name
p ty :: Type
ty = Name -> Type
ConT ''Apply Type -> Type -> Type
`AppT` Name -> Type
VarT Name
p Type -> Type -> Type
`AppT` Type
ty
newNameList :: String -> Int -> Q [Name]
newNameList :: String -> Int -> Q [Name]
newNameList prefix :: String
prefix n :: Int
n = Int -> (Int -> Q Name) -> Q [Name]
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> f a) -> f [a]
ireplicateA Int
n ((Int -> Q Name) -> Q [Name]) -> (Int -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ String -> Q Name
newName (String -> Q Name) -> (Int -> String) -> Int -> Q Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
prefix String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> String) -> (Int -> String) -> Int -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String
forall a. Show a => a -> String
show
eliminatorName :: Name -> String
eliminatorName :: Name -> String
eliminatorName n :: Name
n
| first :: Char
first:_ <- String
nStr
, Char -> Bool
isUpper Char
first
= "elim" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
nStr
| Bool
otherwise
= "~>" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
nStr
where
nStr :: String
nStr = Name -> String
nameBase Name
n
ravel :: [Type] -> Type -> Type
ravel :: [Type] -> Type -> Type
ravel [] res :: Type
res = Type
res
ravel (h :: Type
h:t :: [Type]
t) res :: Type
res = Type -> Type -> Type
AppT (Type -> Type -> Type
AppT Type
ArrowT Type
h) ([Type] -> Type -> Type
ravel [Type]
t Type
res)
foldExp :: Exp -> [Exp] -> Exp
foldExp :: Exp -> [Exp] -> Exp
foldExp = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Exp -> Exp -> Exp
AppE
foldAppType :: Exp -> [Type] -> Exp
foldAppType :: Exp -> [Type] -> Exp
foldAppType = (Exp -> Type -> Exp) -> Exp -> [Type] -> Exp
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Exp -> Type -> Exp
AppTypeE
tyVarBndrName :: TyVarBndr -> Name
tyVarBndrName :: TyVarBndr -> Name
tyVarBndrName (PlainTV n :: Name
n) = Name
n
tyVarBndrName (KindedTV n :: Name
n _) = Name
n
itraverse :: Applicative f => (Int -> a -> f b) -> [a] -> f [b]
itraverse :: (Int -> a -> f b) -> [a] -> f [b]
itraverse f :: Int -> a -> f b
f xs0 :: [a]
xs0 = [a] -> Int -> f [b]
go [a]
xs0 0 where
go :: [a] -> Int -> f [b]
go [] _ = [b] -> f [b]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
go (x :: a
x:xs :: [a]
xs) n :: Int
n = (:) (b -> [b] -> [b]) -> f b -> f ([b] -> [b])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> f b
f Int
n a
x f ([b] -> [b]) -> f [b] -> f [b]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ([a] -> Int -> f [b]
go [a]
xs (Int -> f [b]) -> Int -> f [b]
forall a b. (a -> b) -> a -> b
$! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1))
ireplicateA :: Applicative f => Int -> (Int -> f a) -> f [a]
ireplicateA :: Int -> (Int -> f a) -> f [a]
ireplicateA cnt0 :: Int
cnt0 f :: Int -> f a
f =
Int -> Int -> f [a]
forall t. (Ord t, Num t) => t -> Int -> f [a]
loop Int
cnt0 0
where
loop :: t -> Int -> f [a]
loop cnt :: t
cnt n :: Int
n
| t
cnt t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= 0 = [a] -> f [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
| Bool
otherwise = (a -> [a] -> [a]) -> f a -> f [a] -> f [a]
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (:) (Int -> f a
f Int
n) (t -> Int -> f [a]
loop (t
cnt t -> t -> t
forall a. Num a => a -> a -> a
- 1) (Int -> f [a]) -> Int -> f [a]
forall a b. (a -> b) -> a -> b
$! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1))
singDataConName :: Name -> Name
singDataConName :: Name -> Name
singDataConName nm :: Name
nm
| Name
nm Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== '[] = 'SNil
| Name
nm Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== '(:) = 'SCons
| Just degree :: Int
degree <- Name -> Maybe Int
tupleNameDegree_maybe Name
nm = Int -> Name
mkTupleDataName Int
degree
| Just degree :: Int
degree <- Name -> Maybe Int
unboxedTupleNameDegree_maybe Name
nm = Int -> Name
mkTupleDataName Int
degree
| Bool
otherwise = String -> String -> Name -> Name
prefixConName "S" "%" Name
nm
mkTupleDataName :: Int -> Name
mkTupleDataName :: Int -> Name
mkTupleDataName n :: Int
n = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ "STuple" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int -> String
forall a. Show a => a -> String
show Int
n)
prefixConName :: String -> String -> Name -> Name
prefixConName :: String -> String -> Name -> Name
prefixConName pre :: String
pre tyPre :: String
tyPre n :: Name
n = case (Name -> String
nameBase Name
n) of
(':' : rest :: String
rest) -> String -> Name
mkName (':' Char -> String -> String
forall a. a -> [a] -> [a]
: String
tyPre String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
rest)
alpha :: String
alpha -> String -> Name
mkName (String
pre String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
alpha)