ghc/compiler/basicTypes/UniqSupply.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1996
   3 %
   4 \section[UniqSupply]{The @UniqueSupply@ data type and a (monadic) supply thereof}
   5
   6 \begin{code}
   7 #include "HsVersions.h"
   8
   9 module UniqSupply (
  10
  11         UniqSupply,             -- Abstractly
  12
  13         getUnique, getUniques,  -- basic ops
  14
  15         UniqSM(..),             -- type: unique supply monad
  16         initUs, thenUs, returnUs,
  17         mapUs, mapAndUnzipUs, mapAndUnzip3Us,
  18         thenMaybeUs, mapAccumLUs,
  19
  20         mkSplitUniqSupply,
  21         splitUniqSupply
  22   ) where
  23
  24 IMP_Ubiq(){-uitous-}
  25
  26 import Unique
  27 import Util
  28
  29 import PreludeGlaST
  30
  31 w2i x = word2Int# x
  32 i2w x = int2Word# x
  33 i2w_s x = (x :: Int#)
  34 \end{code}
  35
  36
  37 %************************************************************************
  38 %*                                                                      *
  39 \subsection{Splittable Unique supply: @UniqSupply@}
  40 %*                                                                      *
  41 %************************************************************************
  42
  43 %************************************************************************
  44 %*                                                                      *
  45 \subsubsection[UniqSupply-type]{@UniqSupply@ type and operations}
  46 %*                                                                      *
  47 %************************************************************************
  48
  49 A value of type @UniqSupply@ is unique, and it can
  50 supply {\em one} distinct @Unique@.  Also, from the supply, one can
  51 also manufacture an arbitrary number of further @UniqueSupplies@,
  52 which will be distinct from the first and from all others.
  53
  54 \begin{code}
  55 data UniqSupply
  56   = MkSplitUniqSupply Int       -- make the Unique with this
  57                    UniqSupply UniqSupply
  58                                 -- when split => these two supplies
  59 \end{code}
  60
  61 \begin{code}
  62 mkSplitUniqSupply :: Char -> IO UniqSupply
  63
  64 splitUniqSupply :: UniqSupply -> (UniqSupply, UniqSupply)
  65 getUnique :: UniqSupply -> Unique
  66 getUniques :: Int -> UniqSupply -> [Unique]
  67 \end{code}
  68
  69 \begin{code}
  70 mkSplitUniqSupply (C# c#)
  71   = let
  72         mask# = (i2w (ord# c#)) `shiftL#` (i2w_s 24#)
  73
  74         -- here comes THE MAGIC:
  75
  76         mk_supply#
  77           = unsafe_interleave (
  78                 mk_unique   `thenPrimIO` \ uniq ->
  79                 mk_supply#  `thenPrimIO` \ s1 ->
  80                 mk_supply#  `thenPrimIO` \ s2 ->
  81                 returnPrimIO (MkSplitUniqSupply uniq s1 s2)
  82             )
  83           where
  84             -- inlined copy of unsafeInterleavePrimIO;
  85             -- this is the single-most-hammered bit of code
  86             -- in the compiler....
  87             unsafe_interleave m s
  88               = let
  89                     (r, new_s) = m s
  90                 in
  91                 (r, s)
  92
  93         mk_unique = _ccall_ genSymZh            `thenPrimIO` \ (W# u#) ->
  94                     returnPrimIO (I# (w2i (mask# `or#` u#)))
  95     in
  96     mk_supply#  `thenPrimIO` \ s ->
  97     return s
  98
  99 splitUniqSupply (MkSplitUniqSupply _ s1 s2) = (s1, s2)
 100 \end{code}
 101
 102 \begin{code}
 103 getUnique (MkSplitUniqSupply (I# n) _ _) = mkUniqueGrimily n
 104
 105 getUniques (I# i) supply = i `get_from` supply
 106   where
 107     get_from 0# _ = []
 108     get_from n (MkSplitUniqSupply (I# u) _ s2)
 109       = mkUniqueGrimily u : get_from (n `minusInt#` 1#) s2
 110 \end{code}
 111
 112 %************************************************************************
 113 %*                                                                      *
 114 \subsubsection[UniqSupply-monad]{@UniqSupply@ monad: @UniqSM@}
 115 %*                                                                      *
 116 %************************************************************************
 117
 118 \begin{code}
 119 type UniqSM result = UniqSupply -> result
 120
 121 -- the initUs function also returns the final UniqSupply
 122
 123 initUs :: UniqSupply -> UniqSM a -> (UniqSupply, a)
 124
 125 initUs init_us m
 126   = case (splitUniqSupply init_us) of { (s1, s2) ->
 127     (s2, m s1) }
 128
 129 {-# INLINE thenUs #-}
 130 {-# INLINE returnUs #-}
 131 {-# INLINE splitUniqSupply #-}
 132 \end{code}
 133
 134 @thenUs@ is where we split the @UniqSupply@.
 135 \begin{code}
 136 thenUs :: UniqSM a -> (a -> UniqSM b) -> UniqSM b
 137
 138 thenUs expr cont us
 139   = case (splitUniqSupply us) of { (s1, s2) ->
 140     case (expr s1)            of { result ->
 141     cont result s2 }}
 142 \end{code}
 143
 144 \begin{code}
 145 returnUs :: a -> UniqSM a
 146 returnUs result us = result
 147
 148 mapUs :: (a -> UniqSM b) -> [a] -> UniqSM [b]
 149
 150 mapUs f []     = returnUs []
 151 mapUs f (x:xs)
 152   = f x         `thenUs` \ r  ->
 153     mapUs f xs  `thenUs` \ rs ->
 154     returnUs (r:rs)
 155
 156 mapAndUnzipUs  :: (a -> UniqSM (b,c))   -> [a] -> UniqSM ([b],[c])
 157 mapAndUnzip3Us :: (a -> UniqSM (b,c,d)) -> [a] -> UniqSM ([b],[c],[d])
 158
 159 mapAndUnzipUs f [] = returnUs ([],[])
 160 mapAndUnzipUs f (x:xs)
 161   = f x                 `thenUs` \ (r1,  r2)  ->
 162     mapAndUnzipUs f xs  `thenUs` \ (rs1, rs2) ->
 163     returnUs (r1:rs1, r2:rs2)
 164
 165 mapAndUnzip3Us f [] = returnUs ([],[],[])
 166 mapAndUnzip3Us f (x:xs)
 167   = f x                 `thenUs` \ (r1,  r2,  r3)  ->
 168     mapAndUnzip3Us f xs `thenUs` \ (rs1, rs2, rs3) ->
 169     returnUs (r1:rs1, r2:rs2, r3:rs3)
 170
 171 thenMaybeUs :: UniqSM (Maybe a) -> (a -> UniqSM (Maybe b)) -> UniqSM (Maybe b)
 172 thenMaybeUs m k
 173   = m   `thenUs` \ result ->
 174     case result of
 175       Nothing -> returnUs Nothing
 176       Just x  -> k x
 177
 178 mapAccumLUs :: (acc -> x -> UniqSM (acc, y))
 179             -> acc
 180             -> [x]
 181             -> UniqSM (acc, [y])
 182
 183 mapAccumLUs f b []     = returnUs (b, [])
 184 mapAccumLUs f b (x:xs)
 185   = f b x                           `thenUs` \ (b__2, x__2) ->
 186     mapAccumLUs f b__2 xs           `thenUs` \ (b__3, xs__2) ->
 187     returnUs (b__3, x__2:xs__2)
 188 \end{code}