Clean up handling of PA dictionaries

diff --git a/compiler/package.conf.in b/compiler/package.conf.in
index 8b0beac..24e9d72 100644 (file)
--- a/compiler/package.conf.in
+++ b/compiler/package.conf.in
@@ -260,6 +260,7 @@ exposed-modules:
        VarEnv
        VarSet
         VectMonad
+        VectUtils
         Vectorise
        WorkWrap
        WwLib

diff --git a/compiler/vectorise/VectMonad.hs b/compiler/vectorise/VectMonad.hs
index 289f526..56189f6 100644 (file)
--- a/compiler/vectorise/VectMonad.hs
+++ b/compiler/vectorise/VectMonad.hs
@@ -13,7 +13,7 @@ module VectMonad (
   LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
-  lookupTyCon, extendTyVarPA,
+  lookupTyCon, extendTyVarPA, deleteTyVarPA,
 
   lookupInst, lookupFamInst
 ) where
@@ -239,6 +239,9 @@ lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName
 extendTyVarPA :: Var -> CoreExpr -> VM ()
 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
 
+deleteTyVarPA :: Var -> VM ()
+deleteTyVarPA tv = updLEnv $ \env -> env { local_tyvar_pa = delVarEnv (local_tyvar_pa env) tv }
+
 -- Look up the dfun of a class instance.
 --
 -- The match must be unique - ie, match exactly one instance - but the 

diff --git a/compiler/vectorise/VectUtils.hs b/compiler/vectorise/VectUtils.hs
new file mode 100644 (file)
index 0000000..76d625c
--- /dev/null
+++ b/compiler/vectorise/VectUtils.hs
@@ -0,0 +1,35 @@
+module VectUtils (
+  paDictArgType
+) where
+
+#include "HsVersions.h"
+
+import VectMonad
+
+import Type
+import TypeRep
+import Var
+
+paDictArgType :: TyVar -> VM (Maybe Type)
+paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
+  where
+    go ty k | Just k' <- kindView k = go ty k'
+    go ty (FunTy k1 k2)
+      = do
+          tv   <- newTyVar FSLIT("a") k1
+          mty1 <- go (TyVarTy tv) k1
+          case mty1 of
+            Just ty1 -> do
+                          mty2 <- go (AppTy ty (TyVarTy tv)) k2
+                          return $ fmap (ForAllTy tv . FunTy ty1) mty2
+            Nothing  -> go ty k2
+
+    go ty k
+      | isLiftedTypeKind k
+      = do
+          tc <- builtin paDictTyCon
+          return . Just $ TyConApp tc [ty]
+
+
+    go ty k = return Nothing
+

diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs
index 29774d1..6dde53a 100644 (file)
--- a/compiler/vectorise/Vectorise.hs
+++ b/compiler/vectorise/Vectorise.hs
@@ -4,6 +4,7 @@ where
 #include "HsVersions.h"
 
 import VectMonad
+import VectUtils
 
 import DynFlags
 import HscTypes
@@ -152,49 +153,20 @@ vectExpr lc (_, AnnLet (AnnRec prs) body)
 vectExpr lc (_, AnnLam bndr body)
   | isTyVar bndr
   = do
-      pa_ty          <- paArgType' (TyVarTy bndr) (tyVarKind bndr)
-      pa_var         <- newLocalVar FSLIT("dPA") pa_ty
-      (vbody, lbody) <- localV
-                      $ do
-                          extendTyVarPA bndr (Var pa_var)
-                          -- FIXME: what about shadowing here (bndr in lc)?
-                          vectExpr lc body
-      return (mkLams [bndr, pa_var] vbody,
-              mkLams [bndr, pa_var] lbody)
+      r <- paDictArgType bndr
+      (upd_env, add_lam) <- get_upd r
+      (vbody, lbody) <- localV (upd_env >> vectExpr lc body)
+      return (Lam bndr (add_lam vbody), Lam bndr (add_lam lbody))
+  where
+    get_upd Nothing = return (deleteTyVarPA bndr, id)
+    get_upd (Just pa_ty) = do
+                             pa_var <- newLocalVar FSLIT("dPA") pa_ty
+                             return (extendTyVarPA bndr (Var pa_var),
+                                     Lam pa_var)
 
 -- ----------------------------------------------------------------------------
 -- PA dictionaries
 
-paArgType :: Type -> Kind -> VM (Maybe Type)
-paArgType ty k
-  | Just k' <- kindView k = paArgType ty k'
-
--- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
--- be made up of * and (->), i.e., they can't be coercion kinds or #.
-paArgType ty (FunTy k1 k2)
-  = do
-      tv  <- newTyVar FSLIT("a") k1
-      ty1 <- paArgType' (TyVarTy tv) k1
-      ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
-      return . Just $ ForAllTy tv (FunTy ty1 ty2)
-
-paArgType ty k
-  | isLiftedTypeKind k
-  = do
-      tc <- builtin paDictTyCon
-      return . Just $ TyConApp tc [ty]
-
-  | otherwise
-  = return Nothing 
-
-paArgType' :: Type -> Kind -> VM Type
-paArgType' ty k
-  = do
-      r <- paArgType ty k
-      case r of
-        Just ty' -> return ty'
-        Nothing  -> pprPanic "paArgType'" (ppr ty)
-
 paOfTyCon :: TyCon -> VM CoreExpr
 -- FIXME: just for now
 paOfTyCon tc = maybeV (readGEnv $ \env -> lookupNameEnv (global_tycon_pa env) (tyConName tc))
@@ -244,9 +216,12 @@ vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
                                              (mapM vectType [ty1,ty2])
 vectType (ForAllTy tv ty)
   = do
-      r   <- paArgType (TyVarTy tv) (tyVarKind tv)
+      r   <- paDictArgType tv
       ty' <- vectType ty
-      return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
+      return $ ForAllTy tv (wrap r ty')
+  where
+    wrap Nothing      = id
+    wrap (Just pa_ty) = FunTy pa_ty
 
 vectType ty = pprPanic "vectType:" (ppr ty)