module Vectorise.Utils.PADict (
- mkPADictType,
paDictArgType,
paDictOfType,
- paDFunType,
- paMethod
+ paMethod,
+ prDictOfReprType,
+ prDictOfPReprInstTyCon
)
where
import Vectorise.Monad
import Control.Monad
-mkPADictType :: Type -> VM Type
-mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
-
-
+-- | Construct the PA argument type for the tyvar. For the tyvar (v :: *) it's
+-- just PA v. For (v :: (* -> *) -> *) it's
+--
+-- > forall (a :: * -> *). (forall (b :: *). PA b -> PA (a b)) -> PA (v a)
+--
paDictArgType :: TyVar -> VM (Maybe Type)
paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
where
- go ty k | Just k' <- kindView k = go ty k'
go ty (FunTy k1 k2)
= do
tv <- newTyVar (fsLit "a") k1
go ty k
| isLiftedTypeKind k
- = liftM Just (mkPADictType ty)
+ = do
+ pa_cls <- builtin paClass
+ return $ Just $ PredTy $ ClassP pa_cls [ty]
go _ _ = return Nothing
-- | Get the PA dictionary for some type
+--
paDictOfType :: Type -> VM CoreExpr
paDictOfType ty
= paDictOfTyApp ty_fn ty_args
-- for type variables, look up the dfun and apply to the PA dictionaries
-- of the type arguments
paDictOfTyApp (TyVarTy tv) ty_args
- = do dfun <- maybeV (lookupTyVarPA tv)
+ = do dfun <- maybeCantVectoriseM "No PA dictionary for type variable"
+ (ppr tv <+> text "in" <+> ppr ty)
+ $ lookupTyVarPA tv
dicts <- mapM paDictOfType ty_args
return $ dfun `mkTyApps` ty_args `mkApps` dicts
-- for tycons, we also need to apply the dfun to the PR dictionary of
- -- the representation type
+ -- the representation type if the tycon is polymorphic
paDictOfTyApp (TyConApp tc []) ty_args
= do
- dfun <- maybeV $ lookupTyConPA tc
- pr <- prDictOfPRepr tc ty_args
+ dfun <- maybeCantVectoriseM "No PA dictionary for type constructor"
+ (ppr tc <+> text "in" <+> ppr ty)
+ $ lookupTyConPA tc
+ super <- super_dict tc ty_args
dicts <- mapM paDictOfType ty_args
- return $ Var dfun `mkTyApps` ty_args `mkApps` (pr:dicts)
+ return $ Var dfun `mkTyApps` ty_args `mkApps` super `mkApps` dicts
paDictOfTyApp _ _ = failure
- failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
-
-
+ super_dict _ [] = return []
+ super_dict tycon ty_args
+ = do
+ pr <- prDictOfPReprInst (TyConApp tycon ty_args)
+ return [pr]
-paDFunType :: TyCon -> VM Type
-paDFunType tc
- = do
- margs <- mapM paDictArgType tvs
- res <- mkPADictType (mkTyConApp tc arg_tys)
- return . mkForAllTys tvs
- $ mkFunTys [arg | Just arg <- margs] res
- where
- tvs = tyConTyVars tc
- arg_tys = mkTyVarTys tvs
+ failure = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
paMethod :: (Builtins -> Var) -> String -> Type -> VM CoreExpr
paMethod _ name ty
dict <- paDictOfType ty
return $ mkApps (Var fn) [Type ty, dict]
--- | Get the PR (PRepr t) dictionary, where t is the tycon applied to the type
--- arguments
-prDictOfPRepr :: TyCon -> [Type] -> VM CoreExpr
-prDictOfPRepr tycon tys
+-- | Given a type @ty@, return the PR dictionary for @PRepr ty@.
+prDictOfPReprInst :: Type -> VM CoreExpr
+prDictOfPReprInst ty
= do
- (prepr_tc, prepr_args) <- preprSynTyCon (mkTyConApp tycon tys)
- case coreView (mkTyConApp prepr_tc prepr_args) of
- Just rhs -> do
- dict <- prDictOfReprType rhs
- pr_co <- mkBuiltinCo prTyCon
- let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
- let co = mkAppCoercion pr_co
- $ mkSymCoercion
- $ mkTyConApp arg_co prepr_args
- return $ mkCoerce co dict
- Nothing -> cantVectorise "Invalid PRepr type instance"
- $ ppr $ mkTyConApp prepr_tc prepr_args
+ (prepr_tc, prepr_args) <- preprSynTyCon ty
+ prDictOfPReprInstTyCon ty prepr_tc prepr_args
+
+-- | Given a type @ty@, its PRepr synonym tycon and its type arguments,
+-- return the PR @PRepr ty@. Suppose we have:
+--
+-- > type instance PRepr (T a1 ... an) = t
+--
+-- which is internally translated into
+--
+-- > type :R:PRepr a1 ... an = t
+--
+-- and the corresponding coercion. Then,
+--
+-- > prDictOfPReprInstTyCon (T a1 ... an) :R:PRepr u1 ... un = PR (T u1 ... un)
+--
+-- Note that @ty@ is only used for error messages
+--
+prDictOfPReprInstTyCon :: Type -> TyCon -> [Type] -> VM CoreExpr
+prDictOfPReprInstTyCon ty prepr_tc prepr_args
+ | Just rhs <- coreView (mkTyConApp prepr_tc prepr_args)
+ = do
+ dict <- prDictOfReprType' rhs
+ pr_co <- mkBuiltinCo prTyCon
+ let Just arg_co = tyConFamilyCoercion_maybe prepr_tc
+ let co = mkAppCo pr_co
+ $ mkSymCo
+ $ mkAxInstCo arg_co prepr_args
+ return $ mkCoerce co dict
+
+ | otherwise = cantVectorise "Invalid PRepr type instance" (ppr ty)
-- | Get the PR dictionary for a type. The argument must be a representation
-- type.
prDictOfReprType ty
| Just (tycon, tyargs) <- splitTyConApp_maybe ty
= do
- -- a representation tycon must have a PR instance
- dfun <- maybeV $ lookupTyConPR tycon
- prDFunApply dfun tyargs
+ prepr <- builtin preprTyCon
+ if tycon == prepr
+ then do
+ let [ty'] = tyargs
+ pa <- paDictOfType ty'
+ sel <- builtin paPRSel
+ return $ Var sel `App` Type ty' `App` pa
+ else do
+ -- a representation tycon must have a PR instance
+ dfun <- maybeV $ lookupTyConPR tycon
+ prDFunApply dfun tyargs
| otherwise
= do
prsel <- builtin paPRSel
return $ Var prsel `mkApps` [Type ty, pa]
+prDictOfReprType' :: Type -> VM CoreExpr
+prDictOfReprType' ty = prDictOfReprType ty `orElseV`
+ cantVectorise "No PR dictionary for representation type"
+ (ppr ty)
+
-- | Apply a tycon's PR dfun to dictionary arguments (PR or PA) corresponding
-- to the argument types.
prDFunApply :: Var -> [Type] -> VM CoreExpr
projects / ghc-hetmet.git / blobdiff