Nicer names for hoisted functions

compiler/vectorise/VectMonad.hs

@@ -15,6 +15,8 @@ module VectMonad (
   LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
+  getBindName, inBind,
+
   lookupVar, defGlobalVar,
   lookupTyCon, defTyCon,
   lookupDataCon, defDataCon,
@@ -148,6 +150,9 @@ data LocalEnv = LocalEnv {
 
                  -- Mapping from tyvars to their PA dictionaries
                , local_tyvar_pa :: VarEnv CoreExpr
+
+                 -- Local binding name
+               , local_bind_name :: FastString
                }
               
 
@@ -176,6 +181,7 @@ emptyLocalEnv = LocalEnv {
                    local_vars     = emptyVarEnv
                  , local_tyvars   = []
                  , local_tyvar_pa = emptyVarEnv
+                 , local_bind_name  = FSLIT("fn")
                  }
 
 -- FIXME
@@ -236,7 +242,7 @@ localV p = do
 closedV :: VM a -> VM a
 closedV p = do
               env <- readLEnv id
-              setLEnv emptyLocalEnv
+              setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
               x <- p
               setLEnv env
               return x
@@ -271,6 +277,14 @@ getInstEnv = readGEnv global_inst_env
 getFamInstEnv :: VM FamInstEnvs
 getFamInstEnv = readGEnv global_fam_inst_env
 
+getBindName :: VM FastString
+getBindName = readLEnv local_bind_name
+
+inBind :: Id -> VM a -> VM a
+inBind id p
+  = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
+       p
+
 cloneName :: (OccName -> OccName) -> Name -> VM Name
 cloneName mk_occ name = liftM make (liftDs newUnique)
   where

compiler/vectorise/VectUtils.hs

@@ -216,19 +216,20 @@ hoistExpr fs expr
         env { global_bindings = (var, expr) : global_bindings env }
       return var
 
-hoistVExpr :: FastString -> VExpr -> VM VVar
-hoistVExpr fs (ve, le)
+hoistVExpr :: VExpr -> VM VVar
+hoistVExpr (ve, le)
   = do
+      fs <- getBindName
       vv <- hoistExpr ('v' `consFS` fs) ve
       lv <- hoistExpr ('l' `consFS` fs) le
       return (vv, lv)
 
-hoistPolyVExpr :: FastString -> [TyVar] -> VM VExpr -> VM VExpr
-hoistPolyVExpr fs tvs p
+hoistPolyVExpr :: [TyVar] -> VM VExpr -> VM VExpr
+hoistPolyVExpr tvs p
   = do
       expr <- closedV . polyAbstract tvs $ \abstract ->
               liftM (mapVect abstract) p
-      fn   <- hoistVExpr fs expr
+      fn   <- hoistVExpr expr
       polyVApply (vVar fn) (mkTyVarTys tvs)
 
 takeHoisted :: VM [(Var, CoreExpr)]
@@ -256,7 +257,7 @@ buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body
       res_ty' <- mkClosureTypes arg_tys res_ty
       arg <- newLocalVVar FSLIT("x") arg_ty
       buildClosure tvs lc vars arg_ty res_ty'
-        . hoistPolyVExpr FSLIT("fn") tvs
+        . hoistPolyVExpr tvs
         $ do
             clo <- buildClosures tvs lc (vars ++ [arg]) arg_tys res_ty mk_body
             return $ vLams lc (vars ++ [arg]) clo
@@ -273,7 +274,7 @@ buildClosure tvs lv vars arg_ty res_ty mk_body
       env_bndr <- newLocalVVar FSLIT("env") env_ty
       arg_bndr <- newLocalVVar FSLIT("arg") arg_ty
 
-      fn <- hoistPolyVExpr FSLIT("fn") tvs
+      fn <- hoistPolyVExpr tvs
           $ do
               body  <- mk_body
               body' <- bind (vVar env_bndr)

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)
 
 -- ----------------------------------------------------------------------------
@@ -228,7 +229,7 @@ vectExpr lc (_, AnnCase expr bndr ty alts)
 
 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,12 +237,16 @@ 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
 
+    vect_rhs bndr rhs = localV
+                      . inBind bndr
+                      $ vectExpr lc rhs
+
 vectExpr lc e@(fvs, AnnLam bndr _)
   | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e)
   | otherwise = vectLam lc fvs bs body
@@ -260,7 +265,7 @@ vectLam lc fvs bs body
       res_ty  <- vectType (exprType $ deAnnotate body)
 
       buildClosures tyvars lc vvs arg_tys res_ty
-        . hoistPolyVExpr FSLIT("fn") tyvars
+        . hoistPolyVExpr tyvars
         $ do
             new_lc <- newLocalVar FSLIT("lc") intPrimTy
             (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)