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