ghc/compiler/specialise/SpecMonad.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1995
   3 %
   4 \section[SpecMonad]{Monad for the Specialiser}
   5
   6 \begin{code}
   7 #include "HsVersions.h"
   8
   9 module SpecMonad where
  10
  11 import PlainCore
  12 import SpecTyFuns
  13
  14 IMPORT_Trace
  15 import Outputable       -- ToDo: these may be removable...
  16 import Pretty
  17
  18 import AbsUniType
  19 import Bag
  20 import CmdLineOpts      ( GlobalSwitch(..) )
  21 import CoreLift         ( mkLiftedId, liftExpr, bindUnlift, applyBindUnlifts )
  22 import IdEnv
  23 import Id
  24 import IdInfo
  25 import InstEnv          ( lookupClassInstAtSimpleType )
  26 import Maybes           ( catMaybes, firstJust, maybeToBool, Maybe(..) )
  27 import TyVarEnv         -- ( growTyVarEnvList, nullTyVarEnv, TyVarEnv, TypeEnv(..) )
  28 import Util
  29 import UniqSet
  30 import SplitUniq
  31
  32 infixr 9 `thenSM`
  33 \end{code}
  34
  35 %************************************************************************
  36 %*                                                                      *
  37 \subsection[cloning-binders]{The Specialising IdEnv and CloneInfo}
  38 %*                                                                      *
  39 %************************************************************************
  40
  41 @SpecIdEnv@ maps old Ids to their new "clone". There are three cases:
  42
  43 1) (NoLift CoLitAtom l) : an Id which is bound to a literal
  44
  45 2) (NoLift CoLitAtom l) : an Id bound to a "new" Id
  46    The new Id is a possibly-type-specialised clone of the original
  47
  48 3) Lifted lifted_id unlifted_id :
  49
  50    This indicates that the original Id has been specialised to an
  51    unboxed value which must be lifted (see "Unboxed bindings" above)
  52      @unlifted_id@ is the unboxed clone of the original Id
  53      @lifted_id@ is a *lifted* version of the original Id
  54
  55    When you lookup Ids which are Lifted, you have to insert a case
  56    expression to un-lift the value (done with @bindUnlift@)
  57
  58    You also have to insert a case to lift the value in the binding
  59    (done with @liftExpr@)
  60
  61
  62 \begin{code}
  63 type SpecIdEnv = IdEnv CloneInfo
  64
  65 data CloneInfo
  66  = NoLift PlainCoreAtom -- refers to cloned id or literal
  67
  68  | Lifted Id            -- lifted, cloned id
  69           Id            -- unlifted, cloned id
  70
  71 \end{code}
  72
  73 %************************************************************************
  74 %*                                                                      *
  75 \subsection[monad-Specialise]{Monad used in specialisation}
  76 %*                                                                      *
  77 %************************************************************************
  78
  79 Monad has:
  80
  81  inherited: control flags and
  82             recordInst functions with flags cached
  83
  84             environment mapping tyvars to types
  85             environment mapping Ids to Atoms
  86
  87  threaded in and out: unique supply
  88
  89 \begin{code}
  90 type SpecM result
  91   =  (GlobalSwitch -> Bool)
  92   -> TypeEnv
  93   -> SpecIdEnv
  94   -> SplitUniqSupply
  95   -> result
  96
  97 initSM m sw_chker uniqs
  98   = m sw_chker nullTyVarEnv nullIdEnv uniqs
  99
 100 returnSM :: a -> SpecM a
 101 thenSM   :: SpecM a -> (a -> SpecM b) -> SpecM b
 102 fixSM    :: (a -> SpecM a) -> SpecM a
 103
 104 thenSM m k sw_chkr tvenv idenv us
 105   = case splitUniqSupply us        of { (s1, s2) ->
 106     case (m sw_chkr tvenv idenv s1) of { r ->
 107     k r sw_chkr tvenv idenv s2 }}
 108
 109 returnSM r sw_chkr tvenv idenv us = r
 110
 111 fixSM k sw_chkr tvenv idenv us
 112  = r
 113  where
 114    r = k r sw_chkr tvenv idenv us       -- Recursive in r!
 115 \end{code}
 116
 117
 118 \begin{code}
 119 getSwitchCheckerSM sw_chkr tvenv idenv us = sw_chkr
 120 \end{code}
 121
 122 The only interesting bit is figuring out the type of the SpecId!
 123
 124 \begin{code}
 125 newSpecIds :: [Id]              -- The id of which to make a specialised version
 126            -> [Maybe UniType]   -- Specialise to these types
 127            -> Int               -- No of dicts to specialise
 128            -> SpecM [Id]
 129
 130 newSpecIds new_ids maybe_tys dicts_to_ignore sw_chkr tvenv idenv us
 131   = [ mkSpecId uniq id maybe_tys (spec_id_ty id) (selectIdInfoForSpecId id)
 132       | (id,uniq) <- new_ids `zip` uniqs ]
 133   where
 134     uniqs = getSUniques (length new_ids) us
 135     spec_id_ty id = specialiseTy (getIdUniType id) maybe_tys dicts_to_ignore
 136
 137 newTyVars :: Int -> SpecM [TyVar]
 138 newTyVars n sw_chkr tvenv idenv us
 139  = map mkPolySysTyVar uniqs
 140  where
 141    uniqs = getSUniques n us
 142 \end{code}
 143
 144 @cloneLambdaOrCaseBinders@ and @cloneLetBinders@ take a bunch of
 145 binders, and build ``clones'' for them.  The clones differ from the
 146 originals in three ways:
 147
 148         (a) they have a fresh unique
 149         (b) they have the current type environment applied to their type
 150         (c) for Let binders which have been specialised to unboxed values
 151             the clone will have a lifted type
 152
 153 As well as returning the list of cloned @Id@s they also return a list of
 154 @CloneInfo@s which the original binders should be bound to.
 155
 156 \begin{code}
 157 cloneLambdaOrCaseBinders :: [Id]                        -- Old binders
 158                          -> SpecM ([Id], [CloneInfo])   -- New ones
 159
 160 cloneLambdaOrCaseBinders old_ids sw_chkr tvenv idenv us
 161   = let
 162         uniqs = getSUniques (length old_ids) us
 163     in
 164     unzip (zipWith clone_it old_ids uniqs)
 165   where
 166     clone_it old_id uniq
 167       = (new_id, NoLift (CoVarAtom new_id))
 168       where
 169         new_id = applyTypeEnvToId tvenv (mkIdWithNewUniq old_id uniq)
 170
 171 cloneLetBinders :: Bool                         -- Top level ?
 172                 -> Bool                         -- Recursice
 173                 -> [Id]                         -- Old binders
 174                 -> SpecM ([Id], [CloneInfo])    -- New ones
 175
 176 cloneLetBinders top_lev is_rec old_ids sw_chkr tvenv idenv us
 177   = let
 178         uniqs = getSUniques (2 * length old_ids) us
 179     in
 180     unzip (clone_them old_ids uniqs)
 181   where
 182     clone_them [] [] = []
 183
 184     clone_them (old_id:olds) (u1:u2:uniqs)
 185       | top_lev
 186         = (old_id,
 187            NoLift (CoVarAtom old_id)) : clone_rest
 188
 189          -- Don't clone if it is a top-level thing. Why not?
 190          -- (a) we don't want to change the uniques
 191          --     on such things (see TopLevId in Id.lhs)
 192          -- (b) we don't have to be paranoid about name capture
 193          -- (c) the thing is polymorphic so no need to subst
 194
 195       | otherwise
 196         = if (is_rec && isUnboxedDataType new_ty && not (isUnboxedDataType old_ty))
 197           then (lifted_id,
 198                 Lifted lifted_id unlifted_id) : clone_rest
 199           else (new_id,
 200                 NoLift (CoVarAtom new_id)) : clone_rest
 201
 202       where
 203         clone_rest = clone_them olds uniqs
 204
 205         new_id = applyTypeEnvToId tvenv (mkIdWithNewUniq old_id u1)
 206         new_ty = getIdUniType new_id
 207         old_ty = getIdUniType old_id
 208
 209         (lifted_id, unlifted_id) = mkLiftedId new_id u2
 210
 211
 212 cloneTyVarSM :: TyVar -> SpecM TyVar
 213
 214 cloneTyVarSM old_tyvar sw_chkr tvenv idenv us
 215   = let
 216         uniq = getSUnique us
 217     in
 218     cloneTyVar old_tyvar uniq -- new_tyvar
 219
 220 bindId :: Id -> CloneInfo -> SpecM thing -> SpecM thing
 221
 222 bindId id val specm sw_chkr tvenv idenv us
 223  = specm sw_chkr tvenv (addOneToIdEnv idenv id val) us
 224
 225 bindIds :: [Id] -> [CloneInfo] -> SpecM thing -> SpecM thing
 226
 227 bindIds olds news specm sw_chkr tvenv idenv us
 228  = specm sw_chkr tvenv (growIdEnvList idenv (zip olds news)) us
 229
 230 bindSpecIds :: [Id]                     -- Old
 231             -> [(CloneInfo)]            -- New
 232             -> [[Maybe SpecInfo]]       -- Corresponding specialisations
 233                                         -- Each sub-list corresponds to a different type,
 234                                         -- and contains one Maybe spec_info for each id
 235             -> SpecM thing
 236             -> SpecM thing
 237
 238 bindSpecIds olds clones spec_infos specm sw_chkr tvenv idenv us
 239  = specm sw_chkr tvenv (growIdEnvList idenv old_to_clone) us
 240  where
 241    old_to_clone = mk_old_to_clone olds clones spec_infos
 242
 243    -- The important thing here is that we are *lazy* in spec_infos
 244    mk_old_to_clone [] [] _ = []
 245    mk_old_to_clone (old:rest_olds) (clone:rest_clones) spec_infos
 246      = (old, add_spec_info clone) :
 247        mk_old_to_clone rest_olds rest_clones spec_infos_rest
 248      where
 249        add_spec_info (NoLift (CoVarAtom new))
 250          = NoLift (CoVarAtom (new `addIdSpecialisation`
 251                                   (mkSpecEnv spec_infos_this_id)))
 252        add_spec_info lifted
 253          = lifted               -- no specialised instances for unboxed lifted values
 254
 255        spec_infos_this_id = catMaybes (map head spec_infos)
 256        spec_infos_rest    = map tail spec_infos
 257
 258
 259 bindTyVar :: TyVar -> UniType -> SpecM thing -> SpecM thing
 260
 261 bindTyVar tyvar ty specm sw_chkr tvenv idenv us
 262  = specm sw_chkr (growTyVarEnvList tvenv [(tyvar,ty)]) idenv us
 263 \end{code}
 264
 265 \begin{code}
 266 lookupId :: Id -> SpecM CloneInfo
 267
 268 lookupId id sw_chkr tvenv idenv us
 269   = case lookupIdEnv idenv id of
 270       Nothing   -> NoLift (CoVarAtom id)
 271       Just info -> info
 272 \end{code}
 273
 274 \begin{code}
 275 specTy :: UniType -> SpecM UniType      -- Apply the current type envt to the type
 276
 277 specTy ty sw_chkr tvenv idenv us
 278   = applyTypeEnvToTy tvenv ty
 279 \end{code}
 280
 281 \begin{code}
 282 liftId :: Id -> SpecM (Id, Id)
 283 liftId id sw_chkr tvenv idenv us
 284   = let
 285         uniq = getSUnique us
 286     in
 287     mkLiftedId id uniq
 288 \end{code}
 289
 290 In other monads these @mapSM@ things are usually called @listM@.
 291 I think @mapSM@ is a much better name.  The `2' and `3' variants are
 292 when you want to return two or three results, and get at them
 293 separately.  It saves you having to do an (unzip stuff) right after.
 294
 295 \begin{code}
 296 mapSM          :: (a -> SpecM b)            -> [a] -> SpecM [b]
 297 mapAndUnzipSM  :: (a -> SpecM (b1, b2))     -> [a] -> SpecM ([b1],[b2])
 298 mapAndUnzip3SM :: (a -> SpecM (b1, b2, b3)) -> [a] -> SpecM ([b1],[b2],[b3])
 299 mapAndUnzip4SM :: (a -> SpecM (b1, b2, b3, b4)) -> [a] -> SpecM ([b1],[b2],[b3],[b4])
 300
 301 mapSM f [] = returnSM []
 302 mapSM f (x:xs) = f x            `thenSM` \ r ->
 303                  mapSM f xs     `thenSM` \ rs ->
 304                  returnSM (r:rs)
 305
 306 mapAndUnzipSM f [] = returnSM ([],[])
 307 mapAndUnzipSM f (x:xs) = f x                    `thenSM` \ (r1, r2) ->
 308                          mapAndUnzipSM f xs     `thenSM` \ (rs1,rs2) ->
 309                          returnSM ((r1:rs1),(r2:rs2))
 310
 311 mapAndUnzip3SM f [] = returnSM ([],[],[])
 312 mapAndUnzip3SM f (x:xs) = f x                   `thenSM` \ (r1,r2,r3) ->
 313                           mapAndUnzip3SM f xs   `thenSM` \ (rs1,rs2,rs3) ->
 314                           returnSM ((r1:rs1),(r2:rs2),(r3:rs3))
 315
 316 mapAndUnzip4SM f [] = returnSM ([],[],[],[])
 317 mapAndUnzip4SM f (x:xs) = f x                   `thenSM` \ (r1,r2,r3,r4) ->
 318                           mapAndUnzip4SM f xs   `thenSM` \ (rs1,rs2,rs3,rs4) ->
 319                           returnSM ((r1:rs1),(r2:rs2),(r3:rs3),(r4:rs4))
 320 \end{code}