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-07-07
12 module UsageSPUtils ( AnnotM(AnnotM), initAnnotM,
14 MungeFlags(isSigma,isLocal,isExp,hasUsg,mfLoc),
16 doAnnotBinds, doUnAnnotBinds,
17 annotTy, annotTyN, 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 ( mayHaveNoBinding, 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 )
37 import PrimOp ( PrimOp, primOpUsg )
38 import Maybes ( expectJust )
39 import UniqSupply ( UniqSupply, UniqSM, initUs, getUniqueUs, thenUs, returnUs )
41 import PprCore ( ) -- instances only
44 ======================================================================
46 Walking over (and altering) types
47 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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.
56 genAnnotTy :: Monad m =>
57 (m UsageAnn) -- get new annotation
61 genAnnotTy getAnnM ty = do { u <- getAnnM
62 ; ty' <- genAnnotTyN getAnnM ty
63 ; return (NoteTy (UsgNote u) ty')
66 genAnnotTyN :: Monad m =>
72 (NoteTy (UsgNote _) ty) = panic "genAnnotTyN: unexpected UsgNote"
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')
82 (NoteTy fvn@(FTVNote _) ty) = do { ty' <- genAnnotTyN getAnnM ty
83 ; return (NoteTy fvn ty')
87 ty0@(TyVarTy _) = do { return ty0 }
90 (AppTy ty1 ty2) = do { ty1' <- genAnnotTyN getAnnM ty1
91 ; ty2' <- genAnnotTyN getAnnM ty2
92 ; return (AppTy ty1' ty2')
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')
105 (FunTy ty1 ty2) = do { ty1' <- genAnnotTy getAnnM ty1
106 ; ty2' <- genAnnotTy getAnnM ty2
107 ; return (FunTy ty1' ty2')
111 (ForAllTy v ty) = do { ty' <- genAnnotTyN getAnnM ty
112 ; return (ForAllTy v ty')
118 Walking over (and retyping) terms
119 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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
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.
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.
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.
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
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 }
165 The helper functions @tauTyMF@ and @sigVarTyMF@ create @MungeFlags@
166 for us. @sigVarTyMF@ checks the variable to see how to set the flags.
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 @mayHaveNoBindingId@ 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
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.
183 hasLocalDef :: IdOrTyVar -> Bool
184 hasLocalDef var = isLocallyDefined var
185 && not (mayHaveNoBinding var)
187 hasUsgInfo :: IdOrTyVar -> Bool
188 hasUsgInfo var = (not . isLocallyDefined) var
191 Here's the walk itself.
194 genAnnotBinds :: (MungeFlags -> Type -> AnnotM flexi Type)
195 -> (CoreExpr -> AnnotM flexi CoreExpr) -- see caveats above
197 -> AnnotM flexi [CoreBind]
199 genAnnotBinds _ _ [] = return []
201 genAnnotBinds f g (b:bs) = do { (b',vs,vs') <- genAnnotBind f g b
202 ; bs' <- withAnnVars vs vs' $
207 genAnnotBind :: (MungeFlags -> Type -> AnnotM flexi Type) -- type-altering function
208 -> (CoreExpr -> AnnotM flexi CoreExpr) -- term-altering function
209 -> CoreBind -- original CoreBind
211 (CoreBind, -- annotated CoreBind
212 [IdOrTyVar], -- old variables, to be mapped to...
213 [IdOrTyVar]) -- ... new variables
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'])
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')
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
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
237 Just var -> return var
238 Nothing -> fixedVar v
242 go (Con c args) = -- we know it's saturated
243 do { args' <- mapM go args
244 ; return (Con c args')
247 go (App e arg) = do { e' <- go e
249 ; return (App e' arg')
252 go e0@(Lam v0 _) = do { e1 <- (if isTyVar v0 then return else mungeTerm) e0
254 = case e1 of -- munge may have added note
255 Note tu@(TermUsg _) (Lam v e2)
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'))
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')
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')
274 go (Note scc@(SCC _) e) = do { e' <- go e
275 ; return (Note scc e')
277 go e0@(Note (Coerce ty1 ty0)
278 e) = do { ty1' <- mungeType
279 (tauTyMF (ptext SLIT("coercer of")
282 (tauTyMF (ptext SLIT("coercee of")
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.)
289 ; return (Note (Coerce ty1' ty0') e')
291 go (Note InlineCall e) = do { e' <- go e
292 ; return (Note InlineCall e')
294 go (Note InlineMe e) = do { e' <- go e
295 ; return (Note InlineMe e')
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)
305 go e0@(Type ty) = -- should only occur at toplevel of Arg,
307 do { ty' <- mungeType
308 (tauTyMF (ptext SLIT("tyarg")
313 fixedVar v = ASSERT2( not (hasLocalDef v), text "genAnnotCE: locally defined var" <+> ppr v <+> text "not in varenv" )
314 genAnnotVar mungeType v
316 genAnnotAlt (c,vs,e) = do { vs' <- mapM (genAnnotVar mungeType) vs
317 ; e' <- withAnnVars vs vs' $ go e
318 ; return (c, vs', e')
322 genAnnotVar :: (MungeFlags -> Type -> AnnotM flexi Type)
324 -> AnnotM flexi IdOrTyVar
326 genAnnotVar mungeType v | isTyVar v = return v
327 | otherwise = do { vty' <- mungeType (sigVarTyMF v) (varType v)
328 ; return (setVarType v vty')
332 pprTrace "genAnnotVar" (ppr (tyUsg vty') <+> ppr v) $
338 ======================================================================
340 Some specific things to do to types inside terms
341 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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.
346 We assume there are no annotations already.
349 annotTyM :: MungeFlags -> Type -> AnnotM UniqSupply Type
351 annotTyM mf ty = uniqSMtoAnnotM . uniqSMMToUs $
352 case (hasUsg mf, isLocal mf, isSigma mf) of
353 (True ,_ ,_ ) -> ASSERT( isUsgTy ty )
355 (False,True ,True ) -> if isExp mf then
356 annotTyP (tag 'p') 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)
364 -- specific functions for annotating tau and sigma types
367 annotTy tag = genAnnotTy (newVarUSMM tag)
368 annotTyN tag = genAnnotTyN (newVarUSMM tag)
370 -- ...with uvars and pessimal Manys (for exported ids)
371 annotTyP tag ty = do { ty' <- annotTy tag ty ; return (pessimise True ty') }
374 annotMany, annotManyN :: Type -> Type
379 annotMany ty = unId (genAnnotTy (return UsMany) ty)
380 annotManyN ty = unId (genAnnotTyN (return UsMany) ty)
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 }
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))
391 annotLam (Note (TermUsg _) _) = panic "annotLam: unexpected term usage annot"
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).
399 pessimise :: Bool -> Type -> Type
402 pessimise co ty0@(NoteTy usg@(UsgNote u ) ty)
404 then case u of UsMany -> pty
405 UsVar _ -> pty -- force to UsMany
406 UsOnce -> pprPanic "pessimise:" (ppr ty0)
408 where pty = pessimiseN co ty
410 pessimise co ty0 = pessimiseN co ty0 -- assume UsMany
412 pessimise co ty0@(NoteTy usg@(UsgNote u ) ty)
414 then case u of UsMany -> NoteTy usg pty
415 UsVar _ -> NoteTy (UsgNote UsMany) pty
416 UsOnce -> pprPanic "pessimise:" (ppr ty0)
418 where pty = pessimiseN co ty
420 pessimise co ty0 = pprPanic "pessimise: missing usage note:" $
424 pessimiseN co ty0@(NoteTy usg@(UsgNote _ ) ty) = pprPanic "pessimiseN: unexpected usage note:" $
426 pessimiseN co (NoteTy (SynNote sty) ty) = NoteTy (SynNote (pessimiseN co sty))
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)) )
433 pessimiseN co (FunTy ty1 ty2) = FunTy (pessimise (not co) ty1)
435 pessimiseN co (ForAllTy tyv ty) = ForAllTy tyv (pessimiseN co ty)
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).
444 @unTermUsg@ removes from a term any term usage annotations it finds.
447 unAnnotTyM :: MungeFlags -> Type -> AnnotM a Type
449 unAnnotTyM mf ty = if hasUsg mf then
451 return (fixAnnotTy ty)
452 else return (unannotTy ty)
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"
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 -- IP notes need to be preserved
468 unannotTy ty@(NoteTy (IPNote _) _) = ty
469 unannotTy ty@(TyVarTy _) = ty
470 unannotTy (AppTy ty1 ty2) = AppTy (unannotTy ty1) (unannotTy ty2)
471 unannotTy (TyConApp tc tys) = TyConApp tc (map unannotTy tys)
472 unannotTy (FunTy ty1 ty2) = FunTy (unannotTy ty1) (unannotTy ty2)
473 unannotTy (ForAllTy tyv ty) = ForAllTy tyv (unannotTy ty)
476 fixAnnotTy :: Type -> Type
477 -- put Manys where they are missing
481 fixAnnotTy (NoteTy note@(UsgForAll uv) ty) = NoteTy note (fixAnnotTy ty)
482 fixAnnotTy (NoteTy note@(UsgNote _ ) ty) = NoteTy note (fixAnnotTyN ty)
483 fixAnnotTy ty0 = NoteTy (UsgNote UsMany) (fixAnnotTyN ty0)
485 fixAnnotTyN ty0@(NoteTy note@(UsgNote _ ) ty) = pprPanic "fixAnnotTyN: unexpected usage note:" $
487 fixAnnotTyN (NoteTy (SynNote sty) ty) = NoteTy (SynNote (fixAnnotTyN sty))
489 fixAnnotTyN (NoteTy note@(FTVNote _ ) ty) = NoteTy note (fixAnnotTyN ty)
490 fixAnnotTyN ty0@(TyVarTy _) = ty0
491 fixAnnotTyN (AppTy ty1 ty2) = AppTy (fixAnnotTyN ty1) (fixAnnotTyN ty2)
492 fixAnnotTyN (TyConApp tc tys) = ASSERT( isFunTyCon tc || isAlgTyCon tc || isPrimTyCon tc || isSynTyCon tc )
493 TyConApp tc (map (if isFunTyCon tc then
497 fixAnnotTyN (FunTy ty1 ty2) = FunTy (fixAnnotTy ty1) (fixAnnotTy ty2)
498 fixAnnotTyN (ForAllTy tyv ty) = ForAllTy tyv (fixAnnotTyN ty)
502 The composition (reannotating a type with fresh uvars but the same
503 structure) is useful elsewhere:
506 freshannotTy :: Type -> UniqSMM Type
507 freshannotTy = annotTy (Right "freshannotTy") . unannotTy
511 Wrappers apply these functions to sets of bindings.
514 doAnnotBinds :: UniqSupply
516 -> ([CoreBind],UniqSupply)
518 doAnnotBinds us binds = initAnnotM us (genAnnotBinds annotTyM annotLam binds)
521 doUnAnnotBinds :: [CoreBind]
524 doUnAnnotBinds binds = fst $ initAnnotM () $
525 genAnnotBinds unAnnotTyM unTermUsg binds
528 ======================================================================
533 The @UniqSM@ type is not an instance of @Monad@, and cannot be made so
534 since it is merely a synonym rather than a newtype. Here we define
535 @UniqSMM@, which *is* an instance of @Monad@.
538 newtype UniqSMM a = UsToUniqSMM (UniqSM a)
539 uniqSMMToUs (UsToUniqSMM us) = us
540 usToUniqSMM = UsToUniqSMM
542 instance Monad UniqSMM where
543 m >>= f = UsToUniqSMM $ uniqSMMToUs m `thenUs` \ a ->
545 return = UsToUniqSMM . returnUs
549 For annotation, the monad @AnnotM@, we need to carry around our
550 variable mapping, along with some general state.
553 newtype AnnotM flexi a = AnnotM ( flexi -- UniqSupply etc
554 -> VarEnv IdOrTyVar -- unannotated to annotated variables
555 -> (a,flexi,VarEnv IdOrTyVar))
556 unAnnotM (AnnotM f) = f
558 instance Monad (AnnotM flexi) where
559 a >>= f = AnnotM (\ us ve -> let (r,us',ve') = unAnnotM a us ve
560 in unAnnotM (f r) us' ve')
561 return a = AnnotM (\ us ve -> (a,us,ve))
563 initAnnotM :: fl -> AnnotM fl a -> (a,fl)
564 initAnnotM fl m = case (unAnnotM m) fl emptyVarEnv of { (r,fl',_) -> (r,fl') }
566 withAnnVar :: IdOrTyVar -> IdOrTyVar -> AnnotM fl a -> AnnotM fl a
567 withAnnVar v v' m = AnnotM (\ us ve -> let ve' = extendVarEnv ve v v'
568 (r,us',_) = (unAnnotM m) us ve'
571 withAnnVars :: [IdOrTyVar] -> [IdOrTyVar] -> AnnotM fl a -> AnnotM fl a
572 withAnnVars vs vs' m = AnnotM (\ us ve -> let ve' = plusVarEnv ve (zipVarEnv vs vs')
573 (r,us',_) = (unAnnotM m) us ve'
576 lookupAnnVar :: IdOrTyVar -> AnnotM fl (Maybe IdOrTyVar)
577 lookupAnnVar var = AnnotM (\ us ve -> (lookupVarEnv ve var,
582 A useful helper allows us to turn a computation in the unique supply
583 monad into one in the annotation monad parameterised by a unique
587 uniqSMtoAnnotM :: UniqSM a -> AnnotM UniqSupply a
589 uniqSMtoAnnotM m = AnnotM (\ us ve -> let (r,us') = initUs us m
593 @newVarUs@ and @newVarUSMM@ generate a new usage variable. They take
594 an argument which is used for debugging only, describing what the
595 variable is to annotate.
598 newVarUs :: (Either CoreExpr String) -> UniqSM UsageAnn
599 -- the first arg is for debugging use only
600 newVarUs e = getUniqueUs `thenUs` \ u ->
605 Left (Con (Literal _) _) -> "literal"
606 Left (Con _ _) -> "primop"
607 Left (Lam v e) -> "lambda: " ++ showSDoc (ppr v)
610 in pprTrace "newVarUs:" (ppr uv <+> text src) $
614 newVarUSMM :: (Either CoreExpr String) -> UniqSMM UsageAnn
615 newVarUSMM = usToUniqSMM . newVarUs
618 ======================================================================
620 PrimOps and usage information.
622 Analagously to @DataCon.dataConArgTys@, we determine the argtys and
623 result ty of a primop, *after* substition (which may reveal more args,
624 notably for @CCall@s).
627 primOpUsgTys :: PrimOp -- this primop
628 -> [Type] -- instantiated at these (tau) types
629 -> ([Type],Type) -- requires args of these (sigma) types,
630 -- and returns this (sigma) type
632 primOpUsgTys p tys = let (tyvs,ty0us,rtyu) = primOpUsg p
633 s = mkTyVarSubst tyvs tys
634 (ty1us,rty1u) = splitFunTys (substTy s rtyu)
635 -- substitution may reveal more args
636 in ((map (substTy s) ty0us) ++ ty1us,
640 ======================================================================