2 % (c) The GRASP/AQUA Project, Glasgow University, 1998
4 \section[UsageSPUtils]{UsageSP Utilities}
6 This code is (based on) PhD work of Keith Wansbrough <kw217@cl.cam.ac.uk>,
7 September 1998 .. May 1999.
9 Keith Wansbrough 1998-09-04..1999-05-07
12 module UsageSPUtils ( AnnotM(AnnotM), initAnnotM,
14 MungeFlags(isSigma,isLocal,isExp,hasUsg,mfLoc),
16 doAnnotBinds, doUnAnnotBinds,
17 annotMany, annotManyN, unannotTy, freshannotTy,
20 UniqSMM, usToUniqSMM, uniqSMMToUs,
25 #include "HsVersions.h"
28 import Const ( Con(..), Literal(..) )
29 import Var ( IdOrTyVar, varName, varType, setVarType, mkUVar )
30 import Id ( idMustBeINLINEd )
31 import Name ( isLocallyDefined, isExported )
32 import Type ( Type(..), TyNote(..), UsageAnn(..), isUsgTy, substTy, splitFunTys )
33 import TyCon ( isAlgTyCon, isPrimTyCon, isSynTyCon, isFunTyCon )
35 import PrimOp ( PrimOp, primOpUsg )
36 import Maybes ( expectJust )
37 import UniqSupply ( UniqSupply, UniqSM, initUs, getUniqueUs, thenUs, returnUs )
39 import PprCore ( ) -- instances only
42 ======================================================================
44 Walking over (and altering) types
45 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
47 We often need to fiddle with (i.e., add or remove) usage annotations
48 on a type. We define here a general framework to do this. Usage
49 annotations come from any monad with a function @getAnnM@ which yields
50 a new annotation. We use two mutually recursive functions, one for
51 sigma types and one for tau types.
54 genAnnotTy :: Monad m =>
55 (m UsageAnn) -- get new annotation
59 genAnnotTy getAnnM ty = do { u <- getAnnM
60 ; ty' <- genAnnotTyN getAnnM ty
61 ; return (NoteTy (UsgNote u) ty')
64 genAnnotTyN :: Monad m =>
70 (NoteTy (UsgNote _) ty) = panic "genAnnotTyN: unexpected UsgNote"
72 (NoteTy (SynNote sty) ty) = do { sty' <- genAnnotTyN getAnnM sty
73 -- is this right? shouldn't there be some
74 -- correlation between sty' and ty'?
75 -- But sty is a TyConApp; does this make it safer?
76 ; ty' <- genAnnotTyN getAnnM ty
77 ; return (NoteTy (SynNote sty') ty')
80 (NoteTy fvn@(FTVNote _) ty) = do { ty' <- genAnnotTyN getAnnM ty
81 ; return (NoteTy fvn ty')
85 ty0@(TyVarTy _) = do { return ty0 }
88 (AppTy ty1 ty2) = do { ty1' <- genAnnotTyN getAnnM ty1
89 ; ty2' <- genAnnotTyN getAnnM ty2
90 ; return (AppTy ty1' ty2')
94 (TyConApp tc tys) = ASSERT( isFunTyCon tc || isAlgTyCon tc || isPrimTyCon tc || isSynTyCon tc )
95 do { let gAT = if isFunTyCon tc
96 then genAnnotTy -- sigma for partial apps of (->)
97 else genAnnotTyN -- tau otherwise
98 ; tys' <- mapM (gAT getAnnM) tys
99 ; return (TyConApp tc tys')
103 (FunTy ty1 ty2) = do { ty1' <- genAnnotTy getAnnM ty1
104 ; ty2' <- genAnnotTy getAnnM ty2
105 ; return (FunTy ty1' ty2')
109 (ForAllTy v ty) = do { ty' <- genAnnotTyN getAnnM ty
110 ; return (ForAllTy v ty')
116 Walking over (and retyping) terms
117 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
119 We also often need to play with the types in a term. This is slightly
120 tricky because of redundancy: we want to change binder types, and keep
121 the bound types matching these; then there's a special case also with
122 non-locally-defined bound variables. We generalise over all this
125 The name `annot' is a bit of a misnomer, as the code is parameterised
126 over exactly what it does to the types (and certain terms). Notice
127 also that it is possible for this parameter to use
128 monadically-threaded state: here called `flexi'. For genuine
129 annotation, this state will be a UniqSupply.
131 We may add annotations to the outside of a (term, not type) lambda; a
132 function passed to @genAnnotBinds@ does this, taking the lambda and
133 returning the annotated lambda. It is inside the @AnnotM@ monad.
134 This term-munging function is applied when we see either a term lambda
135 or a usage annotation; *IMPORTANT:* it is applied *before* we recurse
136 down into the term, and it is expected to work only at the top level.
137 Recursion will subsequently be done by genAnnotBinds. It may
138 optionally remove a Note TermUsg, or optionally add one if it is not
139 already present, but it may perform NO OTHER MODIFICATIONS to the
140 structure of the term.
142 We do different things to types of variables bound locally and of
143 variables bound in other modules, in certain cases: the former get
144 uvars and the latter keep their existing annotations when we annotate,
145 for example. To control this, @MungeFlags@ describes what kind of a
146 type this is that we're about to munge.
149 data MungeFlags = MungeFlags { isSigma :: Bool, -- want annotated on top (sigma type)
150 isLocal :: Bool, -- is locally-defined type
151 hasUsg :: Bool, -- has fixed usage info, don't touch
152 isExp :: Bool, -- is exported (and must be pessimised)
153 mfLoc :: SDoc -- location info
156 tauTyMF loc = MungeFlags { isSigma = False, isLocal = True,
157 hasUsg = False, isExp = False, mfLoc = loc }
158 sigVarTyMF v = MungeFlags { isSigma = True, isLocal = hasLocalDef v,
159 hasUsg = hasUsgInfo v, isExp = isExported v,
160 mfLoc = ptext SLIT("type of binder") <+> ppr v }
163 The helper functions @tauTyMF@ and @sigVarTyMF@ create @MungeFlags@
164 for us. @sigVarTyMF@ checks the variable to see how to set the flags.
166 @hasLocalDef@ tells us if the given variable has an actual local
167 definition that we can play with. This is not quite the same as
168 @isLocallyDefined@, since @IMustBeINLINEd@ things (usually) don't have
169 a local definition - the simplifier will inline whatever their
170 unfolding is anyway. We treat these as if they were externally
171 defined, since we don't have access to their definition (at least not
172 easily). This doesn't hurt much, since after the simplifier has run
173 the unfolding will have been inlined and we can access the unfolding
176 @hasUsgInfo@, on the other hand, says if the variable already has
177 usage info in its type that must at all costs be preserved. This is
178 assumed true (exactly) of all imported ids.
181 hasLocalDef :: IdOrTyVar -> Bool
182 hasLocalDef var = isLocallyDefined var
183 && not (idMustBeINLINEd var)
185 hasUsgInfo :: IdOrTyVar -> Bool
186 hasUsgInfo var = (not . isLocallyDefined) var
189 Here's the walk itself.
192 genAnnotBinds :: (MungeFlags -> Type -> AnnotM flexi Type)
193 -> (CoreExpr -> AnnotM flexi CoreExpr) -- see caveats above
195 -> AnnotM flexi [CoreBind]
197 genAnnotBinds _ _ [] = return []
199 genAnnotBinds f g (b:bs) = do { (b',vs,vs') <- genAnnotBind f g b
200 ; bs' <- withAnnVars vs vs' $
205 genAnnotBind :: (MungeFlags -> Type -> AnnotM flexi Type) -- type-altering function
206 -> (CoreExpr -> AnnotM flexi CoreExpr) -- term-altering function
207 -> CoreBind -- original CoreBind
209 (CoreBind, -- annotated CoreBind
210 [IdOrTyVar], -- old variables, to be mapped to...
211 [IdOrTyVar]) -- ... new variables
213 genAnnotBind f g (NonRec v1 e1) = do { v1' <- genAnnotVar f v1
214 ; e1' <- genAnnotCE f g e1
215 ; return (NonRec v1' e1', [v1], [v1'])
218 genAnnotBind f g (Rec ves) = do { let (vs,es) = unzip ves
219 ; vs' <- mapM (genAnnotVar f) vs
220 ; es' <- withAnnVars vs vs' $
221 mapM (genAnnotCE f g) es
222 ; return (Rec (zip vs' es'), vs, vs')
225 genAnnotCE :: (MungeFlags -> Type -> AnnotM flexi Type) -- type-altering function
226 -> (CoreExpr -> AnnotM flexi CoreExpr) -- term-altering function
227 -> CoreExpr -- original expression
228 -> AnnotM flexi CoreExpr -- yields new expression
230 genAnnotCE mungeType mungeTerm = go
231 where go e0@(Var v) | isTyVar v = return e0 -- arises, e.g., as tyargs of Con
232 -- (no it doesn't: (Type (TyVar tyvar))
233 | otherwise = do { mv' <- lookupAnnVar v
235 Just var -> return var
236 Nothing -> fixedVar v
240 go (Con c args) = -- we know it's saturated
241 do { args' <- mapM go args
242 ; return (Con c args')
245 go (App e arg) = do { e' <- go e
247 ; return (App e' arg')
250 go e0@(Lam v0 _) = do { e1 <- (if isTyVar v0 then return else mungeTerm) e0
252 = case e1 of -- munge may have added note
253 Note tu@(TermUsg _) (Lam v e2)
255 Lam v e2 -> (v,e2,id)
256 ; v' <- genAnnotVar mungeType v
257 ; e' <- withAnnVar v v' $ go e2
258 ; return (wrap (Lam v' e'))
261 go (Let bind e) = do { (bind',vs,vs') <- genAnnotBind mungeType mungeTerm bind
262 ; e' <- withAnnVars vs vs' $ go e
263 ; return (Let bind' e')
266 go (Case e v alts) = do { e' <- go e
267 ; v' <- genAnnotVar mungeType v
268 ; alts' <- withAnnVar v v' $ mapM genAnnotAlt alts
269 ; return (Case e' v' alts')
272 go (Note scc@(SCC _) e) = do { e' <- go e
273 ; return (Note scc e')
275 go e0@(Note (Coerce ty1 ty0)
276 e) = do { ty1' <- mungeType
277 (tauTyMF (ptext SLIT("coercer of")
280 (tauTyMF (ptext SLIT("coercee of")
282 -- (Better to specify ty0'
283 -- identical to the type of e, including
284 -- annotations, right at the beginning, but
285 -- not possible at this point.)
287 ; return (Note (Coerce ty1' ty0') e')
289 go (Note InlineCall e) = do { e' <- go e
290 ; return (Note InlineCall e')
292 go e0@(Note (TermUsg _) _) = do { e1 <- mungeTerm e0
293 ; case e1 of -- munge may have removed note
294 Note tu@(TermUsg _) e2 -> do { e3 <- go e2
295 ; return (Note tu e3)
300 go e0@(Type ty) = -- should only occur at toplevel of Arg,
302 do { ty' <- mungeType
303 (tauTyMF (ptext SLIT("tyarg")
308 fixedVar v = ASSERT2( not (hasLocalDef v), text "genAnnotCE: locally defined var" <+> ppr v <+> text "not in varenv" )
309 genAnnotVar mungeType v
311 genAnnotAlt (c,vs,e) = do { vs' <- mapM (genAnnotVar mungeType) vs
312 ; e' <- withAnnVars vs vs' $ go e
313 ; return (c, vs', e')
317 genAnnotVar :: (MungeFlags -> Type -> AnnotM flexi Type)
319 -> AnnotM flexi IdOrTyVar
321 genAnnotVar mungeType v | isTyVar v = return v
322 | otherwise = do { vty' <- mungeType (sigVarTyMF v) (varType v)
323 ; return (setVarType v vty')
327 pprTrace "genAnnotVar" (ppr (tyUsg vty') <+> ppr v) $
333 ======================================================================
335 Some specific things to do to types inside terms
336 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
338 @annotTyM@ annotates a type with fresh uvars everywhere the inference
339 is allowed to go, and leaves alone annotations where it may not go.
341 We assume there are no annotations already.
344 annotTyM :: MungeFlags -> Type -> AnnotM UniqSupply Type
346 annotTyM mf ty = uniqSMtoAnnotM . uniqSMMToUs $
347 case (hasUsg mf, isLocal mf, isSigma mf) of
348 (True ,_ ,_ ) -> ASSERT( isUsgTy ty )
350 (False,True ,True ) -> if isExp mf then
351 annotTyP (tag 'p') ty
354 (False,True ,False) -> annotTyN (tag 't') ty
355 (False,False,True ) -> return $ annotMany ty -- assume worst
356 (False,False,False) -> return $ annotManyN ty
357 where tag c = Right $ "annotTyM:" ++ [c] ++ ": " ++ showSDoc (ppr ty)
359 -- specific functions for annotating tau and sigma types
362 annotTy tag = genAnnotTy (newVarUSMM tag)
363 annotTyN tag = genAnnotTyN (newVarUSMM tag)
365 -- ...with uvars and pessimal Manys (for exported ids)
366 annotTyP tag ty = do { ty' <- annotTy tag ty ; return (pessimise True ty') }
369 annotMany, annotManyN :: Type -> Type
374 annotMany ty = unId (genAnnotTy (return UsMany) ty)
375 annotManyN ty = unId (genAnnotTyN (return UsMany) ty)
378 -- monad required for the above
379 newtype Id a = Id a ; unId (Id a) = a
380 instance Monad Id where { a >>= f = f (unId a) ; return a = Id a }
382 -- lambda-annotating function for use along with the above
383 annotLam e0@(Lam v e) = do { uv <- uniqSMtoAnnotM $ newVarUs (Left e0)
384 ; return (Note (TermUsg uv) (Lam v e))
386 annotLam (Note (TermUsg _) _) = panic "annotLam: unexpected term usage annot"
389 The above requires a `pessimising' translation. This is applied to
390 types of exported ids, and ensures that they have a fully general
391 type (since we don't know how they will be used in other modules).
394 pessimise :: Bool -> Type -> Type
397 pessimise co ty0@(NoteTy usg@(UsgNote u ) ty)
399 then case u of UsMany -> pty
400 UsVar _ -> pty -- force to UsMany
401 UsOnce -> pprPanic "pessimise:" (ppr ty0)
403 where pty = pessimiseN co ty
405 pessimise co ty0 = pessimiseN co ty0 -- assume UsMany
407 pessimise co ty0@(NoteTy usg@(UsgNote u ) ty)
409 then case u of UsMany -> NoteTy usg pty
410 UsVar _ -> NoteTy (UsgNote UsMany) pty
411 UsOnce -> pprPanic "pessimise:" (ppr ty0)
413 where pty = pessimiseN co ty
415 pessimise co ty0 = pprPanic "pessimise: missing usage note:" $
419 pessimiseN co ty0@(NoteTy usg@(UsgNote _ ) ty) = pprPanic "pessimiseN: unexpected usage note:" $
421 pessimiseN co (NoteTy (SynNote sty) ty) = NoteTy (SynNote (pessimiseN co sty))
423 pessimiseN co (NoteTy note@(FTVNote _ ) ty) = NoteTy note (pessimiseN co ty)
424 pessimiseN co ty0@(TyVarTy _) = ty0
425 pessimiseN co ty0@(AppTy _ _) = ty0
426 pessimiseN co ty0@(TyConApp tc tys) = ASSERT( not ((isFunTyCon tc) && (length tys > 1)) )
428 pessimiseN co (FunTy ty1 ty2) = FunTy (pessimise (not co) ty1)
430 pessimiseN co (ForAllTy tyv ty) = ForAllTy tyv (pessimiseN co ty)
434 @unAnnotTyM@ strips annotations (that the inference is allowed to
435 touch) from a term, and `fixes' those it isn't permitted to touch (by
436 putting @Many@ annotations where they are missing, but leaving
437 existing annotations in the type).
439 @unTermUsg@ removes from a term any term usage annotations it finds.
442 unAnnotTyM :: MungeFlags -> Type -> AnnotM a Type
444 unAnnotTyM mf ty = if hasUsg mf then
446 return (fixAnnotTy ty)
447 else return (unannotTy ty)
450 unTermUsg :: CoreExpr -> AnnotM a CoreExpr
451 -- strip all term annotations
452 unTermUsg e@(Lam _ _) = return e
453 unTermUsg (Note (TermUsg _) e) = return e
454 unTermUsg _ = panic "unTermUsg"
456 unannotTy :: Type -> Type
457 -- strip all annotations
458 unannotTy (NoteTy (UsgNote _ ) ty) = unannotTy ty
459 unannotTy (NoteTy (SynNote sty) ty) = NoteTy (SynNote (unannotTy sty)) (unannotTy ty)
460 unannotTy (NoteTy note@(FTVNote _ ) ty) = NoteTy note (unannotTy ty)
461 unannotTy ty@(TyVarTy _) = ty
462 unannotTy (AppTy ty1 ty2) = AppTy (unannotTy ty1) (unannotTy ty2)
463 unannotTy (TyConApp tc tys) = TyConApp tc (map unannotTy tys)
464 unannotTy (FunTy ty1 ty2) = FunTy (unannotTy ty1) (unannotTy ty2)
465 unannotTy (ForAllTy tyv ty) = ForAllTy tyv (unannotTy ty)
468 fixAnnotTy :: Type -> Type
469 -- put Manys where they are missing
473 fixAnnotTy (NoteTy note@(UsgNote _ ) ty) = NoteTy note (fixAnnotTyN ty)
474 fixAnnotTy ty0 = NoteTy (UsgNote UsMany) (fixAnnotTyN ty0)
476 fixAnnotTyN ty0@(NoteTy note@(UsgNote _ ) ty) = pprPanic "fixAnnotTyN: unexpected usage note:" $
478 fixAnnotTyN (NoteTy (SynNote sty) ty) = NoteTy (SynNote (fixAnnotTyN sty))
480 fixAnnotTyN (NoteTy note@(FTVNote _ ) ty) = NoteTy note (fixAnnotTyN ty)
481 fixAnnotTyN ty0@(TyVarTy _) = ty0
482 fixAnnotTyN (AppTy ty1 ty2) = AppTy (fixAnnotTyN ty1) (fixAnnotTyN ty2)
483 fixAnnotTyN (TyConApp tc tys) = ASSERT( isFunTyCon tc || isAlgTyCon tc || isPrimTyCon tc || isSynTyCon tc )
484 TyConApp tc (map (if isFunTyCon tc then
488 fixAnnotTyN (FunTy ty1 ty2) = FunTy (fixAnnotTy ty1) (fixAnnotTy ty2)
489 fixAnnotTyN (ForAllTy tyv ty) = ForAllTy tyv (fixAnnotTyN ty)
493 The composition (reannotating a type with fresh uvars but the same
494 structure) is useful elsewhere:
497 freshannotTy :: Type -> UniqSMM Type
498 freshannotTy = annotTy (Right "freshannotTy") . unannotTy
502 Wrappers apply these functions to sets of bindings.
505 doAnnotBinds :: UniqSupply
507 -> ([CoreBind],UniqSupply)
509 doAnnotBinds us binds = initAnnotM us (genAnnotBinds annotTyM annotLam binds)
512 doUnAnnotBinds :: [CoreBind]
515 doUnAnnotBinds binds = fst $ initAnnotM () $
516 genAnnotBinds unAnnotTyM unTermUsg binds
519 ======================================================================
524 The @UniqSM@ type is not an instance of @Monad@, and cannot be made so
525 since it is merely a synonym rather than a newtype. Here we define
526 @UniqSMM@, which *is* an instance of @Monad@.
529 newtype UniqSMM a = UsToUniqSMM (UniqSM a)
530 uniqSMMToUs (UsToUniqSMM us) = us
531 usToUniqSMM = UsToUniqSMM
533 instance Monad UniqSMM where
534 m >>= f = UsToUniqSMM $ uniqSMMToUs m `thenUs` \ a ->
536 return = UsToUniqSMM . returnUs
540 For annotation, the monad @AnnotM@, we need to carry around our
541 variable mapping, along with some general state.
544 newtype AnnotM flexi a = AnnotM ( flexi -- UniqSupply etc
545 -> VarEnv IdOrTyVar -- unannotated to annotated variables
546 -> (a,flexi,VarEnv IdOrTyVar))
547 unAnnotM (AnnotM f) = f
549 instance Monad (AnnotM flexi) where
550 a >>= f = AnnotM (\ us ve -> let (r,us',ve') = unAnnotM a us ve
551 in unAnnotM (f r) us' ve')
552 return a = AnnotM (\ us ve -> (a,us,ve))
554 initAnnotM :: fl -> AnnotM fl a -> (a,fl)
555 initAnnotM fl m = case (unAnnotM m) fl emptyVarEnv of { (r,fl',_) -> (r,fl') }
557 withAnnVar :: IdOrTyVar -> IdOrTyVar -> AnnotM fl a -> AnnotM fl a
558 withAnnVar v v' m = AnnotM (\ us ve -> let ve' = extendVarEnv ve v v'
559 (r,us',_) = (unAnnotM m) us ve'
562 withAnnVars :: [IdOrTyVar] -> [IdOrTyVar] -> AnnotM fl a -> AnnotM fl a
563 withAnnVars vs vs' m = AnnotM (\ us ve -> let ve' = plusVarEnv ve (zipVarEnv vs vs')
564 (r,us',_) = (unAnnotM m) us ve'
567 lookupAnnVar :: IdOrTyVar -> AnnotM fl (Maybe IdOrTyVar)
568 lookupAnnVar var = AnnotM (\ us ve -> (lookupVarEnv ve var,
573 A useful helper allows us to turn a computation in the unique supply
574 monad into one in the annotation monad parameterised by a unique
578 uniqSMtoAnnotM :: UniqSM a -> AnnotM UniqSupply a
580 uniqSMtoAnnotM m = AnnotM (\ us ve -> let (r,us') = initUs us m
584 @newVarUs@ and @newVarUSMM@ generate a new usage variable. They take
585 an argument which is used for debugging only, describing what the
586 variable is to annotate.
589 newVarUs :: (Either CoreExpr String) -> UniqSM UsageAnn
590 -- the first arg is for debugging use only
591 newVarUs e = getUniqueUs `thenUs` \ u ->
596 Left (Con (Literal _) _) -> "literal"
597 Left (Con _ _) -> "primop"
598 Left (Lam v e) -> "lambda: " ++ showSDoc (ppr v)
601 in pprTrace "newVarUs:" (ppr uv <+> text src) $
605 newVarUSMM :: (Either CoreExpr String) -> UniqSMM UsageAnn
606 newVarUSMM = usToUniqSMM . newVarUs
609 ======================================================================
611 PrimOps and usage information.
613 Analagously to @DataCon.dataConArgTys@, we determine the argtys and
614 result ty of a primop, *after* substition (which may reveal more args,
615 notably for @CCall@s).
618 primOpUsgTys :: PrimOp -- this primop
619 -> [Type] -- instantiated at these (tau) types
620 -> ([Type],Type) -- requires args of these (sigma) types,
621 -- and returns this (sigma) type
623 primOpUsgTys p tys = let (tyvs,ty0us,rtyu) = primOpUsg p
624 s = zipVarEnv tyvs tys
625 (ty1us,rty1u) = splitFunTys (substTy s rtyu)
626 -- substitution may reveal more args
627 in ((map (substTy s) ty0us) ++ ty1us,
631 ======================================================================