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 Literal ( Literal(..) )
29 import Var ( Var, 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 PprType ( {- instance Outputable Type -} )
35 import Subst ( substTy, mkTyVarSubst )
36 import TyCon ( isAlgTyCon, isPrimTyCon, isSynTyCon, isFunTyCon )
38 import PrimOp ( PrimOp, primOpUsg )
39 import Maybes ( expectJust )
40 import UniqSupply ( UniqSupply, UniqSM, initUs, getUniqueUs, thenUs, returnUs )
42 import PprCore ( ) -- instances only
45 ======================================================================
47 Walking over (and altering) types
48 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
50 We often need to fiddle with (i.e., add or remove) usage annotations
51 on a type. We define here a general framework to do this. Usage
52 annotations come from any monad with a function @getAnnM@ which yields
53 a new annotation. We use two mutually recursive functions, one for
54 sigma types and one for tau types.
57 genAnnotTy :: Monad m =>
58 (m UsageAnn) -- get new annotation
62 genAnnotTy getAnnM ty = do { u <- getAnnM
63 ; ty' <- genAnnotTyN getAnnM ty
64 ; return (NoteTy (UsgNote u) ty')
67 genAnnotTyN :: Monad m =>
73 (NoteTy (UsgNote _) ty) = panic "genAnnotTyN: unexpected UsgNote"
75 (NoteTy (SynNote sty) ty) = do { sty' <- genAnnotTyN getAnnM sty
76 -- is this right? shouldn't there be some
77 -- correlation between sty' and ty'?
78 -- But sty is a TyConApp; does this make it safer?
79 ; ty' <- genAnnotTyN getAnnM ty
80 ; return (NoteTy (SynNote sty') ty')
83 (NoteTy fvn@(FTVNote _) ty) = do { ty' <- genAnnotTyN getAnnM ty
84 ; return (NoteTy fvn ty')
88 ty0@(TyVarTy _) = do { return ty0 }
91 (AppTy ty1 ty2) = do { ty1' <- genAnnotTyN getAnnM ty1
92 ; ty2' <- genAnnotTyN getAnnM ty2
93 ; return (AppTy ty1' ty2')
97 (TyConApp tc tys) = ASSERT( isFunTyCon tc || isAlgTyCon tc || isPrimTyCon tc || isSynTyCon tc )
98 do { let gAT = if isFunTyCon tc
99 then genAnnotTy -- sigma for partial apps of (->)
100 else genAnnotTyN -- tau otherwise
101 ; tys' <- mapM (gAT getAnnM) tys
102 ; return (TyConApp tc tys')
106 (FunTy ty1 ty2) = do { ty1' <- genAnnotTy getAnnM ty1
107 ; ty2' <- genAnnotTy getAnnM ty2
108 ; return (FunTy ty1' ty2')
112 (ForAllTy v ty) = do { ty' <- genAnnotTyN getAnnM ty
113 ; return (ForAllTy v ty')
119 Walking over (and retyping) terms
120 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
122 We also often need to play with the types in a term. This is slightly
123 tricky because of redundancy: we want to change binder types, and keep
124 the bound types matching these; then there's a special case also with
125 non-locally-defined bound variables. We generalise over all this
128 The name `annot' is a bit of a misnomer, as the code is parameterised
129 over exactly what it does to the types (and certain terms). Notice
130 also that it is possible for this parameter to use
131 monadically-threaded state: here called `flexi'. For genuine
132 annotation, this state will be a UniqSupply.
134 We may add annotations to the outside of a (term, not type) lambda; a
135 function passed to @genAnnotBinds@ does this, taking the lambda and
136 returning the annotated lambda. It is inside the @AnnotM@ monad.
137 This term-munging function is applied when we see either a term lambda
138 or a usage annotation; *IMPORTANT:* it is applied *before* we recurse
139 down into the term, and it is expected to work only at the top level.
140 Recursion will subsequently be done by genAnnotBinds. It may
141 optionally remove a Note TermUsg, or optionally add one if it is not
142 already present, but it may perform NO OTHER MODIFICATIONS to the
143 structure of the term.
145 We do different things to types of variables bound locally and of
146 variables bound in other modules, in certain cases: the former get
147 uvars and the latter keep their existing annotations when we annotate,
148 for example. To control this, @MungeFlags@ describes what kind of a
149 type this is that we're about to munge.
152 data MungeFlags = MungeFlags { isSigma :: Bool, -- want annotated on top (sigma type)
153 isLocal :: Bool, -- is locally-defined type
154 hasUsg :: Bool, -- has fixed usage info, don't touch
155 isExp :: Bool, -- is exported (and must be pessimised)
156 mfLoc :: SDoc -- location info
159 tauTyMF loc = MungeFlags { isSigma = False, isLocal = True,
160 hasUsg = False, isExp = False, mfLoc = loc }
161 sigVarTyMF v = MungeFlags { isSigma = True, isLocal = hasLocalDef v,
162 hasUsg = hasUsgInfo v, isExp = isExportedId v,
163 mfLoc = ptext SLIT("type of binder") <+> ppr v }
166 The helper functions @tauTyMF@ and @sigVarTyMF@ create @MungeFlags@
167 for us. @sigVarTyMF@ checks the variable to see how to set the flags.
169 @hasLocalDef@ tells us if the given variable has an actual local
170 definition that we can play with. This is not quite the same as
171 @isLocallyDefined@, since @mayHaveNoBindingId@ things (usually) don't have
172 a local definition - the simplifier will inline whatever their
173 unfolding is anyway. We treat these as if they were externally
174 defined, since we don't have access to their definition (at least not
175 easily). This doesn't hurt much, since after the simplifier has run
176 the unfolding will have been inlined and we can access the unfolding
179 @hasUsgInfo@, on the other hand, says if the variable already has
180 usage info in its type that must at all costs be preserved. This is
181 assumed true (exactly) of all imported ids.
184 hasLocalDef :: Var -> Bool
185 hasLocalDef var = isLocallyDefined var
186 && not (mayHaveNoBinding var)
188 hasUsgInfo :: Var -> Bool
189 hasUsgInfo var = (not . isLocallyDefined) var
192 Here's the walk itself.
195 genAnnotBinds :: (MungeFlags -> Type -> AnnotM flexi Type)
196 -> (CoreExpr -> AnnotM flexi CoreExpr) -- see caveats above
198 -> AnnotM flexi [CoreBind]
200 genAnnotBinds _ _ [] = return []
202 genAnnotBinds f g (b:bs) = do { (b',vs,vs') <- genAnnotBind f g b
203 ; bs' <- withAnnVars vs vs' $
208 genAnnotBind :: (MungeFlags -> Type -> AnnotM flexi Type) -- type-altering function
209 -> (CoreExpr -> AnnotM flexi CoreExpr) -- term-altering function
210 -> CoreBind -- original CoreBind
212 (CoreBind, -- annotated CoreBind
213 [Var], -- old variables, to be mapped to...
214 [Var]) -- ... new variables
216 genAnnotBind f g (NonRec v1 e1) = do { v1' <- genAnnotVar f v1
217 ; e1' <- genAnnotCE f g e1
218 ; return (NonRec v1' e1', [v1], [v1'])
221 genAnnotBind f g (Rec ves) = do { let (vs,es) = unzip ves
222 ; vs' <- mapM (genAnnotVar f) vs
223 ; es' <- withAnnVars vs vs' $
224 mapM (genAnnotCE f g) es
225 ; return (Rec (zip vs' es'), vs, vs')
228 genAnnotCE :: (MungeFlags -> Type -> AnnotM flexi Type) -- type-altering function
229 -> (CoreExpr -> AnnotM flexi CoreExpr) -- term-altering function
230 -> CoreExpr -- original expression
231 -> AnnotM flexi CoreExpr -- yields new expression
233 genAnnotCE mungeType mungeTerm = go
234 where go e0@(Var v) | isTyVar v = return e0 -- arises, e.g., as tyargs of constructor
235 -- (no it doesn't: (Type (TyVar tyvar))
236 | otherwise = do { mv' <- lookupAnnVar v
238 Just var -> return var
239 Nothing -> fixedVar v
243 go (Lit l) = -- we know it's saturated
246 go (App e arg) = do { e' <- go e
248 ; return (App e' arg')
251 go e0@(Lam v0 _) = do { e1 <- (if isTyVar v0 then return else mungeTerm) e0
253 = case e1 of -- munge may have added note
254 Note tu@(TermUsg _) (Lam v e2)
256 Lam v e2 -> (v,e2,id)
257 ; v' <- genAnnotVar mungeType v
258 ; e' <- withAnnVar v v' $ go e2
259 ; return (wrap (Lam v' e'))
262 go (Let bind e) = do { (bind',vs,vs') <- genAnnotBind mungeType mungeTerm bind
263 ; e' <- withAnnVars vs vs' $ go e
264 ; return (Let bind' e')
267 go (Case e v alts) = do { e' <- go e
268 ; v' <- genAnnotVar mungeType v
269 ; alts' <- withAnnVar v v' $ mapM genAnnotAlt alts
270 ; return (Case e' v' alts')
273 go (Note scc@(SCC _) e) = do { e' <- go e
274 ; return (Note scc e')
276 go e0@(Note (Coerce ty1 ty0)
277 e) = do { ty1' <- mungeType
278 (tauTyMF (ptext SLIT("coercer of")
281 (tauTyMF (ptext SLIT("coercee of")
283 -- (Better to specify ty0'
284 -- identical to the type of e, including
285 -- annotations, right at the beginning, but
286 -- not possible at this point.)
288 ; return (Note (Coerce ty1' ty0') e')
290 go (Note InlineCall e) = do { e' <- go e
291 ; return (Note InlineCall e')
293 go (Note InlineMe e) = do { e' <- go e
294 ; return (Note InlineMe e')
296 go e0@(Note (TermUsg _) _) = do { e1 <- mungeTerm e0
297 ; case e1 of -- munge may have removed note
298 Note tu@(TermUsg _) e2 -> do { e3 <- go e2
299 ; return (Note tu e3)
304 go e0@(Type ty) = -- should only occur at toplevel of Arg,
306 do { ty' <- mungeType
307 (tauTyMF (ptext SLIT("tyarg")
312 fixedVar v = ASSERT2( not (hasLocalDef v), text "genAnnotCE: locally defined var" <+> ppr v <+> text "not in varenv" )
313 genAnnotVar mungeType v
315 genAnnotAlt (c,vs,e) = do { vs' <- mapM (genAnnotVar mungeType) vs
316 ; e' <- withAnnVars vs vs' $ go e
317 ; return (c, vs', e')
321 genAnnotVar :: (MungeFlags -> Type -> AnnotM flexi Type)
325 genAnnotVar mungeType v | isTyVar v = return v
326 | otherwise = do { vty' <- mungeType (sigVarTyMF v) (varType v)
327 ; return (setVarType v vty')
331 pprTrace "genAnnotVar" (ppr (tyUsg vty') <+> ppr v) $
337 ======================================================================
339 Some specific things to do to types inside terms
340 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
342 @annotTyM@ annotates a type with fresh uvars everywhere the inference
343 is allowed to go, and leaves alone annotations where it may not go.
345 We assume there are no annotations already.
348 annotTyM :: MungeFlags -> Type -> AnnotM UniqSupply Type
350 annotTyM mf ty = uniqSMtoAnnotM . uniqSMMToUs $
351 case (hasUsg mf, isLocal mf, isSigma mf) of
352 (True ,_ ,_ ) -> ASSERT( isUsgTy ty )
354 (False,True ,True ) -> if isExp mf then
355 annotTyP (tag 'p') ty
358 (False,True ,False) -> annotTyN (tag 't') ty
359 (False,False,True ) -> return $ annotMany ty -- assume worst
360 (False,False,False) -> return $ annotManyN ty
361 where tag c = Right $ "annotTyM:" ++ [c] ++ ": " ++ showSDoc (ppr ty)
363 -- specific functions for annotating tau and sigma types
366 annotTy tag = genAnnotTy (newVarUSMM tag)
367 annotTyN tag = genAnnotTyN (newVarUSMM tag)
369 -- ...with uvars and pessimal Manys (for exported ids)
370 annotTyP tag ty = do { ty' <- annotTy tag ty ; return (pessimise True ty') }
373 annotMany, annotManyN :: Type -> Type
378 annotMany ty = unId (genAnnotTy (return UsMany) ty)
379 annotManyN ty = unId (genAnnotTyN (return UsMany) ty)
382 -- monad required for the above
383 newtype Id a = Id a ; unId (Id a) = a
384 instance Monad Id where { a >>= f = f (unId a) ; return a = Id a }
386 -- lambda-annotating function for use along with the above
387 annotLam e0@(Lam v e) = do { uv <- uniqSMtoAnnotM $ newVarUs (Left e0)
388 ; return (Note (TermUsg uv) (Lam v e))
390 annotLam (Note (TermUsg _) _) = panic "annotLam: unexpected term usage annot"
393 The above requires a `pessimising' translation. This is applied to
394 types of exported ids, and ensures that they have a fully general
395 type (since we don't know how they will be used in other modules).
398 pessimise :: Bool -> Type -> Type
401 pessimise co ty0@(NoteTy usg@(UsgNote u ) ty)
403 then case u of UsMany -> pty
404 UsVar _ -> pty -- force to UsMany
405 UsOnce -> pprPanic "pessimise:" (ppr ty0)
407 where pty = pessimiseN co ty
409 pessimise co ty0 = pessimiseN co ty0 -- assume UsMany
411 pessimise co ty0@(NoteTy usg@(UsgNote u ) ty)
413 then case u of UsMany -> NoteTy usg pty
414 UsVar _ -> NoteTy (UsgNote UsMany) pty
415 UsOnce -> pprPanic "pessimise:" (ppr ty0)
417 where pty = pessimiseN co ty
419 pessimise co ty0 = pprPanic "pessimise: missing usage note:" $
423 pessimiseN co ty0@(NoteTy usg@(UsgNote _ ) ty) = pprPanic "pessimiseN: unexpected usage note:" $
425 pessimiseN co (NoteTy (SynNote sty) ty) = NoteTy (SynNote (pessimiseN co sty))
427 pessimiseN co (NoteTy note@(FTVNote _ ) ty) = NoteTy note (pessimiseN co ty)
428 pessimiseN co ty0@(TyVarTy _) = ty0
429 pessimiseN co ty0@(AppTy _ _) = ty0
430 pessimiseN co ty0@(TyConApp tc tys) = ASSERT( not ((isFunTyCon tc) && (length tys > 1)) )
432 pessimiseN co (FunTy ty1 ty2) = FunTy (pessimise (not co) ty1)
434 pessimiseN co (ForAllTy tyv ty) = ForAllTy tyv (pessimiseN co ty)
438 @unAnnotTyM@ strips annotations (that the inference is allowed to
439 touch) from a term, and `fixes' those it isn't permitted to touch (by
440 putting @Many@ annotations where they are missing, but leaving
441 existing annotations in the type).
443 @unTermUsg@ removes from a term any term usage annotations it finds.
446 unAnnotTyM :: MungeFlags -> Type -> AnnotM a Type
448 unAnnotTyM mf ty = if hasUsg mf then
450 return (fixAnnotTy ty)
451 else return (unannotTy ty)
454 unTermUsg :: CoreExpr -> AnnotM a CoreExpr
455 -- strip all term annotations
456 unTermUsg e@(Lam _ _) = return e
457 unTermUsg (Note (TermUsg _) e) = return e
458 unTermUsg _ = panic "unTermUsg"
460 unannotTy :: Type -> Type
461 -- strip all annotations
462 unannotTy (NoteTy (UsgForAll uv) ty) = unannotTy ty
463 unannotTy (NoteTy (UsgNote _ ) ty) = unannotTy ty
464 unannotTy (NoteTy (SynNote sty) ty) = NoteTy (SynNote (unannotTy sty)) (unannotTy ty)
465 unannotTy (NoteTy note@(FTVNote _ ) ty) = NoteTy note (unannotTy ty)
466 -- IP notes need to be preserved
467 unannotTy ty@(NoteTy (IPNote _) _) = ty
468 unannotTy ty@(TyVarTy _) = ty
469 unannotTy (AppTy ty1 ty2) = AppTy (unannotTy ty1) (unannotTy ty2)
470 unannotTy (TyConApp tc tys) = TyConApp tc (map unannotTy tys)
471 unannotTy (FunTy ty1 ty2) = FunTy (unannotTy ty1) (unannotTy ty2)
472 unannotTy (ForAllTy tyv ty) = ForAllTy tyv (unannotTy ty)
475 fixAnnotTy :: Type -> Type
476 -- put Manys where they are missing
480 fixAnnotTy (NoteTy note@(UsgForAll uv) ty) = NoteTy note (fixAnnotTy ty)
481 fixAnnotTy (NoteTy note@(UsgNote _ ) ty) = NoteTy note (fixAnnotTyN ty)
482 fixAnnotTy ty0 = NoteTy (UsgNote UsMany) (fixAnnotTyN ty0)
484 fixAnnotTyN ty0@(NoteTy note@(UsgNote _ ) ty) = pprPanic "fixAnnotTyN: unexpected usage note:" $
486 fixAnnotTyN (NoteTy (SynNote sty) ty) = NoteTy (SynNote (fixAnnotTyN sty))
488 fixAnnotTyN (NoteTy note@(FTVNote _ ) ty) = NoteTy note (fixAnnotTyN ty)
489 fixAnnotTyN ty0@(TyVarTy _) = ty0
490 fixAnnotTyN (AppTy ty1 ty2) = AppTy (fixAnnotTyN ty1) (fixAnnotTyN ty2)
491 fixAnnotTyN (TyConApp tc tys) = ASSERT( isFunTyCon tc || isAlgTyCon tc || isPrimTyCon tc || isSynTyCon tc )
492 TyConApp tc (map (if isFunTyCon tc then
496 fixAnnotTyN (FunTy ty1 ty2) = FunTy (fixAnnotTy ty1) (fixAnnotTy ty2)
497 fixAnnotTyN (ForAllTy tyv ty) = ForAllTy tyv (fixAnnotTyN ty)
501 The composition (reannotating a type with fresh uvars but the same
502 structure) is useful elsewhere:
505 freshannotTy :: Type -> UniqSMM Type
506 freshannotTy = annotTy (Right "freshannotTy") . unannotTy
510 Wrappers apply these functions to sets of bindings.
513 doAnnotBinds :: UniqSupply
515 -> ([CoreBind],UniqSupply)
517 doAnnotBinds us binds = initAnnotM us (genAnnotBinds annotTyM annotLam binds)
520 doUnAnnotBinds :: [CoreBind]
523 doUnAnnotBinds binds = fst $ initAnnotM () $
524 genAnnotBinds unAnnotTyM unTermUsg binds
527 ======================================================================
532 The @UniqSM@ type is not an instance of @Monad@, and cannot be made so
533 since it is merely a synonym rather than a newtype. Here we define
534 @UniqSMM@, which *is* an instance of @Monad@.
537 newtype UniqSMM a = UsToUniqSMM (UniqSM a)
538 uniqSMMToUs (UsToUniqSMM us) = us
539 usToUniqSMM = UsToUniqSMM
541 instance Monad UniqSMM where
542 m >>= f = UsToUniqSMM $ uniqSMMToUs m `thenUs` \ a ->
544 return = UsToUniqSMM . returnUs
548 For annotation, the monad @AnnotM@, we need to carry around our
549 variable mapping, along with some general state.
552 newtype AnnotM flexi a = AnnotM ( flexi -- UniqSupply etc
553 -> VarEnv Var -- unannotated to annotated variables
554 -> (a,flexi,VarEnv Var))
555 unAnnotM (AnnotM f) = f
557 instance Monad (AnnotM flexi) where
558 a >>= f = AnnotM (\ us ve -> let (r,us',ve') = unAnnotM a us ve
559 in unAnnotM (f r) us' ve')
560 return a = AnnotM (\ us ve -> (a,us,ve))
562 initAnnotM :: fl -> AnnotM fl a -> (a,fl)
563 initAnnotM fl m = case (unAnnotM m) fl emptyVarEnv of { (r,fl',_) -> (r,fl') }
565 withAnnVar :: Var -> Var -> AnnotM fl a -> AnnotM fl a
566 withAnnVar v v' m = AnnotM (\ us ve -> let ve' = extendVarEnv ve v v'
567 (r,us',_) = (unAnnotM m) us ve'
570 withAnnVars :: [Var] -> [Var] -> AnnotM fl a -> AnnotM fl a
571 withAnnVars vs vs' m = AnnotM (\ us ve -> let ve' = plusVarEnv ve (zipVarEnv vs vs')
572 (r,us',_) = (unAnnotM m) us ve'
575 lookupAnnVar :: Var -> AnnotM fl (Maybe Var)
576 lookupAnnVar var = AnnotM (\ us ve -> (lookupVarEnv ve var,
581 A useful helper allows us to turn a computation in the unique supply
582 monad into one in the annotation monad parameterised by a unique
586 uniqSMtoAnnotM :: UniqSM a -> AnnotM UniqSupply a
588 uniqSMtoAnnotM m = AnnotM (\ us ve -> let (r,us') = initUs us m
592 @newVarUs@ and @newVarUSMM@ generate a new usage variable. They take
593 an argument which is used for debugging only, describing what the
594 variable is to annotate.
597 newVarUs :: (Either CoreExpr String) -> UniqSM UsageAnn
598 -- the first arg is for debugging use only
599 newVarUs e = getUniqueUs `thenUs` \ u ->
604 Left (Lit _) -> "literal"
605 Left (Lam v e) -> "lambda: " ++ showSDoc (ppr v)
608 in pprTrace "newVarUs:" (ppr uv <+> text src) $
612 newVarUSMM :: (Either CoreExpr String) -> UniqSMM UsageAnn
613 newVarUSMM = usToUniqSMM . newVarUs
616 ======================================================================
618 PrimOps and usage information.
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).
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
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,
638 ======================================================================