module Data.TrieMap.Representation.TH (genRepr, genOptRepr, genOrdRepr) where
import Language.Haskell.TH.Syntax
import Language.Haskell.TH.ExpandSyns
import qualified Data.Vector as V
import Data.TrieMap.Representation.Class
import Data.TrieMap.Representation.TH.Utils
import Data.TrieMap.Representation.TH.Representation
import Data.TrieMap.Representation.TH.Factorized
import Data.TrieMap.Representation.TH.ReprMonad
genOrdRepr :: Name -> Q [Dec]
genOrdRepr tycon = execReprMonad $ do
(cxt, ty, _) <- getDataForName tycon
outputRepr cxt ty =<< ordRepr ty
getDataForName :: Quasi m => Name -> m (Cxt, Type, [AlgCon])
getDataForName tycon = do
TyConI dec <- qReify tycon
let theTyp = compose tycon . map (mkName . nameBase . tyVarBndrVar)
case dec of
DataD cxt _ tyvars cons _ ->
return (cxt, theTyp tyvars, map algCon cons)
NewtypeD cxt _ tyvars con _ ->
return (cxt, theTyp tyvars, [algCon con])
_ -> error "Error: could not get kind of type constructor"
getDataForType :: Quasi m => Type -> m (Cxt, [AlgCon])
getDataForType ty
| (ConT tyCon, args) <- decompose ty
= do TyConI dec <- qReify tyCon
let subAll tyvars cxt cons = let subs = zip (map tyVarBndrVar tyvars) args in
([foldr substInPred p subs | p <- cxt], [foldr substInAlgCon (algCon con) subs | con <- cons])
case dec of
DataD cxt _ tyvars cons _ ->
return (subAll tyvars cxt cons)
NewtypeD cxt _ tyvars con _ ->
return (subAll tyvars cxt [con])
_ -> failure
| otherwise = failure
where failure = fail "Error: could not reify type constructor"
genRepr :: Name -> Q [Dec]
genRepr tyCon = execReprMonad $ do
(_, ty, _) <- getDataForName tyCon
let ?combine = mergeWith sumRepr
genReprMain ty
genOptRepr :: Name -> Q [Dec]
genOptRepr tyCon = execReprMonad $ do
(_, ty, _) <- getDataForName tyCon
let ?combine = unify
genReprMain ty
mustBreakTy :: Type -> ReprMonad Bool
mustBreakTy ty = case decompose ty of
(ConT tyCon, _) -> mustBreak tyCon
_ -> return False
recurseTy :: Type -> ReprMonad a -> ReprMonad a
recurseTy ty m = case decompose ty of
(ConT tyCon, _) -> recurse tyCon m
_ -> m
genReprMain :: (?combine :: [Representation] -> Representation) => Type -> ReprMonad Type
genReprMain ty = do
breakTy <- mustBreakTy ty
if breakTy then fail "Cannot recurse here"
else do
knownInst <- getInstance ty
case knownInst of
Just known -> return known
Nothing -> do
(cxt, cons) <- getDataForType ty
conReprs <- mapM (recurseTy ty . conRepr) cons
outputRepr cxt ty (checkEnumRepr $ ?combine conReprs)
conRepr :: (?combine :: [Representation] -> Representation) => AlgCon -> ReprMonad Representation
conRepr (con, []) = return $ conify con unitRepr
conRepr (con, args) = do
argReprs <- mapM typeRepr args
return (conify con (foldr1 prodRepr argReprs))
typeRepr :: (?combine :: [Representation] -> Representation) => Type -> ReprMonad Representation
typeRepr t00 = liftQuasi (expandSyns t00) >>= \ t0 -> case decompose t0 of
(ListT, [t]) -> do
tRepr <- typeRepr t
vectorizeRepr (VarE 'V.fromList) tRepr
(TupleT 0, _) -> return unitRepr
(TupleT _, ts) -> do
reps <- mapM typeRepr ts
return $ mapReprInput TupP $ mergeWith prodRepr reps
(ConT con, ts)
| con == ''() -> return unitRepr
| con == ''Either, [tL, tR] <- ts
-> do reprL <- typeRepr tL
reprR <- typeRepr tR
return (mapReprInput (ConP leftN) reprL `sumRepr` mapReprInput (ConP rightN) reprR)
| con == ''Maybe, [t] <- ts
-> do tRepr <- typeRepr t
return (conify 'Nothing unitRepr `sumRepr` conify 'Just tRepr)
_ -> bootstrapRepr t0
bootstrapRepr :: (?combine :: [Representation] -> Representation) => Type -> ReprMonad Representation
bootstrapRepr t0 = qRecover fallback
(do _tRep <- genReprMain t0
recursiveRepr (ConT ''Rep `AppT` t0) (VarE 'toRep))
where fallback = keyRepr t0