e41609aff9323accf6bef1cb5ce15ab05a1ff9e9
[ghc-hetmet.git] / ghc / compiler / usageSP / UsageSPUtils.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1998
3 %
4 \section[UsageSPUtils]{UsageSP Utilities}
5
6 This code is (based on) PhD work of Keith Wansbrough <kw217@cl.cam.ac.uk>,
7 September 1998 .. May 1999.
8
9 Keith Wansbrough 1998-09-04..1999-07-07
10
11 \begin{code}
12 module UsageSPUtils ( AnnotM(AnnotM), initAnnotM,
13                       genAnnotBinds,
14                       MungeFlags(isSigma,isLocal,isExp,hasUsg,mfLoc),
15
16                       doAnnotBinds, doUnAnnotBinds,
17                       annotTy, annotTyN, annotMany, annotManyN, unannotTy, freshannotTy,
18
19                       newVarUs, newVarUSMM,
20                       UniqSMM, usToUniqSMM, uniqSMMToUs,
21
22                       primOpUsgTys,
23                     ) where
24
25 #include "HsVersions.h"
26
27 import CoreSyn
28 import Const            ( Con(..), Literal(..) )
29 import Var              ( IdOrTyVar, varName, varType, setVarType, mkUVar )
30 import Id               ( idMustBeINLINEd, isExportedId )
31 import Name             ( isLocallyDefined )
32 import TypeRep          ( Type(..), TyNote(..) )  -- friend
33 import Type             ( UsageAnn(..), isUsgTy, splitFunTys )
34 import Subst            ( substTy, mkTyVarSubst )
35 import TyCon            ( isAlgTyCon, isPrimTyCon, isSynTyCon, isFunTyCon )
36 import VarEnv
37 import PrimOp           ( PrimOp, primOpUsg )
38 import Maybes           ( expectJust )
39 import UniqSupply       ( UniqSupply, UniqSM, initUs, getUniqueUs, thenUs, returnUs )
40 import Outputable
41 import PprCore          ( )  -- instances only
42 \end{code}
43
44 ======================================================================
45
46 Walking over (and altering) types
47 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
48
49 We often need to fiddle with (i.e., add or remove) usage annotations
50 on a type.  We define here a general framework to do this.  Usage
51 annotations come from any monad with a function @getAnnM@ which yields
52 a new annotation.  We use two mutually recursive functions, one for
53 sigma types and one for tau types.
54
55 \begin{code}
56 genAnnotTy :: Monad m =>
57               (m UsageAnn)  -- get new annotation
58            -> Type          -- old type
59            -> m Type        -- new type
60
61 genAnnotTy getAnnM ty = do { u   <- getAnnM
62                            ; ty' <- genAnnotTyN getAnnM ty
63                            ; return (NoteTy (UsgNote u) ty')
64                            }
65
66 genAnnotTyN :: Monad m =>
67                (m UsageAnn)
68             -> Type
69             -> m Type
70
71 genAnnotTyN getAnnM
72   (NoteTy (UsgNote _) ty)     = panic "genAnnotTyN: unexpected UsgNote"
73 genAnnotTyN getAnnM
74   (NoteTy (SynNote sty) ty)   = do { sty' <- genAnnotTyN getAnnM sty
75                                 -- is this right? shouldn't there be some
76                                 -- correlation between sty' and ty'?
77                                 -- But sty is a TyConApp; does this make it safer?
78                                    ; ty'  <- genAnnotTyN getAnnM ty
79                                    ; return (NoteTy (SynNote sty') ty')
80                                    }
81 genAnnotTyN getAnnM
82   (NoteTy fvn@(FTVNote _) ty) = do { ty' <- genAnnotTyN getAnnM ty
83                                    ; return (NoteTy fvn ty')
84                                    }
85
86 genAnnotTyN getAnnM
87   ty0@(TyVarTy _)             = do { return ty0 }
88
89 genAnnotTyN getAnnM
90   (AppTy ty1 ty2)             = do { ty1' <- genAnnotTyN getAnnM ty1
91                                    ; ty2' <- genAnnotTyN getAnnM ty2
92                                    ; return (AppTy ty1' ty2')
93                                    }
94
95 genAnnotTyN getAnnM
96   (TyConApp tc tys)           = ASSERT( isFunTyCon tc || isAlgTyCon tc || isPrimTyCon tc || isSynTyCon tc )
97                                 do { let gAT = if isFunTyCon tc
98                                                then genAnnotTy  -- sigma for partial apps of (->)
99                                                else genAnnotTyN -- tau otherwise
100                                    ; tys' <- mapM (gAT getAnnM) tys
101                                    ; return (TyConApp tc tys')
102                                    }
103
104 genAnnotTyN getAnnM
105   (FunTy ty1 ty2)             = do { ty1' <- genAnnotTy getAnnM ty1
106                                    ; ty2' <- genAnnotTy getAnnM ty2
107                                    ; return (FunTy ty1' ty2')
108                                    }
109
110 genAnnotTyN getAnnM
111   (ForAllTy v ty)             = do { ty' <- genAnnotTyN getAnnM ty
112                                    ; return (ForAllTy v ty')
113                                    }
114 \end{code}
115
116
117
118 Walking over (and retyping) terms
119 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
120
121 We also often need to play with the types in a term.  This is slightly
122 tricky because of redundancy: we want to change binder types, and keep
123 the bound types matching these; then there's a special case also with
124 non-locally-defined bound variables.  We generalise over all this
125 here.
126
127 The name `annot' is a bit of a misnomer, as the code is parameterised
128 over exactly what it does to the types (and certain terms).  Notice
129 also that it is possible for this parameter to use
130 monadically-threaded state: here called `flexi'.  For genuine
131 annotation, this state will be a UniqSupply.
132
133 We may add annotations to the outside of a (term, not type) lambda; a
134 function passed to @genAnnotBinds@ does this, taking the lambda and
135 returning the annotated lambda.  It is inside the @AnnotM@ monad.
136 This term-munging function is applied when we see either a term lambda
137 or a usage annotation; *IMPORTANT:* it is applied *before* we recurse
138 down into the term, and it is expected to work only at the top level.
139 Recursion will subsequently be done by genAnnotBinds.  It may
140 optionally remove a Note TermUsg, or optionally add one if it is not
141 already present, but it may perform NO OTHER MODIFICATIONS to the
142 structure of the term.
143
144 We do different things to types of variables bound locally and of
145 variables bound in other modules, in certain cases: the former get
146 uvars and the latter keep their existing annotations when we annotate,
147 for example.  To control this, @MungeFlags@ describes what kind of a
148 type this is that we're about to munge.
149
150 \begin{code}
151 data MungeFlags = MungeFlags { isSigma :: Bool,  -- want annotated on top (sigma type)
152                                isLocal :: Bool,  -- is locally-defined type
153                                hasUsg  :: Bool,  -- has fixed usage info, don't touch
154                                isExp   :: Bool,  -- is exported (and must be pessimised)
155                                mfLoc   :: SDoc   -- location info
156                              }
157
158 tauTyMF loc  = MungeFlags { isSigma = False, isLocal = True,
159                             hasUsg = False,  isExp = False,  mfLoc = loc }
160 sigVarTyMF v = MungeFlags { isSigma = True,  isLocal = hasLocalDef v, 
161                             hasUsg = hasUsgInfo v, isExp = isExportedId v,
162                             mfLoc = ptext SLIT("type of binder") <+> ppr v }
163 \end{code}
164
165 The helper functions @tauTyMF@ and @sigVarTyMF@ create @MungeFlags@
166 for us.  @sigVarTyMF@ checks the variable to see how to set the flags.
167
168 @hasLocalDef@ tells us if the given variable has an actual local
169 definition that we can play with.  This is not quite the same as
170 @isLocallyDefined@, since @IMustBeINLINEd@ things (usually) don't have
171 a local definition - the simplifier will inline whatever their
172 unfolding is anyway.  We treat these as if they were externally
173 defined, since we don't have access to their definition (at least not
174 easily).  This doesn't hurt much, since after the simplifier has run
175 the unfolding will have been inlined and we can access the unfolding
176 directly.
177
178 @hasUsgInfo@, on the other hand, says if the variable already has
179 usage info in its type that must at all costs be preserved.  This is
180 assumed true (exactly) of all imported ids.
181
182 \begin{code}
183 hasLocalDef :: IdOrTyVar -> Bool
184 hasLocalDef var = isLocallyDefined var
185                   && not (idMustBeINLINEd var)
186
187 hasUsgInfo :: IdOrTyVar -> Bool
188 hasUsgInfo var = (not . isLocallyDefined) var
189 \end{code}
190
191 Here's the walk itself.
192
193 \begin{code}
194 genAnnotBinds :: (MungeFlags -> Type -> AnnotM flexi Type)
195               -> (CoreExpr -> AnnotM flexi CoreExpr)       -- see caveats above
196               -> [CoreBind]
197               -> AnnotM flexi [CoreBind]
198
199 genAnnotBinds _ _ []     = return []
200
201 genAnnotBinds f g (b:bs) = do { (b',vs,vs') <- genAnnotBind f g b
202                               ; bs' <- withAnnVars vs vs' $
203                                          genAnnotBinds f g bs
204                               ; return (b':bs')
205                               }
206
207 genAnnotBind :: (MungeFlags -> Type -> AnnotM flexi Type)  -- type-altering function
208              -> (CoreExpr -> AnnotM flexi CoreExpr)        -- term-altering function
209              -> CoreBind                          -- original CoreBind
210              -> AnnotM flexi
211                        (CoreBind,                 -- annotated CoreBind
212                         [IdOrTyVar],              -- old variables, to be mapped to...
213                         [IdOrTyVar])              -- ... new variables
214
215 genAnnotBind f g (NonRec v1 e1) = do { v1' <- genAnnotVar f v1
216                                      ; e1' <- genAnnotCE f g e1
217                                      ; return (NonRec v1' e1', [v1], [v1'])
218                                      }
219
220 genAnnotBind f g (Rec ves)      = do { let (vs,es) = unzip ves
221                                      ; vs' <- mapM (genAnnotVar f) vs
222                                      ; es' <- withAnnVars vs vs' $
223                                                 mapM (genAnnotCE f g) es
224                                      ; return (Rec (zip vs' es'), vs, vs')
225                                      }
226
227 genAnnotCE :: (MungeFlags -> Type -> AnnotM flexi Type)  -- type-altering function
228            -> (CoreExpr -> AnnotM flexi CoreExpr)        -- term-altering function
229            -> CoreExpr                             -- original expression
230            -> AnnotM flexi CoreExpr                -- yields new expression
231
232 genAnnotCE mungeType mungeTerm = go
233   where go e0@(Var v) | isTyVar v    = return e0  -- arises, e.g., as tyargs of Con
234                                                   -- (no it doesn't: (Type (TyVar tyvar))
235                       | otherwise    = do { mv' <- lookupAnnVar v
236                                           ; v'  <- case mv' of
237                                                      Just var -> return var
238                                                      Nothing  -> fixedVar v
239                                           ; return (Var v')
240                                           }
241
242         go (Con c args)              = -- we know it's saturated
243                                        do { args' <- mapM go args
244                                           ; return (Con c args')
245                                           }
246
247         go (App e arg)               = do { e' <- go e
248                                           ; arg' <- go arg
249                                           ; return (App e' arg')
250                                           }
251
252         go e0@(Lam v0 _)              = do { e1 <- (if isTyVar v0 then return else mungeTerm) e0
253                                           ; let (v,e2,wrap)
254                                                   = case e1 of  -- munge may have added note
255                                                       Note tu@(TermUsg _) (Lam v e2)
256                                                                -> (v,e2,Note tu)
257                                                       Lam v e2 -> (v,e2,id)
258                                           ; v' <- genAnnotVar mungeType v
259                                           ; e' <- withAnnVar v v' $ go e2
260                                           ; return (wrap (Lam v' e'))
261                                           }
262
263         go (Let bind e)              = do { (bind',vs,vs') <- genAnnotBind mungeType mungeTerm bind
264                                           ; e' <- withAnnVars vs vs' $ go e
265                                           ; return (Let bind' e')
266                                           }
267
268         go (Case e v alts)           = do { e' <- go e
269                                           ; v' <- genAnnotVar mungeType v
270                                           ; alts' <- withAnnVar v v' $ mapM genAnnotAlt alts
271                                           ; return (Case e' v' alts')
272                                           }
273
274         go (Note scc@(SCC _)      e) = do { e' <- go e
275                                           ; return (Note scc e')
276                                           }
277         go e0@(Note (Coerce ty1 ty0)
278                                   e) = do { ty1' <- mungeType
279                                                       (tauTyMF (ptext SLIT("coercer of")
280                                                                 <+> ppr e0)) ty1
281                                           ; ty0' <- mungeType
282                                                       (tauTyMF (ptext SLIT("coercee of")
283                                                                 <+> ppr e0)) ty0
284                                                  -- (Better to specify ty0'
285                                                  --  identical to the type of e, including
286                                                  --  annotations, right at the beginning, but
287                                                  --  not possible at this point.)
288                                           ; e' <- go e
289                                           ; return (Note (Coerce ty1' ty0') e')
290                                           }
291         go (Note InlineCall       e) = do { e' <- go e
292                                           ; return (Note InlineCall e')
293                                           }
294         go (Note InlineMe         e) = do { e' <- go e
295                                           ; return (Note InlineMe e')
296                                           }
297         go e0@(Note (TermUsg _)   _) = do { e1 <- mungeTerm e0
298                                           ; case e1 of  -- munge may have removed note
299                                               Note tu@(TermUsg _) e2 -> do { e3 <- go e2
300                                                                            ; return (Note tu e3)
301                                                                            }
302                                               e2                     -> go e2
303                                           }
304
305         go e0@(Type ty)              = -- should only occur at toplevel of Arg,
306                                        -- hence tau-type
307                                        do { ty' <- mungeType
308                                                      (tauTyMF (ptext SLIT("tyarg")
309                                                                <+> ppr e0)) ty
310                                           ; return (Type ty')
311                                           }
312
313         fixedVar v = ASSERT2( not (hasLocalDef v), text "genAnnotCE: locally defined var" <+> ppr v <+> text "not in varenv" )
314                      genAnnotVar mungeType v
315
316         genAnnotAlt (c,vs,e)         = do { vs' <- mapM (genAnnotVar mungeType) vs
317                                           ; e' <- withAnnVars vs vs' $ go e
318                                           ; return (c, vs', e')
319                                           }
320
321
322 genAnnotVar :: (MungeFlags -> Type -> AnnotM flexi Type)
323             -> IdOrTyVar
324             -> AnnotM flexi IdOrTyVar
325
326 genAnnotVar mungeType v | isTyVar v = return v
327                         | otherwise = do { vty' <- mungeType (sigVarTyMF v) (varType v)
328                                          ; return (setVarType v vty')
329                                          }
330 {- ifdef DEBUG
331                                          ; return $
332                                              pprTrace "genAnnotVar" (ppr (tyUsg vty') <+> ppr v) $
333                                              (setVarType v vty')
334    endif
335  -}
336 \end{code}
337
338 ======================================================================
339
340 Some specific things to do to types inside terms
341 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
342
343 @annotTyM@ annotates a type with fresh uvars everywhere the inference
344 is allowed to go, and leaves alone annotations where it may not go.
345
346 We assume there are no annotations already.
347
348 \begin{code}
349 annotTyM :: MungeFlags -> Type -> AnnotM UniqSupply Type
350 -- general function
351 annotTyM mf ty = uniqSMtoAnnotM . uniqSMMToUs $
352                    case (hasUsg mf, isLocal mf, isSigma mf) of
353                      (True ,_    ,_    ) -> ASSERT( isUsgTy ty )
354                                             return ty
355                      (False,True ,True ) -> if isExp mf then
356                                               annotTyP (tag 'p') ty
357                                             else
358                                               annotTy (tag 's') ty
359                      (False,True ,False) -> annotTyN (tag 't') ty
360                      (False,False,True ) -> return $ annotMany  ty -- assume worst
361                      (False,False,False) -> return $ annotManyN ty
362   where tag c = Right $ "annotTyM:" ++ [c] ++ ": " ++ showSDoc (ppr ty)
363
364 -- specific functions for annotating tau and sigma types
365
366 -- ...with uvars
367 annotTy  tag = genAnnotTy  (newVarUSMM tag)
368 annotTyN tag = genAnnotTyN (newVarUSMM tag)
369
370 -- ...with uvars and pessimal Manys (for exported ids)
371 annotTyP tag ty = do { ty' <- annotTy tag ty ; return (pessimise True ty') }
372
373 -- ...with Many
374 annotMany, annotManyN :: Type -> Type
375 #ifndef USMANY
376 annotMany  = id
377 annotManyN = id
378 #else
379 annotMany  ty = unId (genAnnotTy  (return UsMany) ty)
380 annotManyN ty = unId (genAnnotTyN (return UsMany) ty)
381 #endif
382
383 -- monad required for the above
384 newtype Id a = Id a ; unId (Id a) = a
385 instance Monad Id where { a >>= f  = f (unId a) ; return a = Id a }
386
387 -- lambda-annotating function for use along with the above
388 annotLam e0@(Lam v e) = do { uv <- uniqSMtoAnnotM $ newVarUs (Left e0)
389                            ; return (Note (TermUsg uv) (Lam v e))
390                            }
391 annotLam (Note (TermUsg _) _) = panic "annotLam: unexpected term usage annot"
392 \end{code}
393
394 The above requires a `pessimising' translation.  This is applied to
395 types of exported ids, and ensures that they have a fully general
396 type (since we don't know how they will be used in other modules).
397
398 \begin{code}
399 pessimise :: Bool -> Type -> Type
400
401 #ifndef USMANY
402 pessimise  co ty0@(NoteTy  usg@(UsgNote u  ) ty)
403   = if co
404     then case u of UsMany  -> pty
405                    UsVar _ -> pty  -- force to UsMany
406                    UsOnce  -> pprPanic "pessimise:" (ppr ty0)
407     else NoteTy usg pty
408   where pty = pessimiseN co ty
409                  
410 pessimise  co ty0 = pessimiseN co ty0  -- assume UsMany
411 #else
412 pessimise  co ty0@(NoteTy  usg@(UsgNote u  ) ty)
413   = if co
414     then case u of UsMany  -> NoteTy usg pty
415                    UsVar _ -> NoteTy (UsgNote UsMany) pty
416                    UsOnce  -> pprPanic "pessimise:" (ppr ty0)
417     else NoteTy usg pty
418   where pty = pessimiseN co ty
419                  
420 pessimise  co ty0                                = pprPanic "pessimise: missing usage note:" $
421                                                             ppr ty0
422 #endif
423
424 pessimiseN co ty0@(NoteTy  usg@(UsgNote _  ) ty) = pprPanic "pessimiseN: unexpected usage note:" $
425                                                             ppr ty0
426 pessimiseN co     (NoteTy      (SynNote sty) ty) = NoteTy (SynNote (pessimiseN co sty))
427                                                                    (pessimiseN co ty )
428 pessimiseN co     (NoteTy note@(FTVNote _  ) ty) = NoteTy note (pessimiseN co ty)
429 pessimiseN co ty0@(TyVarTy _)                    = ty0
430 pessimiseN co ty0@(AppTy _ _)                    = ty0
431 pessimiseN co ty0@(TyConApp tc tys)              = ASSERT( not ((isFunTyCon tc) && (length tys > 1)) )
432                                                    ty0
433 pessimiseN co     (FunTy ty1 ty2)                = FunTy (pessimise (not co) ty1)
434                                                          (pessimise      co  ty2)
435 pessimiseN co     (ForAllTy tyv ty)              = ForAllTy tyv (pessimiseN co ty)
436 \end{code}
437
438
439 @unAnnotTyM@ strips annotations (that the inference is allowed to
440 touch) from a term, and `fixes' those it isn't permitted to touch (by
441 putting @Many@ annotations where they are missing, but leaving
442 existing annotations in the type).
443
444 @unTermUsg@ removes from a term any term usage annotations it finds.
445
446 \begin{code}
447 unAnnotTyM :: MungeFlags -> Type -> AnnotM a Type
448
449 unAnnotTyM mf ty = if hasUsg mf then
450                      ASSERT( isSigma mf )
451                      return (fixAnnotTy ty)
452                    else return (unannotTy ty)
453
454
455 unTermUsg :: CoreExpr -> AnnotM a CoreExpr
456 -- strip all term annotations
457 unTermUsg e@(Lam _ _)          = return e
458 unTermUsg (Note (TermUsg _) e) = return e
459 unTermUsg _                    = panic "unTermUsg"
460
461 unannotTy :: Type -> Type
462 -- strip all annotations
463 unannotTy    (NoteTy     (UsgForAll uv) ty) = unannotTy ty
464 unannotTy    (NoteTy      (UsgNote _  ) ty) = unannotTy ty
465 unannotTy    (NoteTy      (SynNote sty) ty) = NoteTy (SynNote (unannotTy sty)) (unannotTy ty)
466 unannotTy    (NoteTy note@(FTVNote _  ) ty) = NoteTy note (unannotTy ty)
467 unannotTy ty@(TyVarTy _)                    = ty
468 unannotTy    (AppTy ty1 ty2)                = AppTy (unannotTy ty1) (unannotTy ty2)
469 unannotTy    (TyConApp tc tys)              = TyConApp tc (map unannotTy tys)
470 unannotTy    (FunTy ty1 ty2)                = FunTy (unannotTy ty1) (unannotTy ty2)
471 unannotTy    (ForAllTy tyv ty)              = ForAllTy tyv (unannotTy ty)
472
473
474 fixAnnotTy :: Type -> Type
475 -- put Manys where they are missing
476 #ifndef USMANY
477 fixAnnotTy = id
478 #else
479 fixAnnotTy     (NoteTy note@(UsgForAll uv) ty) = NoteTy note (fixAnnotTy  ty)
480 fixAnnotTy      (NoteTy note@(UsgNote _  ) ty) = NoteTy note (fixAnnotTyN ty)
481 fixAnnotTy  ty0                                = NoteTy (UsgNote UsMany) (fixAnnotTyN ty0)
482
483 fixAnnotTyN ty0@(NoteTy note@(UsgNote _  ) ty) = pprPanic "fixAnnotTyN: unexpected usage note:" $
484                                                           ppr ty0
485 fixAnnotTyN     (NoteTy      (SynNote sty) ty) = NoteTy (SynNote (fixAnnotTyN sty))
486                                                                  (fixAnnotTyN ty )
487 fixAnnotTyN     (NoteTy note@(FTVNote _  ) ty) = NoteTy note (fixAnnotTyN ty)
488 fixAnnotTyN ty0@(TyVarTy _)                    = ty0
489 fixAnnotTyN     (AppTy ty1 ty2)                = AppTy (fixAnnotTyN ty1) (fixAnnotTyN ty2)
490 fixAnnotTyN     (TyConApp tc tys)              = ASSERT( isFunTyCon tc || isAlgTyCon tc || isPrimTyCon tc || isSynTyCon tc )
491                                                  TyConApp tc (map (if isFunTyCon tc then
492                                                                      fixAnnotTy
493                                                                    else
494                                                                      fixAnnotTyN) tys)
495 fixAnnotTyN     (FunTy ty1 ty2)                = FunTy (fixAnnotTy ty1) (fixAnnotTy ty2)
496 fixAnnotTyN     (ForAllTy tyv ty)              = ForAllTy tyv (fixAnnotTyN ty)
497 #endif
498 \end{code}
499
500 The composition (reannotating a type with fresh uvars but the same
501 structure) is useful elsewhere:
502
503 \begin{code}
504 freshannotTy :: Type -> UniqSMM Type
505 freshannotTy = annotTy (Right "freshannotTy") . unannotTy
506 \end{code}
507
508
509 Wrappers apply these functions to sets of bindings.
510
511 \begin{code}
512 doAnnotBinds :: UniqSupply
513              -> [CoreBind]
514              -> ([CoreBind],UniqSupply)
515
516 doAnnotBinds us binds = initAnnotM us (genAnnotBinds annotTyM annotLam binds)
517
518
519 doUnAnnotBinds :: [CoreBind]
520                -> [CoreBind]
521
522 doUnAnnotBinds binds = fst $ initAnnotM () $
523                          genAnnotBinds unAnnotTyM unTermUsg binds
524 \end{code}
525
526 ======================================================================
527
528 Monadic machinery
529 ~~~~~~~~~~~~~~~~~
530
531 The @UniqSM@ type is not an instance of @Monad@, and cannot be made so
532 since it is merely a synonym rather than a newtype.  Here we define
533 @UniqSMM@, which *is* an instance of @Monad@.
534
535 \begin{code}
536 newtype UniqSMM a = UsToUniqSMM (UniqSM a)
537 uniqSMMToUs (UsToUniqSMM us) = us
538 usToUniqSMM = UsToUniqSMM
539
540 instance Monad UniqSMM where
541   m >>= f  = UsToUniqSMM $ uniqSMMToUs m `thenUs` \ a ->
542                            uniqSMMToUs (f a)
543   return   = UsToUniqSMM . returnUs
544 \end{code}
545
546
547 For annotation, the monad @AnnotM@, we need to carry around our
548 variable mapping, along with some general state.
549
550 \begin{code}
551 newtype AnnotM flexi a = AnnotM (   flexi                     -- UniqSupply etc
552                                   -> VarEnv IdOrTyVar         -- unannotated to annotated variables
553                                   -> (a,flexi,VarEnv IdOrTyVar))
554 unAnnotM (AnnotM f) = f
555
556 instance Monad (AnnotM flexi) where
557   a >>= f  = AnnotM (\ us ve -> let (r,us',ve') = unAnnotM a us ve
558                                 in  unAnnotM (f r) us' ve')
559   return a = AnnotM (\ us ve -> (a,us,ve))
560
561 initAnnotM :: fl -> AnnotM fl a -> (a,fl)
562 initAnnotM fl m = case (unAnnotM m) fl emptyVarEnv of { (r,fl',_) -> (r,fl') }
563
564 withAnnVar :: IdOrTyVar -> IdOrTyVar -> AnnotM fl a -> AnnotM fl a
565 withAnnVar v v' m = AnnotM (\ us ve -> let ve'          = extendVarEnv ve v v'
566                                            (r,us',_)    = (unAnnotM m) us ve'
567                                        in  (r,us',ve))
568
569 withAnnVars :: [IdOrTyVar] -> [IdOrTyVar] -> AnnotM fl a -> AnnotM fl a
570 withAnnVars vs vs' m = AnnotM (\ us ve -> let ve'          = plusVarEnv ve (zipVarEnv vs vs')
571                                               (r,us',_)    = (unAnnotM m) us ve'
572                                           in  (r,us',ve))
573
574 lookupAnnVar :: IdOrTyVar -> AnnotM fl (Maybe IdOrTyVar)
575 lookupAnnVar var = AnnotM (\ us ve -> (lookupVarEnv ve var,
576                                        us,
577                                        ve))
578 \end{code}
579
580 A useful helper allows us to turn a computation in the unique supply
581 monad into one in the annotation monad parameterised by a unique
582 supply.
583
584 \begin{code}
585 uniqSMtoAnnotM :: UniqSM a -> AnnotM UniqSupply a
586
587 uniqSMtoAnnotM m = AnnotM (\ us ve -> let (r,us') = initUs us m
588                                       in  (r,us',ve))
589 \end{code}
590
591 @newVarUs@ and @newVarUSMM@ generate a new usage variable.  They take
592 an argument which is used for debugging only, describing what the
593 variable is to annotate.
594
595 \begin{code}
596 newVarUs :: (Either CoreExpr String) -> UniqSM UsageAnn
597 -- the first arg is for debugging use only
598 newVarUs e = getUniqueUs `thenUs` \ u ->
599              let uv = mkUVar u in
600              returnUs (UsVar uv)
601 {- #ifdef DEBUG
602              let src = case e of
603                          Left (Con (Literal _) _) -> "literal"
604                          Left (Con _           _) -> "primop"
605                          Left (Lam v e)           -> "lambda: " ++ showSDoc (ppr v)
606                          Left _                   -> "unknown"
607                          Right s                  -> s
608              in pprTrace "newVarUs:" (ppr uv <+> text src) $
609    #endif
610  -}
611
612 newVarUSMM :: (Either CoreExpr String) -> UniqSMM UsageAnn
613 newVarUSMM = usToUniqSMM . newVarUs
614 \end{code}
615
616 ======================================================================
617
618 PrimOps and usage information.
619
620 Analagously to @DataCon.dataConArgTys@, we determine the argtys and
621 result ty of a primop, *after* substition (which may reveal more args,
622 notably for @CCall@s).
623
624 \begin{code}
625 primOpUsgTys :: PrimOp         -- this primop
626              -> [Type]         -- instantiated at these (tau) types
627              -> ([Type],Type)  -- requires args of these (sigma) types,
628                                --  and returns this (sigma) type
629
630 primOpUsgTys p tys = let (tyvs,ty0us,rtyu) = primOpUsg p
631                          s                 = mkTyVarSubst tyvs tys
632                          (ty1us,rty1u)     = splitFunTys (substTy s rtyu)
633                                              -- substitution may reveal more args
634                      in  ((map (substTy s) ty0us) ++ ty1us,
635                           rty1u)
636 \end{code}
637
638 ======================================================================
639
640 EOF