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               ( 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 )
  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 @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
 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 (mayHaveNoBinding 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 -- 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)
 474
 475
 476 fixAnnotTy :: Type -> Type
 477 -- put Manys where they are missing
 478 #ifndef USMANY
 479 fixAnnotTy = id
 480 #else
 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)
 484
 485 fixAnnotTyN ty0@(NoteTy note@(UsgNote _  ) ty) = pprPanic "fixAnnotTyN: unexpected usage note:" $
 486                                                           ppr ty0
 487 fixAnnotTyN     (NoteTy      (SynNote sty) ty) = NoteTy (SynNote (fixAnnotTyN sty))
 488                                                                  (fixAnnotTyN ty )
 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
 494                                                                      fixAnnotTy
 495                                                                    else
 496                                                                      fixAnnotTyN) tys)
 497 fixAnnotTyN     (FunTy ty1 ty2)                = FunTy (fixAnnotTy ty1) (fixAnnotTy ty2)
 498 fixAnnotTyN     (ForAllTy tyv ty)              = ForAllTy tyv (fixAnnotTyN ty)
 499 #endif
 500 \end{code}
 501
 502 The composition (reannotating a type with fresh uvars but the same
 503 structure) is useful elsewhere:
 504
 505 \begin{code}
 506 freshannotTy :: Type -> UniqSMM Type
 507 freshannotTy = annotTy (Right "freshannotTy") . unannotTy
 508 \end{code}
 509
 510
 511 Wrappers apply these functions to sets of bindings.
 512
 513 \begin{code}
 514 doAnnotBinds :: UniqSupply
 515              -> [CoreBind]
 516              -> ([CoreBind],UniqSupply)
 517
 518 doAnnotBinds us binds = initAnnotM us (genAnnotBinds annotTyM annotLam binds)
 519
 520
 521 doUnAnnotBinds :: [CoreBind]
 522                -> [CoreBind]
 523
 524 doUnAnnotBinds binds = fst $ initAnnotM () $
 525                          genAnnotBinds unAnnotTyM unTermUsg binds
 526 \end{code}
 527
 528 ======================================================================
 529
 530 Monadic machinery
 531 ~~~~~~~~~~~~~~~~~
 532
 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@.
 536
 537 \begin{code}
 538 newtype UniqSMM a = UsToUniqSMM (UniqSM a)
 539 uniqSMMToUs (UsToUniqSMM us) = us
 540 usToUniqSMM = UsToUniqSMM
 541
 542 instance Monad UniqSMM where
 543   m >>= f  = UsToUniqSMM $ uniqSMMToUs m `thenUs` \ a ->
 544                            uniqSMMToUs (f a)
 545   return   = UsToUniqSMM . returnUs
 546 \end{code}
 547
 548
 549 For annotation, the monad @AnnotM@, we need to carry around our
 550 variable mapping, along with some general state.
 551
 552 \begin{code}
 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
 557
 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))
 562
 563 initAnnotM :: fl -> AnnotM fl a -> (a,fl)
 564 initAnnotM fl m = case (unAnnotM m) fl emptyVarEnv of { (r,fl',_) -> (r,fl') }
 565
 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'
 569                                        in  (r,us',ve))
 570
 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'
 574                                           in  (r,us',ve))
 575
 576 lookupAnnVar :: IdOrTyVar -> AnnotM fl (Maybe IdOrTyVar)
 577 lookupAnnVar var = AnnotM (\ us ve -> (lookupVarEnv ve var,
 578                                        us,
 579                                        ve))
 580 \end{code}
 581
 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
 584 supply.
 585
 586 \begin{code}
 587 uniqSMtoAnnotM :: UniqSM a -> AnnotM UniqSupply a
 588
 589 uniqSMtoAnnotM m = AnnotM (\ us ve -> let (r,us') = initUs us m
 590                                       in  (r,us',ve))
 591 \end{code}
 592
 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.
 596
 597 \begin{code}
 598 newVarUs :: (Either CoreExpr String) -> UniqSM UsageAnn
 599 -- the first arg is for debugging use only
 600 newVarUs e = getUniqueUs `thenUs` \ u ->
 601              let uv = mkUVar u in
 602              returnUs (UsVar uv)
 603 {- #ifdef DEBUG
 604              let src = case e of
 605                          Left (Con (Literal _) _) -> "literal"
 606                          Left (Con _           _) -> "primop"
 607                          Left (Lam v e)           -> "lambda: " ++ showSDoc (ppr v)
 608                          Left _                   -> "unknown"
 609                          Right s                  -> s
 610              in pprTrace "newVarUs:" (ppr uv <+> text src) $
 611    #endif
 612  -}
 613
 614 newVarUSMM :: (Either CoreExpr String) -> UniqSMM UsageAnn
 615 newVarUSMM = usToUniqSMM . newVarUs
 616 \end{code}
 617
 618 ======================================================================
 619
 620 PrimOps and usage information.
 621
 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).
 625
 626 \begin{code}
 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
 631
 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,
 637                           rty1u)
 638 \end{code}
 639
 640 ======================================================================
 641
 642 EOF