+vectPolyExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
+vectPolyExpr lc expr
+ = polyAbstract tvs $ \mk_lams ->
+ -- FIXME: shadowing (tvs in lc)
+ do
+ (vmono, lmono) <- vectExpr lc mono
+ return $ (mk_lams vmono, mk_lams lmono)
+ where
+ (tvs, mono) = collectAnnTypeBinders expr
+
+vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
+vectExpr lc (_, AnnType ty)
+ = do
+ vty <- vectType ty
+ return (Type vty, Type vty)
+
+vectExpr lc (_, AnnVar v) = vectVar lc v
+
+vectExpr lc (_, AnnLit lit)
+ = do
+ let vexpr = Lit lit
+ lexpr <- replicatePA vexpr lc
+ return (vexpr, lexpr)
+
+vectExpr lc (_, AnnNote note expr)
+ = do
+ (vexpr, lexpr) <- vectExpr lc expr
+ return (Note note vexpr, Note note lexpr)
+
+vectExpr lc e@(_, AnnApp _ arg)
+ | isAnnTypeArg arg
+ = vectTyAppExpr lc fn tys
+ where
+ (fn, tys) = collectAnnTypeArgs e
+
+vectExpr lc (_, AnnApp fn arg)
+ = do
+ fn' <- vectExpr lc fn
+ arg' <- vectExpr lc arg
+ capply fn' arg'
+
+vectExpr lc (_, AnnCase expr bndr ty alts)
+ = panic "vectExpr: case"
+
+vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
+ = do
+ (vrhs, lrhs) <- vectPolyExpr lc rhs
+ (vbndr, lbndr, (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body)
+ return (Let (NonRec vbndr vrhs) vbody,
+ Let (NonRec lbndr lrhs) lbody)
+
+vectExpr lc (_, AnnLet (AnnRec prs) body)
+ = do
+ (vbndrs, lbndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect
+ return (Let (Rec (zip vbndrs vrhss)) vbody,
+ Let (Rec (zip lbndrs lrhss)) lbody)
+ where
+ (bndrs, rhss) = unzip prs
+
+ vect = do
+ (vrhss, lrhss) <- mapAndUnzipM (vectExpr lc) rhss
+ (vbody, lbody) <- vectPolyExpr lc body
+ return (vrhss, vbody, lrhss, lbody)
+
+vectExpr lc e@(_, AnnLam bndr body)
+ | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)
+
+vectExpr lc (fvs, AnnLam bndr body)
+ = do
+ tyvars <- localTyVars
+ info <- mkCEnvInfo fvs bndr body
+ (poly_vfn, poly_lfn) <- mkClosureFns info tyvars bndr body
+
+ vfn_var <- hoistExpr FSLIT("vfn") poly_vfn
+ lfn_var <- hoistExpr FSLIT("lfn") poly_lfn
+
+ let (venv, lenv) = mkClosureEnvs info lc
+
+ let env_ty = cenv_vty info
+
+ pa_dict <- paDictOfType env_ty
+
+ arg_ty <- vectType (varType bndr)
+ res_ty <- vectType (exprType $ deAnnotate body)
+
+ -- FIXME: move the functions to the top level
+ mono_vfn <- polyApply (Var vfn_var) (mkTyVarTys tyvars)
+ mono_lfn <- polyApply (Var lfn_var) (mkTyVarTys tyvars)
+
+ mk_clo <- builtin mkClosureVar
+ mk_cloP <- builtin mkClosurePVar
+
+ let vclo = Var mk_clo `mkTyApps` [arg_ty, res_ty, env_ty]
+ `mkApps` [pa_dict, mono_vfn, mono_lfn, venv]
+
+ lclo = Var mk_cloP `mkTyApps` [arg_ty, res_ty, env_ty]
+ `mkApps` [pa_dict, mono_vfn, mono_lfn, lenv]
+
+ return (vclo, lclo)
+
+data CEnvInfo = CEnvInfo {
+ cenv_vars :: [Var]
+ , cenv_values :: [(CoreExpr, CoreExpr)]
+ , cenv_vty :: Type
+ , cenv_lty :: Type
+ , cenv_repr_tycon :: TyCon
+ , cenv_repr_tyargs :: [Type]
+ , cenv_repr_datacon :: DataCon