Improve closure generation for functions with multiple parameters

2 files changed, 49 insertions, 10 deletions

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 3abbe44b49..57571ab630 100644
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -1,5 +1,6 @@
 module VectUtils (
   collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
+  collectAnnValBinders,
   splitClosureTy,
   mkPADictType, mkPArrayType,
   paDictArgType, paDictOfType,
@@ -7,7 +8,7 @@ module VectUtils (
   polyAbstract, polyApply, polyVApply,
   lookupPArrayFamInst,
   hoistExpr, hoistPolyVExpr, takeHoisted,
-  buildClosure
+  buildClosure, buildClosures
 ) where
 
 #include "HsVersions.h"
@@ -46,6 +47,12 @@ collectAnnTypeBinders expr = go [] expr
     go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
     go bs e                           = (reverse bs, e)
 
+collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
+collectAnnValBinders expr = go [] expr
+  where
+    go bs (_, AnnLam b e) | isId b = go (b:bs) e
+    go bs e                        = (reverse bs, e)
+
 isAnnTypeArg :: AnnExpr b ann -> Bool
 isAnnTypeArg (_, AnnType t) = True
 isAnnTypeArg _              = False
@@ -72,6 +79,20 @@ splitPArrayTy ty
 
   | otherwise = pprPanic "splitPArrayTy" (ppr ty)
 
+mkClosureType :: Type -> Type -> VM Type
+mkClosureType arg_ty res_ty
+  = do
+      tc <- builtin closureTyCon
+      return $ mkTyConApp tc [arg_ty, res_ty]
+
+mkClosureTypes :: [Type] -> Type -> VM Type
+mkClosureTypes arg_tys res_ty
+  = do
+      tc <- builtin closureTyCon
+      return $ foldr (mk tc) res_ty arg_tys
+  where
+    mk tc arg_ty res_ty = mkTyConApp tc [arg_ty, res_ty]
+
 mkPADictType :: Type -> VM Type
 mkPADictType ty
   = do
@@ -227,11 +248,24 @@ mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
       return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
               Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
 
+buildClosures :: [TyVar] -> Var -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
+buildClosures tvs lc vars [arg_ty] res_ty mk_body
+  = buildClosure tvs lc vars arg_ty res_ty mk_body
+buildClosures tvs lc vars (arg_ty : arg_tys) res_ty mk_body
+  = do
+      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
+        $ do
+            clo <- buildClosures tvs lc (vars ++ [arg]) arg_tys res_ty mk_body
+            return $ vLams lc (vars ++ [arg]) clo
+
 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
 --   where
 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
-
+--
 buildClosure :: [TyVar] -> Var -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
 buildClosure tvs lv vars arg_ty res_ty mk_body
   = do
diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 055137abe4..89ee166dd4 100644
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -242,26 +242,31 @@ vectExpr lc (_, AnnLet (AnnRec bs) body)
   where
     (bndrs, rhss) = unzip bs
 
-vectExpr lc e@(_, AnnLam bndr body)
-  | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)
+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
 
-vectExpr lc (fvs, AnnLam bndr body)
+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
+      arg_tys <- mapM (vectType . idType) bs
+      res_ty  <- vectType (exprType $ deAnnotate body)
+
+      buildClosures tyvars lc vvs arg_tys res_ty
         . hoistPolyVExpr FSLIT("fn") tyvars
         $ do
             new_lc <- newLocalVar FSLIT("lc") intPrimTy
-            (vbndrs, vbody) <- vectBndrsIn (vs ++ [bndr])
+            (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
                                            (vectExpr new_lc body)
             return $ vLams new_lc vbndrs vbody
-
+  
 vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr)
 vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
 vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)