Super-monster patch implementing the new typechecker -- at last
[ghc-hetmet.git] / compiler / vectorise / VectMonad.hs
index 2649716..65a3489 100644 (file)
@@ -1,7 +1,7 @@
+{-# LANGUAGE NamedFieldPuns #-}
 
 -- | The Vectorisation monad.
 module VectMonad (
-  Scope(..),
   VM,
 
   noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
@@ -16,11 +16,9 @@ module VectMonad (
   combinePDVar, scalarZip, closureCtrFun,
   builtin, builtins,
 
-  GlobalEnv(..),
   setFamInstEnv,
   readGEnv, setGEnv, updGEnv,
 
-  LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
   getBindName, inBind,
@@ -40,6 +38,7 @@ module VectMonad (
 #include "HsVersions.h"
 
 import VectBuiltIn
+import Vectorise.Env
 
 import HscTypes hiding  ( MonadThings(..) )
 import Module           ( PackageId )
@@ -210,12 +209,15 @@ updVectInfo env tyenv info
     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
     }
   where
+    mk_env :: NamedThing from =>
+              (TypeEnv -> [from])
+           -> (GlobalEnv -> NameEnv to)
+           -> NameEnv (from,to)
     mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
                                    | from <- from_tyenv tyenv
                                    , let name = getName from
                                    , Just to <- [lookupNameEnv (from_env env) name]]
 
-
 -- The Vectorisation Monad ----------------------------------------------------
 
 -- Vectorisation can either succeed with new envionment and a value,
@@ -364,9 +366,11 @@ updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 readLEnv :: (LocalEnv -> a) -> VM a
 readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
 
+-- | Set the local environment.
 setLEnv :: LocalEnv -> VM ()
 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 
+-- | Update the enviroment using a provided function.
 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 
@@ -439,6 +443,8 @@ newTyVar fs k
       u <- liftDs newUnique
       return $ mkTyVar (mkSysTvName u fs) k
 
+
+-- | Add a mapping between a global var and its vectorised version to the state.
 defGlobalVar :: Var -> Var -> VM ()
 defGlobalVar v v' = updGEnv $ \env ->
   env { global_vars = extendVarEnv (global_vars env) v v'
@@ -448,20 +454,36 @@ defGlobalVar v v' = updGEnv $ \env ->
     upd env | isExportedId v = extendVarEnv env v (v, v')
             | otherwise      = env
 
+-- Var ------------------------------------------------------------------------
 -- | Lookup the vectorised and\/or lifted versions of this variable.
 --     If it's in the global environment we get the vectorised version.
 --      If it's in the local environment we get both the vectorised and lifted version.
 --     
 lookupVar :: Var -> VM (Scope Var (Var, Var))
 lookupVar v
-  = do
-      r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
+ = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
       case r of
         Just e  -> return (Local e)
         Nothing -> liftM Global
-                . maybeCantVectoriseM "Variable not vectorised:" (ppr v)
+                . maybeCantVectoriseVarM v
                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
 
+maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
+maybeCantVectoriseVarM v p
+ = do r <- p
+      case r of
+        Just x  -> return x
+        Nothing -> dumpVar v
+
+dumpVar :: Var -> a
+dumpVar var
+       | Just _                <- isClassOpId_maybe var
+       = cantVectorise "ClassOpId not vectorised:" (ppr var)
+
+       | otherwise
+       = cantVectorise "Variable not vectorised:" (ppr var)
+
+-------------------------------------------------------------------------------
 globalScalars :: VM VarSet
 globalScalars = readGEnv global_scalars
 
@@ -572,7 +594,7 @@ lookupInst cls tys
 --
 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 lookupFamInst tycon tys
-  = ASSERT( isOpenTyCon tycon )
+  = ASSERT( isFamilyTyCon tycon )
     do { instEnv <- getFamInstEnv
        ; case lookupFamInstEnv instEnv tycon tys of
           [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
@@ -581,6 +603,8 @@ lookupFamInst tycon tys
                       (ppr $ mkTyConApp tycon tys)
        }
 
+
+-- | Run a vectorisation computation.
 initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 initV pkg hsc_env guts info p
   = do