Refactor slightly

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

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 796a265..20f19b6 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -124,69 +124,84 @@ 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)
   = 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 +214,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)