X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=5734ae1abef0126c3c2064c1e22ab891d970da2a;hb=e1364f66b4e743237e942e0826ed096f5e06de76;hp=055137abe41c5abb102e22d1aababaa1760b72a8;hpb=0a21de62e274acc8e8e260298da4f6c1ee18ecc2;p=ghc-hetmet.git

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 055137a..5734ae1 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -42,7 +42,7 @@ import BasicTypes           ( Boxity(..) )
 
 import Outputable
 import FastString
-import Control.Monad        ( liftM, liftM2, mapAndUnzipM )
+import Control.Monad        ( liftM, liftM2, zipWithM, mapAndUnzipM )
 
 vectorise :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
           -> IO (SimplCount, ModGuts)
@@ -81,7 +81,7 @@ vectTopBind :: CoreBind -> VM CoreBind
 vectTopBind b@(NonRec var expr)
   = do
       var'  <- vectTopBinder var
-      expr' <- vectTopRhs expr
+      expr' <- vectTopRhs var expr
       hs    <- takeHoisted
       return . Rec $ (var, expr) : (var', expr') : hs
   `orElseV`
@@ -90,7 +90,7 @@ vectTopBind b@(NonRec var expr)
 vectTopBind b@(Rec bs)
   = do
       vars'  <- mapM vectTopBinder vars
-      exprs' <- mapM vectTopRhs exprs
+      exprs' <- zipWithM vectTopRhs vars exprs
       hs     <- takeHoisted
       return . Rec $ bs ++ zip vars' exprs' ++ hs
   `orElseV`
@@ -108,11 +108,12 @@ vectTopBinder var
       defGlobalVar var var'
       return var'
     
-vectTopRhs :: CoreExpr -> VM CoreExpr
-vectTopRhs expr
+vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
+vectTopRhs var expr
   = do
       lc <- newLocalVar FSLIT("lc") intPrimTy
       closedV . liftM vectorised
+              . inBind var
               $ vectPolyExpr lc (freeVars expr)
 
 -- ----------------------------------------------------------------------------
@@ -149,17 +150,6 @@ vectBndrsIn vs p
 -- ----------------------------------------------------------------------------
 -- Expressions
 
-capply :: VExpr -> VExpr -> VM VExpr
-capply (vfn, lfn) (varg, larg)
-  = do
-      apply  <- builtin applyClosureVar
-      applyP <- builtin applyClosurePVar
-      return (mkApps (Var apply)  [Type arg_ty, Type res_ty, vfn, varg],
-              mkApps (Var applyP) [Type arg_ty, Type res_ty, lfn, larg])
-  where
-    fn_ty            = exprType vfn
-    (arg_ty, res_ty) = splitClosureTy fn_ty
-
 vectVar :: Var -> Var -> VM VExpr
 vectVar lc v
   = do
@@ -221,14 +211,14 @@ vectExpr lc (_, AnnApp fn arg)
   = do
       fn'  <- vectExpr lc fn
       arg' <- vectExpr lc arg
-      capply fn' arg'
+      mkClosureApp fn' arg'
 
 vectExpr lc (_, AnnCase expr bndr ty alts)
   = panic "vectExpr: case"
 
 vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
   = do
-      vrhs <- vectPolyExpr lc rhs
+      vrhs <- localV . inBind bndr $ vectPolyExpr lc rhs
       (vbndr, vbody) <- vectBndrIn bndr (vectExpr lc body)
       return $ vLet (vNonRec vbndr vrhs) vbody
 
@@ -236,33 +226,41 @@ vectExpr lc (_, AnnLet (AnnRec bs) body)
   = do
       (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
                                 $ liftM2 (,)
-                                  (mapM (vectExpr lc) rhss)
+                                  (zipWithM vect_rhs bndrs rhss)
                                   (vectPolyExpr lc body)
       return $ vLet (vRec vbndrs vrhss) vbody
   where
     (bndrs, rhss) = unzip bs
 
-vectExpr lc e@(_, AnnLam bndr body)
-  | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)
+    vect_rhs bndr rhs = localV
+                      . inBind bndr
+                      $ vectExpr lc rhs
 
-vectExpr lc (fvs, AnnLam bndr body)
+vectExpr lc e@(fvs, AnnLam bndr _)
+  | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
+  | otherwise = vectLam lc fvs bs body
+  where
+    (bs,body) = collectAnnValBinders e
+
+vectLam :: Var -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
+vectLam lc fvs bs body
   = do
       tyvars <- localTyVars
       (vs, vvs) <- readLEnv $ \env ->
                    unzip [(var, vv) | var <- varSetElems fvs
                                     , Just vv <- [lookupVarEnv (local_vars env) var]]
 
-      arg_ty <- vectType (idType bndr)
-      res_ty <- vectType (exprType $ deAnnotate body)
-      buildClosure tyvars lc vvs arg_ty res_ty
-        . hoistPolyVExpr FSLIT("fn") tyvars
-        $ do
-            new_lc <- newLocalVar FSLIT("lc") intPrimTy
-            (vbndrs, vbody) <- vectBndrsIn (vs ++ [bndr])
-                                           (vectExpr new_lc body)
-            return $ vLams new_lc vbndrs vbody
+      arg_tys <- mapM (vectType . idType) bs
+      res_ty  <- vectType (exprType $ deAnnotate body)
 
-vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr)
+      buildClosures tyvars lc vvs arg_tys res_ty
+        . hoistPolyVExpr tyvars
+        $ do
+            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
+                                           (vectExpr lc body)
+            return $ vLams lc vbndrs vbody
+  
+vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM VExpr
 vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
 vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)