New closure generation code

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 7b0e4af..545460d 100644 (file)
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -13,13 +13,19 @@ module VectUtils (
 
 import VectMonad
 
+import DsUtils
 import CoreSyn
 import CoreUtils
 import Type
 import TypeRep
 import TyCon
+import DataCon            ( dataConWrapId )
 import Var
+import Id                 ( mkWildId )
+import MkId               ( unwrapFamInstScrut )
 import PrelNames
+import TysWiredIn
+import BasicTypes         ( Boxity(..) )
 
 import Outputable
 import FastString
@@ -173,6 +179,13 @@ hoistExpr fs expr
         env { global_bindings = (var, expr) : global_bindings env }
       return var
 
+hoistPolyExpr :: FastString -> [TyVar] -> CoreExpr -> VM CoreExpr
+hoistPolyExpr fs tvs expr
+  = do
+      poly_expr <- closedV . polyAbstract tvs $ \abstract -> return (abstract expr)
+      fn        <- hoistExpr fs poly_expr
+      polyApply (Var fn) (mkTyVarTys tvs)
+
 takeHoisted :: VM [(Var, CoreExpr)]
 takeHoisted
   = do
@@ -180,3 +193,91 @@ takeHoisted
       setGEnv $ env { global_bindings = [] }
       return $ global_bindings env
 
+
+mkClosure :: Type -> Type -> Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
+mkClosure arg_ty res_ty env_ty pa_dict vfn lfn env
+  = do
+      mk <- builtin mkClosureVar
+      return $ Var mk `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [pa_dict, vfn, lfn, env]
+
+mkClosureP :: Type -> Type -> Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
+mkClosureP arg_ty res_ty env_ty pa_dict vfn lfn env
+  = do
+      mk <- builtin mkClosurePVar
+      return $ Var mk `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [pa_dict, vfn, lfn, env]
+
+-- (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,Var)] -> (Var,Var) -> Var -> CoreExpr -> CoreExpr -> VM (CoreExpr, CoreExpr)
+buildClosure tvs env (varg, larg) lv vbody lbody
+  = do
+      let (venv_ty, venv, bind_venv) = mkVectEnv tys vs
+      (lenv, bind_lenv) <- mkLiftEnv (Var lv) tys ls
+      lenv_ty <- mkPArrayType venv_ty
+
+      venv_bndr <- newLocalVar FSLIT("env") venv_ty
+      lenv_bndr <- newLocalVar FSLIT("env") lenv_ty
+
+      let mono_vfn = mkLams [venv_bndr, varg]
+                   . bind_venv (Var venv_bndr)
+                   $ vbody `mkVarApps` vs `mkVarApps` [varg]
+          mono_lfn = mkLams [lenv_bndr, larg]
+                   . bind_lenv (Var lenv_bndr) lv
+                   $ lbody `mkVarApps` (lv:ls) `mkVarApps` [larg]
+
+      vfn <- hoistPolyExpr FSLIT("vfn") tvs mono_vfn
+      lfn <- hoistPolyExpr FSLIT("lfn") tvs mono_lfn
+
+      pa_dict <- paDictOfType venv_ty
+
+      vclo <- mkClosure  arg_ty res_ty venv_ty pa_dict vfn lfn venv
+      lclo <- mkClosureP arg_ty res_ty venv_ty pa_dict vfn lfn lenv
+
+      return (vclo, lclo)
+
+  where
+    vs  = map fst env
+    ls  = map snd env
+    tys = map idType vs
+
+    arg_ty = idType varg
+    res_ty = exprType vbody
+    
+mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
+mkVectEnv []   []  = (unitTy, Var unitDataConId, \env body -> body)
+mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
+mkVectEnv tys  vs  = (ty, mkCoreTup (map Var vs),
+                        \env body -> Case env (mkWildId ty) (exprType body)
+                                       [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
+  where
+    ty = mkCoreTupTy tys
+
+mkLiftEnv :: CoreExpr -> [Type] -> [Var]
+          -> VM (CoreExpr, CoreExpr -> Var -> CoreExpr -> CoreExpr)
+mkLiftEnv lc [ty] [v]
+  = do
+      len <- lengthPA (Var v)
+      return (Var v, \env lv body -> Let (NonRec v env)
+                                   $ Case len lv (exprType body) [(DEFAULT, [], body)])
+
+-- NOTE: this transparently deals with empty environments
+mkLiftEnv lc tys vs
+  = do
+      (env_tc, env_tyargs) <- lookupPArrayFamInst vty
+      let [env_con] = tyConDataCons env_tc
+          
+          env = Var (dataConWrapId env_con)
+                `mkTyApps` env_tyargs
+                `mkApps`   (lc : map Var vs)
+
+          bind env lv body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
+                             in
+                             Case scrut (mkWildId (exprType scrut)) (exprType body)
+                               [(DataAlt env_con, lv : vs, body)]
+      return (env, bind)
+  where
+    vty = mkCoreTupTy tys
+