Add utility function for vectorisation

[ghc-hetmet.git] / compiler / vectorise / Vectorise.hs

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 796a265..6ac3d48 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -130,63 +130,75 @@ vectPolyVar lc v tys
                  return $ mkApps e [arg | (vty, dict) <- zip vtys dicts
                                         , arg <- [Type vty, dict]]
 
-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]
+    
+
+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)
   = 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 <- replicateP 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
@@ -199,6 +211,7 @@ vectExpr lc (_, AnnLet (AnnRec prs) body)
              (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)