Add generation of PR dictionaries

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

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
index 0789688..b00206d 100644 (file)
--- a/compiler/vectorise/VectUtils.hs
+++ b/compiler/vectorise/VectUtils.hs
@@ -6,6 +6,7 @@ module VectUtils (
   mkPRepr, mkToPRepr, mkFromPRepr,
   mkPADictType, mkPArrayType, mkPReprType,
   parrayReprTyCon, parrayReprDataCon, mkVScrut,
+  prDictOfType,
   paDictArgType, paDictOfType, paDFunType,
   paMethod, lengthPA, replicatePA, emptyPA, liftPA,
   polyAbstract, polyApply, polyVApply,
@@ -153,9 +154,10 @@ mkToPRepr ess
       left_dc  <- builtin leftDataCon
       right_dc <- builtin rightDataCon
 
-      let mk_embed (expr, ty, pa)
-            = (mkConApp   embed_dc [Type ty, pa, expr],
+      let mk_embed expr
+            = (mkConApp   embed_dc [Type ty, expr],
                mkTyConApp embed_tc [ty])
+            where ty = exprType expr
 
           mk_cross (expr1, ty1) (expr2, ty2)
             = (mkConApp   cross_dc [Type ty1, Type ty2, expr1, expr2],
@@ -172,14 +174,8 @@ mkToPRepr ess
               (mkConApp left_dc [Type lty, Type rty, expr]
                  : [mkConApp right_dc [Type lty, Type rty, alt] | alt <- alts],
                mkTyConApp plus_tc [lty, rty])
-      
-      liftM (mk_sum . map (mk_tup . map mk_embed))
-            (mapM (mapM init) ess)
-  where
-    init expr = let ty = exprType expr
-                in do
-                     pa <- paDictOfType ty
-                     return (expr, ty, pa)
+
+      return . mk_sum $ map (mk_tup . map mk_embed) ess
 
 mkFromPRepr :: CoreExpr -> Type -> [([Var], CoreExpr)] -> VM CoreExpr
 mkFromPRepr scrut res_ty alts
@@ -191,24 +187,22 @@ mkFromPRepr scrut res_ty alts
       pa_tc    <- builtin paTyCon
 
       let un_embed expr ty var res
-            = do
-                pa <- newLocalVar FSLIT("pa") (mkTyConApp pa_tc [idType var])
-                return $ Case expr (mkWildId ty) res_ty
-                         [(DataAlt embed_dc, [pa, var], res)]
+            = Case expr (mkWildId ty) res_ty
+                   [(DataAlt embed_dc, [var], res)]
 
           un_cross expr ty var1 var2 res
             = Case expr (mkWildId ty) res_ty
                 [(DataAlt cross_dc, [var1, var2], res)]
 
           un_tup expr ty []    res = return res
-          un_tup expr ty [var] res = un_embed expr ty var res
+          un_tup expr ty [var] res = return $ un_embed expr ty var res
           un_tup expr ty (var : vars) res
             = do
                 lv <- newLocalVar FSLIT("x") lty
                 rv <- newLocalVar FSLIT("y") rty
-                liftM (un_cross expr ty lv rv)
-                        (un_embed (Var lv) lty var
-                         =<< un_tup (Var rv) rty vars res)
+                liftM (un_cross expr ty lv rv
+                      . un_embed (Var lv) lty var)
+                      (un_tup (Var rv) rty vars res)
             where
               (lty, rty) = splitCrossTy ty
 
@@ -261,6 +255,37 @@ mkVScrut (ve, le)
       (tc, arg_tys) <- parrayReprTyCon (exprType ve)
       return ((ve, unwrapFamInstScrut tc arg_tys le), tc, arg_tys)
 
+
+prDictOfType :: Type -> VM CoreExpr
+prDictOfType orig_ty
+  | Just (tycon, ty_args) <- splitTyConApp_maybe orig_ty
+  = do
+      dfun <- traceMaybeV "prDictOfType" (ppr tycon) (lookupTyConPR tycon)
+      prDFunApply (Var dfun) ty_args
+
+prDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
+prDFunApply dfun tys
+  = do
+      args <- mapM mkDFunArg arg_tys
+      return $ mkApps mono_dfun args
+  where
+    mono_dfun    = mkTyApps dfun tys
+    (arg_tys, _) = splitFunTys (exprType mono_dfun)
+
+mkDFunArg :: Type -> VM CoreExpr
+mkDFunArg ty
+  | Just (tycon, [arg]) <- splitTyConApp_maybe ty
+
+  = let name = tyConName tycon
+
+        get_dict | name == paTyConName = paDictOfType
+                 | name == prTyConName = prDictOfType
+                 | otherwise           = pprPanic "mkDFunArg" (ppr ty)
+
+    in get_dict arg
+
+mkDFunArg ty = pprPanic "mkDFunArg" (ppr ty)
+
 paDictArgType :: TyVar -> VM (Maybe Type)
 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
   where