8fdfcdb99992e199e62d499fd99bd0bb751c0618
[ghc-hetmet.git] / compiler / vectorise / VectMonad.hs
1 module VectMonad (
2   Scope(..),
3   VM,
4
5   noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
6   onlyIfV, fixV, localV, closedV,
7   initV, cantVectorise, maybeCantVectorise, maybeCantVectoriseM,
8   liftDs,
9   cloneName, cloneId, cloneVar,
10   newExportedVar, newLocalVar, newDummyVar, newTyVar,
11   
12   Builtins(..), sumTyCon, prodTyCon,
13   combinePAVar, scalarZip, closureCtrFun,
14   builtin, builtins,
15
16   GlobalEnv(..),
17   setFamInstEnv,
18   readGEnv, setGEnv, updGEnv,
19
20   LocalEnv(..),
21   readLEnv, setLEnv, updLEnv,
22
23   getBindName, inBind,
24
25   lookupVar, defGlobalVar, globalScalars,
26   lookupTyCon, defTyCon,
27   lookupDataCon, defDataCon,
28   lookupTyConPA, defTyConPA, defTyConPAs,
29   lookupTyConPR,
30   lookupBoxedTyCon,
31   lookupPrimMethod, lookupPrimPArray,
32   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
33
34   lookupInst, lookupFamInst
35 ) where
36
37 #include "HsVersions.h"
38
39 import VectBuiltIn
40
41 import HscTypes hiding  ( MonadThings(..) )
42 import Module           ( PackageId )
43 import CoreSyn
44 import Class
45 import TyCon
46 import DataCon
47 import Type
48 import Var
49 import VarSet
50 import VarEnv
51 import Id
52 import Name
53 import NameEnv
54 import IOEnv         ( liftIO )
55
56 import DsMonad
57
58 import InstEnv
59 import FamInstEnv
60
61 import Outputable
62 import FastString
63 import SrcLoc        ( noSrcSpan )
64
65 import Control.Monad
66
67 data Scope a b = Global a | Local b
68
69 -- ----------------------------------------------------------------------------
70 -- Vectorisation monad
71
72 data GlobalEnv = GlobalEnv {
73                   -- Mapping from global variables to their vectorised versions.
74                   -- 
75                   global_vars :: VarEnv Var
76
77                   -- Purely scalar variables. Code which mentions only these
78                   -- variables doesn't have to be lifted.
79                 , global_scalars :: VarSet
80
81                   -- Exported variables which have a vectorised version
82                   --
83                 , global_exported_vars :: VarEnv (Var, Var)
84
85                   -- Mapping from TyCons to their vectorised versions.
86                   -- TyCons which do not have to be vectorised are mapped to
87                   -- themselves.
88                   --
89                 , global_tycons :: NameEnv TyCon
90
91                   -- Mapping from DataCons to their vectorised versions
92                   --
93                 , global_datacons :: NameEnv DataCon
94
95                   -- Mapping from TyCons to their PA dfuns
96                   --
97                 , global_pa_funs :: NameEnv Var
98
99                   -- Mapping from TyCons to their PR dfuns
100                 , global_pr_funs :: NameEnv Var
101
102                   -- Mapping from unboxed TyCons to their boxed versions
103                 , global_boxed_tycons :: NameEnv TyCon
104
105                 -- External package inst-env & home-package inst-env for class
106                 -- instances
107                 --
108                 , global_inst_env :: (InstEnv, InstEnv)
109
110                 -- External package inst-env & home-package inst-env for family
111                 -- instances
112                 --
113                 , global_fam_inst_env :: FamInstEnvs
114
115                 -- Hoisted bindings
116                 , global_bindings :: [(Var, CoreExpr)]
117                 }
118
119 data LocalEnv = LocalEnv {
120                  -- Mapping from local variables to their vectorised and
121                  -- lifted versions
122                  --
123                  local_vars :: VarEnv (Var, Var)
124
125                  -- In-scope type variables
126                  --
127                , local_tyvars :: [TyVar]
128
129                  -- Mapping from tyvars to their PA dictionaries
130                , local_tyvar_pa :: VarEnv CoreExpr
131
132                  -- Local binding name
133                , local_bind_name :: FastString
134                }
135
136 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
137 initGlobalEnv info instEnvs famInstEnvs
138   = GlobalEnv {
139       global_vars          = mapVarEnv snd $ vectInfoVar info
140     , global_scalars   = emptyVarSet
141     , global_exported_vars = emptyVarEnv
142     , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
143     , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
144     , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
145     , global_pr_funs       = emptyNameEnv
146     , global_boxed_tycons  = emptyNameEnv
147     , global_inst_env      = instEnvs
148     , global_fam_inst_env  = famInstEnvs
149     , global_bindings      = []
150     }
151
152 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
153 extendImportedVarsEnv ps genv
154   = genv { global_vars = extendVarEnvList (global_vars genv) ps }
155
156 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
157 extendScalars vs genv
158   = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
159
160 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
161 setFamInstEnv l_fam_inst genv
162   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
163   where
164     (g_fam_inst, _) = global_fam_inst_env genv
165
166 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
167 extendTyConsEnv ps genv
168   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
169
170 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
171 extendDataConsEnv ps genv
172   = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
173
174 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
175 extendPAFunsEnv ps genv
176   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
177
178 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
179 setPRFunsEnv ps genv
180   = genv { global_pr_funs = mkNameEnv ps }
181
182 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
183 setBoxedTyConsEnv ps genv
184   = genv { global_boxed_tycons = mkNameEnv ps }
185
186 emptyLocalEnv :: LocalEnv
187 emptyLocalEnv = LocalEnv {
188                    local_vars     = emptyVarEnv
189                  , local_tyvars   = []
190                  , local_tyvar_pa = emptyVarEnv
191                  , local_bind_name  = fsLit "fn"
192                  }
193
194 -- FIXME
195 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
196 updVectInfo env tyenv info
197   = info {
198       vectInfoVar     = global_exported_vars env
199     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
200     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
201     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
202     }
203   where
204     mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
205                                    | from <- from_tyenv tyenv
206                                    , let name = getName from
207                                    , Just to <- [lookupNameEnv (from_env env) name]]
208
209 data VResult a = Yes GlobalEnv LocalEnv a | No
210
211 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
212
213 instance Monad VM where
214   return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
215   VM p >>= f = VM $ \bi genv lenv -> do
216                                       r <- p bi genv lenv
217                                       case r of
218                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
219                                         No                -> return No
220
221
222 cantVectorise :: String -> SDoc -> a
223 cantVectorise s d = pgmError
224                   . showSDocDump
225                   $ vcat [text "*** Vectorisation error ***",
226                           nest 4 $ sep [text s, nest 4 d]]
227
228 maybeCantVectorise :: String -> SDoc -> Maybe a -> a
229 maybeCantVectorise s d Nothing  = cantVectorise s d
230 maybeCantVectorise _ _ (Just x) = x
231
232 maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
233 maybeCantVectoriseM s d p
234   = do
235       r <- p
236       case r of
237         Just x  -> return x
238         Nothing -> cantVectorise s d
239
240 noV :: VM a
241 noV = VM $ \_ _ _ -> return No
242
243 traceNoV :: String -> SDoc -> VM a
244 traceNoV s d = pprTrace s d noV
245
246 ensureV :: Bool -> VM ()
247 ensureV False = noV
248 ensureV True  = return ()
249
250 onlyIfV :: Bool -> VM a -> VM a
251 onlyIfV b p = ensureV b >> p
252
253 traceEnsureV :: String -> SDoc -> Bool -> VM ()
254 traceEnsureV s d False = traceNoV s d
255 traceEnsureV _ _ True  = return ()
256
257 tryV :: VM a -> VM (Maybe a)
258 tryV (VM p) = VM $ \bi genv lenv ->
259   do
260     r <- p bi genv lenv
261     case r of
262       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
263       No                -> return (Yes genv  lenv  Nothing)
264
265 maybeV :: VM (Maybe a) -> VM a
266 maybeV p = maybe noV return =<< p
267
268 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
269 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
270
271 orElseV :: VM a -> VM a -> VM a
272 orElseV p q = maybe q return =<< tryV p
273
274 fixV :: (a -> VM a) -> VM a
275 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
276   where
277     -- NOTE: It is essential that we are lazy in r above so do not replace
278     --       calls to this function by an explicit case.
279     unYes (Yes _ _ x) = x
280     unYes No          = panic "VectMonad.fixV: no result"
281
282 localV :: VM a -> VM a
283 localV p = do
284              env <- readLEnv id
285              x <- p
286              setLEnv env
287              return x
288
289 closedV :: VM a -> VM a
290 closedV p = do
291               env <- readLEnv id
292               setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
293               x <- p
294               setLEnv env
295               return x
296
297 liftDs :: DsM a -> VM a
298 liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
299
300 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
301 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
302
303 builtin :: (Builtins -> a) -> VM a
304 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
305
306 builtins :: (a -> Builtins -> b) -> VM (a -> b)
307 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
308
309 readGEnv :: (GlobalEnv -> a) -> VM a
310 readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
311
312 setGEnv :: GlobalEnv -> VM ()
313 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
314
315 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
316 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
317
318 readLEnv :: (LocalEnv -> a) -> VM a
319 readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
320
321 setLEnv :: LocalEnv -> VM ()
322 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
323
324 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
325 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
326
327 getInstEnv :: VM (InstEnv, InstEnv)
328 getInstEnv = readGEnv global_inst_env
329
330 getFamInstEnv :: VM FamInstEnvs
331 getFamInstEnv = readGEnv global_fam_inst_env
332
333 getBindName :: VM FastString
334 getBindName = readLEnv local_bind_name
335
336 inBind :: Id -> VM a -> VM a
337 inBind id p
338   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
339        p
340
341 cloneName :: (OccName -> OccName) -> Name -> VM Name
342 cloneName mk_occ name = liftM make (liftDs newUnique)
343   where
344     occ_name = mk_occ (nameOccName name)
345
346     make u | isExternalName name = mkExternalName u (nameModule name)
347                                                     occ_name
348                                                     (nameSrcSpan name)
349            | otherwise           = mkSystemName u occ_name
350
351 cloneId :: (OccName -> OccName) -> Id -> Type -> VM Id
352 cloneId mk_occ id ty
353   = do
354       name <- cloneName mk_occ (getName id)
355       let id' | isExportedId id = Id.mkExportedLocalId name ty
356               | otherwise       = Id.mkLocalId         name ty
357       return id'
358
359 cloneVar :: Var -> VM Var
360 cloneVar var = liftM (setIdUnique var) (liftDs newUnique)
361
362 newExportedVar :: OccName -> Type -> VM Var
363 newExportedVar occ_name ty 
364   = do
365       mod <- liftDs getModuleDs
366       u   <- liftDs newUnique
367
368       let name = mkExternalName u mod occ_name noSrcSpan
369       
370       return $ Id.mkExportedLocalId name ty
371
372 newLocalVar :: FastString -> Type -> VM Var
373 newLocalVar fs ty
374   = do
375       u <- liftDs newUnique
376       return $ mkSysLocal fs u ty
377
378 newDummyVar :: Type -> VM Var
379 newDummyVar = newLocalVar (fsLit "vv")
380
381 newTyVar :: FastString -> Kind -> VM Var
382 newTyVar fs k
383   = do
384       u <- liftDs newUnique
385       return $ mkTyVar (mkSysTvName u fs) k
386
387 defGlobalVar :: Var -> Var -> VM ()
388 defGlobalVar v v' = updGEnv $ \env ->
389   env { global_vars = extendVarEnv (global_vars env) v v'
390       , global_exported_vars = upd (global_exported_vars env)
391       }
392   where
393     upd env | isExportedId v = extendVarEnv env v (v, v')
394             | otherwise      = env
395
396 lookupVar :: Var -> VM (Scope Var (Var, Var))
397 lookupVar v
398   = do
399       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
400       case r of
401         Just e  -> return (Local e)
402         Nothing -> liftM Global
403                 . maybeCantVectoriseM "Variable not vectorised:" (ppr v)
404                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
405
406 globalScalars :: VM VarSet
407 globalScalars = readGEnv global_scalars
408
409 lookupTyCon :: TyCon -> VM (Maybe TyCon)
410 lookupTyCon tc
411   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
412
413   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
414
415 defTyCon :: TyCon -> TyCon -> VM ()
416 defTyCon tc tc' = updGEnv $ \env ->
417   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
418
419 lookupDataCon :: DataCon -> VM (Maybe DataCon)
420 lookupDataCon dc
421   | isTupleTyCon (dataConTyCon dc) = return (Just dc)
422   | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
423
424 defDataCon :: DataCon -> DataCon -> VM ()
425 defDataCon dc dc' = updGEnv $ \env ->
426   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
427
428 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
429 lookupPrimPArray = liftBuiltinDs . primPArray
430
431 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
432 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
433
434 lookupTyConPA :: TyCon -> VM (Maybe Var)
435 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
436
437 defTyConPA :: TyCon -> Var -> VM ()
438 defTyConPA tc pa = updGEnv $ \env ->
439   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
440
441 defTyConPAs :: [(TyCon, Var)] -> VM ()
442 defTyConPAs ps = updGEnv $ \env ->
443   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
444                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
445
446 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
447 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
448
449 lookupTyConPR :: TyCon -> VM (Maybe Var)
450 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
451
452 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
453 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
454                                                        (tyConName tc)
455
456 defLocalTyVar :: TyVar -> VM ()
457 defLocalTyVar tv = updLEnv $ \env ->
458   env { local_tyvars   = tv : local_tyvars env
459       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
460       }
461
462 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
463 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
464   env { local_tyvars   = tv : local_tyvars env
465       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
466       }
467
468 localTyVars :: VM [TyVar]
469 localTyVars = readLEnv (reverse . local_tyvars)
470
471 -- Look up the dfun of a class instance.
472 --
473 -- The match must be unique - ie, match exactly one instance - but the 
474 -- type arguments used for matching may be more specific than those of 
475 -- the class instance declaration.  The found class instances must not have
476 -- any type variables in the instance context that do not appear in the
477 -- instances head (i.e., no flexi vars); for details for what this means,
478 -- see the docs at InstEnv.lookupInstEnv.
479 --
480 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
481 lookupInst cls tys
482   = do { instEnv <- getInstEnv
483        ; case lookupInstEnv instEnv cls tys of
484            ([(inst, inst_tys)], _) 
485              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
486              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: " 
487                                       (ppr $ mkTyConApp (classTyCon cls) tys)
488              where
489                inst_tys'  = [ty | Right ty <- inst_tys]
490                noFlexiVar = all isRight inst_tys
491            _other         ->
492              pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
493        }
494   where
495     isRight (Left  _) = False
496     isRight (Right _) = True
497
498 -- Look up the representation tycon of a family instance.
499 --
500 -- The match must be unique - ie, match exactly one instance - but the 
501 -- type arguments used for matching may be more specific than those of 
502 -- the family instance declaration.
503 --
504 -- Return the instance tycon and its type instance.  For example, if we have
505 --
506 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
507 --
508 -- then we have a coercion (ie, type instance of family instance coercion)
509 --
510 --  :Co:R42T Int :: T [Int] ~ :R42T Int
511 --
512 -- which implies that :R42T was declared as 'data instance T [a]'.
513 --
514 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
515 lookupFamInst tycon tys
516   = ASSERT( isOpenTyCon tycon )
517     do { instEnv <- getFamInstEnv
518        ; case lookupFamInstEnv instEnv tycon tys of
519            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
520            _other                -> 
521              pprPanic "VectMonad.lookupFamInst: not found: " 
522                       (ppr $ mkTyConApp tycon tys)
523        }
524
525 initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
526 initV pkg hsc_env guts info p
527   = do
528          -- XXX: ignores error messages and warnings, check that this is
529          -- indeed ok (the use of "Just r" suggests so)
530       (_,Just r) <- initDs hsc_env (mg_module guts)
531                                (mg_rdr_env guts)
532                                (mg_types guts)
533                                go
534       return r
535   where
536
537     go =
538       do
539         builtins       <- initBuiltins pkg
540         builtin_vars   <- initBuiltinVars builtins
541         builtin_tycons <- initBuiltinTyCons builtins
542         let builtin_datacons = initBuiltinDataCons builtins
543         builtin_pas    <- initBuiltinPAs builtins
544         builtin_prs    <- initBuiltinPRs builtins
545         builtin_boxed  <- initBuiltinBoxedTyCons builtins
546         builtin_scalars <- initBuiltinScalars builtins
547
548         eps <- liftIO $ hscEPS hsc_env
549         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
550             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
551
552         let genv = extendImportedVarsEnv builtin_vars
553                  . extendScalars builtin_scalars
554                  . extendTyConsEnv builtin_tycons
555                  . extendDataConsEnv builtin_datacons
556                  . extendPAFunsEnv builtin_pas
557                  . setPRFunsEnv    builtin_prs
558                  . setBoxedTyConsEnv builtin_boxed
559                  $ initGlobalEnv info instEnvs famInstEnvs
560
561         r <- runVM p builtins genv emptyLocalEnv
562         case r of
563           Yes genv _ x -> return $ Just (new_info genv, x)
564           No           -> return Nothing
565
566     new_info genv = updVectInfo genv (mg_types guts) info
567