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         SYN_IE(UniqSM),         -- type: unique supply monad
  16         initUs, thenUs, returnUs, fixUs,
  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
  30 #if __GLASGOW_HASKELL__ == 201
  31 import PreludeGlaST
  32 # define WHASH      GHCbase.W#
  33 #elif __GLASGOW_HASKELL__ >= 202
  34 import GlaExts
  35 import STBase
  36 # if __GLASGOW_HASKELL__ == 202
  37 import PrelBase ( Char(..) )
  38 # endif
  39 # define WHASH      GlaExts.W#
  40 #else
  41 import PreludeGlaST
  42 # define WHASH      W#
  43 #endif
  44
  45 w2i x = word2Int# x
  46 i2w x = int2Word# x
  47 i2w_s x = (x :: Int#)
  48 \end{code}
  49
  50
  51 %************************************************************************
  52 %*                                                                      *
  53 \subsection{Splittable Unique supply: @UniqSupply@}
  54 %*                                                                      *
  55 %************************************************************************
  56
  57 %************************************************************************
  58 %*                                                                      *
  59 \subsubsection[UniqSupply-type]{@UniqSupply@ type and operations}
  60 %*                                                                      *
  61 %************************************************************************
  62
  63 A value of type @UniqSupply@ is unique, and it can
  64 supply {\em one} distinct @Unique@.  Also, from the supply, one can
  65 also manufacture an arbitrary number of further @UniqueSupplies@,
  66 which will be distinct from the first and from all others.
  67
  68 \begin{code}
  69 data UniqSupply
  70   = MkSplitUniqSupply Int       -- make the Unique with this
  71                    UniqSupply UniqSupply
  72                                 -- when split => these two supplies
  73 \end{code}
  74
  75 \begin{code}
  76 mkSplitUniqSupply :: Char -> IO UniqSupply
  77
  78 splitUniqSupply :: UniqSupply -> (UniqSupply, UniqSupply)
  79 getUnique :: UniqSupply -> Unique
  80 getUniques :: Int -> UniqSupply -> [Unique]
  81 \end{code}
  82
  83 \begin{code}
  84 mkSplitUniqSupply (C# c#)
  85   = let
  86         mask# = (i2w (ord# c#)) `shiftL#` (i2w_s 24#)
  87
  88         -- here comes THE MAGIC:
  89
  90         mk_supply#
  91           = unsafe_interleave (
  92                 mk_unique   `thenPrimIO` \ uniq ->
  93                 mk_supply#  `thenPrimIO` \ s1 ->
  94                 mk_supply#  `thenPrimIO` \ s2 ->
  95                 returnPrimIO (MkSplitUniqSupply uniq s1 s2)
  96             )
  97           where
  98 --
  99             -- inlined copy of unsafeInterleavePrimIO;
 100             -- this is the single-most-hammered bit of code
 101             -- in the compiler....
 102             -- Too bad it's not 1.3-portable...
 103             unsafe_interleave m =
 104                MkST ( \ s ->
 105                 let
 106                     (MkST m') = m
 107                     (r, new_s) = m' s
 108                 in
 109                 (r, s))
 110 --
 111
 112         mk_unique = _ccall_ genSymZh            `thenPrimIO` \ (WHASH u#) ->
 113                     returnPrimIO (I# (w2i (mask# `or#` u#)))
 114     in
 115 #if __GLASGOW_HASKELL__ >= 200
 116     primIOToIO mk_supply#       >>= \ s ->
 117     return s
 118 #else
 119     mk_supply#  `thenPrimIO` \ s ->
 120     return s
 121 #endif
 122
 123 splitUniqSupply (MkSplitUniqSupply _ s1 s2) = (s1, s2)
 124 \end{code}
 125
 126 \begin{code}
 127 getUnique (MkSplitUniqSupply (I# n) _ _) = mkUniqueGrimily n
 128
 129 getUniques (I# i) supply = i `get_from` supply
 130   where
 131     get_from 0# _ = []
 132     get_from n (MkSplitUniqSupply (I# u) _ s2)
 133       = mkUniqueGrimily u : get_from (n -# 1#) s2
 134 \end{code}
 135
 136 %************************************************************************
 137 %*                                                                      *
 138 \subsubsection[UniqSupply-monad]{@UniqSupply@ monad: @UniqSM@}
 139 %*                                                                      *
 140 %************************************************************************
 141
 142 \begin{code}
 143 type UniqSM result = UniqSupply -> result
 144
 145 -- the initUs function also returns the final UniqSupply
 146
 147 initUs :: UniqSupply -> UniqSM a -> a
 148
 149 initUs init_us m = m init_us
 150
 151 {-# INLINE thenUs #-}
 152 {-# INLINE returnUs #-}
 153 {-# INLINE splitUniqSupply #-}
 154 \end{code}
 155
 156 @thenUs@ is where we split the @UniqSupply@.
 157 \begin{code}
 158 fixUs :: (a -> UniqSM a) -> UniqSM a
 159 fixUs m us
 160   = r  where  r = m r us
 161
 162 thenUs :: UniqSM a -> (a -> UniqSM b) -> UniqSM b
 163
 164 thenUs expr cont us
 165   = case (splitUniqSupply us) of { (s1, s2) ->
 166     case (expr s1)            of { result ->
 167     cont result s2 }}
 168 \end{code}
 169
 170 \begin{code}
 171 returnUs :: a -> UniqSM a
 172 returnUs result us = result
 173
 174 mapUs :: (a -> UniqSM b) -> [a] -> UniqSM [b]
 175
 176 mapUs f []     = returnUs []
 177 mapUs f (x:xs)
 178   = f x         `thenUs` \ r  ->
 179     mapUs f xs  `thenUs` \ rs ->
 180     returnUs (r:rs)
 181
 182 mapAndUnzipUs  :: (a -> UniqSM (b,c))   -> [a] -> UniqSM ([b],[c])
 183 mapAndUnzip3Us :: (a -> UniqSM (b,c,d)) -> [a] -> UniqSM ([b],[c],[d])
 184
 185 mapAndUnzipUs f [] = returnUs ([],[])
 186 mapAndUnzipUs f (x:xs)
 187   = f x                 `thenUs` \ (r1,  r2)  ->
 188     mapAndUnzipUs f xs  `thenUs` \ (rs1, rs2) ->
 189     returnUs (r1:rs1, r2:rs2)
 190
 191 mapAndUnzip3Us f [] = returnUs ([],[],[])
 192 mapAndUnzip3Us f (x:xs)
 193   = f x                 `thenUs` \ (r1,  r2,  r3)  ->
 194     mapAndUnzip3Us f xs `thenUs` \ (rs1, rs2, rs3) ->
 195     returnUs (r1:rs1, r2:rs2, r3:rs3)
 196
 197 thenMaybeUs :: UniqSM (Maybe a) -> (a -> UniqSM (Maybe b)) -> UniqSM (Maybe b)
 198 thenMaybeUs m k
 199   = m   `thenUs` \ result ->
 200     case result of
 201       Nothing -> returnUs Nothing
 202       Just x  -> k x
 203
 204 mapAccumLUs :: (acc -> x -> UniqSM (acc, y))
 205             -> acc
 206             -> [x]
 207             -> UniqSM (acc, [y])
 208
 209 mapAccumLUs f b []     = returnUs (b, [])
 210 mapAccumLUs f b (x:xs)
 211   = f b x                           `thenUs` \ (b__2, x__2) ->
 212     mapAccumLUs f b__2 xs           `thenUs` \ (b__3, xs__2) ->
 213     returnUs (b__3, x__2:xs__2)
 214 \end{code}