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