Minor refactoring of placeHolderPunRhs
[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, newLocalVars, newDummyVar, newTyVar,
11   
12   Builtins(..), sumTyCon, prodTyCon, prodDataCon,
13   selTy, selReplicate, selPick, selTags, selElements,
14   combinePDVar, scalarZip, closureCtrFun,
15   builtin, builtins,
16
17   GlobalEnv(..),
18   setFamInstEnv,
19   readGEnv, setGEnv, updGEnv,
20
21   LocalEnv(..),
22   readLEnv, setLEnv, updLEnv,
23
24   getBindName, inBind,
25
26   lookupVar, defGlobalVar, globalScalars,
27   lookupTyCon, defTyCon,
28   lookupDataCon, defDataCon,
29   lookupTyConPA, defTyConPA, defTyConPAs,
30   lookupTyConPR,
31   lookupBoxedTyCon,
32   lookupPrimMethod, lookupPrimPArray,
33   lookupTyVarPA, defLocalTyVar, defLocalTyVarWithPA, localTyVars,
34
35   lookupInst, lookupFamInst
36 ) where
37
38 #include "HsVersions.h"
39
40 import VectBuiltIn
41
42 import HscTypes hiding  ( MonadThings(..) )
43 import Module           ( PackageId )
44 import CoreSyn
45 import Class
46 import TyCon
47 import DataCon
48 import Type
49 import Var
50 import VarSet
51 import VarEnv
52 import Id
53 import Name
54 import NameEnv
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 newLocalVars :: FastString -> [Type] -> VM [Var]
379 newLocalVars fs = mapM (newLocalVar fs)
380
381 newDummyVar :: Type -> VM Var
382 newDummyVar = newLocalVar (fsLit "vv")
383
384 newTyVar :: FastString -> Kind -> VM Var
385 newTyVar fs k
386   = do
387       u <- liftDs newUnique
388       return $ mkTyVar (mkSysTvName u fs) k
389
390 defGlobalVar :: Var -> Var -> VM ()
391 defGlobalVar v v' = updGEnv $ \env ->
392   env { global_vars = extendVarEnv (global_vars env) v v'
393       , global_exported_vars = upd (global_exported_vars env)
394       }
395   where
396     upd env | isExportedId v = extendVarEnv env v (v, v')
397             | otherwise      = env
398
399 lookupVar :: Var -> VM (Scope Var (Var, Var))
400 lookupVar v
401   = do
402       r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
403       case r of
404         Just e  -> return (Local e)
405         Nothing -> liftM Global
406                 . maybeCantVectoriseM "Variable not vectorised:" (ppr v)
407                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
408
409 globalScalars :: VM VarSet
410 globalScalars = readGEnv global_scalars
411
412 lookupTyCon :: TyCon -> VM (Maybe TyCon)
413 lookupTyCon tc
414   | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc)
415
416   | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
417
418 defTyCon :: TyCon -> TyCon -> VM ()
419 defTyCon tc tc' = updGEnv $ \env ->
420   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
421
422 lookupDataCon :: DataCon -> VM (Maybe DataCon)
423 lookupDataCon dc
424   | isTupleTyCon (dataConTyCon dc) = return (Just dc)
425   | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
426
427 defDataCon :: DataCon -> DataCon -> VM ()
428 defDataCon dc dc' = updGEnv $ \env ->
429   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
430
431 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
432 lookupPrimPArray = liftBuiltinDs . primPArray
433
434 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
435 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
436
437 lookupTyConPA :: TyCon -> VM (Maybe Var)
438 lookupTyConPA tc = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
439
440 defTyConPA :: TyCon -> Var -> VM ()
441 defTyConPA tc pa = updGEnv $ \env ->
442   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
443
444 defTyConPAs :: [(TyCon, Var)] -> VM ()
445 defTyConPAs ps = updGEnv $ \env ->
446   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
447                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
448
449 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
450 lookupTyVarPA tv = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
451
452 lookupTyConPR :: TyCon -> VM (Maybe Var)
453 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
454
455 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
456 lookupBoxedTyCon tc = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
457                                                        (tyConName tc)
458
459 defLocalTyVar :: TyVar -> VM ()
460 defLocalTyVar tv = updLEnv $ \env ->
461   env { local_tyvars   = tv : local_tyvars env
462       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
463       }
464
465 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
466 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
467   env { local_tyvars   = tv : local_tyvars env
468       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
469       }
470
471 localTyVars :: VM [TyVar]
472 localTyVars = readLEnv (reverse . local_tyvars)
473
474 -- Look up the dfun of a class instance.
475 --
476 -- The match must be unique - ie, match exactly one instance - but the 
477 -- type arguments used for matching may be more specific than those of 
478 -- the class instance declaration.  The found class instances must not have
479 -- any type variables in the instance context that do not appear in the
480 -- instances head (i.e., no flexi vars); for details for what this means,
481 -- see the docs at InstEnv.lookupInstEnv.
482 --
483 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
484 lookupInst cls tys
485   = do { instEnv <- getInstEnv
486        ; case lookupInstEnv instEnv cls tys of
487            ([(inst, inst_tys)], _) 
488              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
489              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: " 
490                                       (ppr $ mkTyConApp (classTyCon cls) tys)
491              where
492                inst_tys'  = [ty | Right ty <- inst_tys]
493                noFlexiVar = all isRight inst_tys
494            _other         ->
495              pprPanic "VectMonad.lookupInst: not found " (ppr cls <+> ppr tys)
496        }
497   where
498     isRight (Left  _) = False
499     isRight (Right _) = True
500
501 -- Look up the representation tycon of a family instance.
502 --
503 -- The match must be unique - ie, match exactly one instance - but the 
504 -- type arguments used for matching may be more specific than those of 
505 -- the family instance declaration.
506 --
507 -- Return the instance tycon and its type instance.  For example, if we have
508 --
509 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
510 --
511 -- then we have a coercion (ie, type instance of family instance coercion)
512 --
513 --  :Co:R42T Int :: T [Int] ~ :R42T Int
514 --
515 -- which implies that :R42T was declared as 'data instance T [a]'.
516 --
517 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
518 lookupFamInst tycon tys
519   = ASSERT( isOpenTyCon tycon )
520     do { instEnv <- getFamInstEnv
521        ; case lookupFamInstEnv instEnv tycon tys of
522            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
523            _other                -> 
524              pprPanic "VectMonad.lookupFamInst: not found: " 
525                       (ppr $ mkTyConApp tycon tys)
526        }
527
528 initV :: PackageId -> HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
529 initV pkg hsc_env guts info p
530   = do
531          -- XXX: ignores error messages and warnings, check that this is
532          -- indeed ok (the use of "Just r" suggests so)
533       (_,Just r) <- initDs hsc_env (mg_module guts)
534                                (mg_rdr_env guts)
535                                (mg_types guts)
536                                go
537       return r
538   where
539
540     go =
541       do
542         builtins       <- initBuiltins pkg
543         builtin_vars   <- initBuiltinVars builtins
544         builtin_tycons <- initBuiltinTyCons builtins
545         let builtin_datacons = initBuiltinDataCons builtins
546         builtin_boxed  <- initBuiltinBoxedTyCons builtins
547         builtin_scalars <- initBuiltinScalars builtins
548
549         eps <- liftIO $ hscEPS hsc_env
550         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
551             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
552
553         builtin_prs    <- initBuiltinPRs builtins instEnvs
554         builtin_pas    <- initBuiltinPAs builtins instEnvs
555
556         let genv = extendImportedVarsEnv builtin_vars
557                  . extendScalars builtin_scalars
558                  . extendTyConsEnv builtin_tycons
559                  . extendDataConsEnv builtin_datacons
560                  . extendPAFunsEnv builtin_pas
561                  . setPRFunsEnv    builtin_prs
562                  . setBoxedTyConsEnv builtin_boxed
563                  $ initGlobalEnv info instEnvs famInstEnvs
564
565         r <- runVM p builtins genv emptyLocalEnv
566         case r of
567           Yes genv _ x -> return $ Just (new_info genv, x)
568           No           -> return Nothing
569
570     new_info genv = updVectInfo genv (mg_types guts) info
571