X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=20f19b6037022c299e98211b10ccde2900389866;hb=abaea16d603c0cc167df4c873e0e7e57697ba736;hp=e533650d5841ba4943b824fd222f62f4ab338435;hpb=39466c4fe6d9e49f5000b113f7fda4c9afcfb592;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index e533650..20f19b6 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -124,35 +124,42 @@ vectPolyVar lc v tys lexpr <- replicateP vexpr lc return (vexpr, lexpr) where - mk_app e = do - vtys <- mapM vectType tys - dicts <- mapM paDictOfType vtys - return $ mkApps e [arg | (vty, dict) <- zip vtys dicts - , arg <- [Type vty, dict]] + mk_app e = applyToTypes e =<< mapM vectType tys -vectPolyExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr) -vectPolyExpr lc expr +abstractOverTyVars :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a +abstractOverTyVars tvs p = do mdicts <- mapM mk_dict_var tvs - - -- FIXME: shadowing (tvs in lc) - (vmono, lmono) <- localV - $ do - zipWithM_ (\tv -> maybe (deleteTyVarPA tv) (extendTyVarPA tv . Var)) - tvs mdicts - vectExpr lc mono - return (mk_lams tvs mdicts vmono, mk_lams tvs mdicts lmono) + zipWithM_ (\tv -> maybe (deleteTyVarPA tv) (extendTyVarPA tv . Var)) tvs mdicts + p (mk_lams mdicts) where - (tvs, mono) = collectAnnTypeBinders expr - mk_dict_var tv = do r <- paDictArgType tv case r of Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty) Nothing -> return Nothing - mk_lams tvs mdicts = mkLams [arg | (tv, mdict) <- zip tvs mdicts - , arg <- tv : maybeToList mdict] + mk_lams mdicts = mkLams [arg | (tv, mdict) <- zip tvs mdicts + , arg <- tv : maybeToList mdict] + +applyToTypes :: CoreExpr -> [Type] -> VM CoreExpr +applyToTypes expr tys + = do + dicts <- mapM paDictOfType tys + return $ mkApps expr [arg | (ty, dict) <- zip tys dicts + , arg <- [Type ty, dict]] + + +vectPolyExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr) +vectPolyExpr lc expr + = localV + . abstractOverTyVars 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)