UniqSupply.lhs 5.34 KB
Newer Older
1
%
2
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3
4
5
6
7
8
9
10
%
\section[UniqSupply]{The @UniqueSupply@ data type and a (monadic) supply thereof}

\begin{code}
module UniqSupply (

	UniqSupply,		-- Abstractly

11
	uniqFromSupply, uniqsFromSupply,	-- basic ops
12

13
	UniqSM,		-- type: unique supply monad
14
	initUs, initUs_, thenUs, thenUs_, returnUs, fixUs, getUs, withUs,
15
	getUniqueUs, getUniquesUs,
16
	mapUs, mapAndUnzipUs, mapAndUnzip3Us,
17
	thenMaybeUs, mapAccumLUs,
18
	lazyThenUs, lazyMapUs,
19
20

	mkSplitUniqSupply,
21
	splitUniqSupply
22
23
  ) where

24
#include "HsVersions.h"
25
26

import Unique
27

28
29
import GLAEXTS
import UNSAFE_IO	( unsafeInterleaveIO )
sof's avatar
sof committed
30

31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
w2i x = word2Int# x
i2w x = int2Word# x
i2w_s x = (x :: Int#)
\end{code}


%************************************************************************
%*									*
\subsection{Splittable Unique supply: @UniqSupply@}
%*									*
%************************************************************************

A value of type @UniqSupply@ is unique, and it can
supply {\em one} distinct @Unique@.  Also, from the supply, one can
also manufacture an arbitrary number of further @UniqueSupplies@,
which will be distinct from the first and from all others.

\begin{code}
data UniqSupply
50
  = MkSplitUniqSupply Int#	-- make the Unique with this
51
52
53
54
55
		   UniqSupply UniqSupply
				-- when split => these two supplies
\end{code}

\begin{code}
56
mkSplitUniqSupply :: Char -> IO UniqSupply
57
58

splitUniqSupply :: UniqSupply -> (UniqSupply, UniqSupply)
59
uniqFromSupply  :: UniqSupply -> Unique
60
uniqsFromSupply :: UniqSupply -> [Unique]	-- Infinite
61
62
63
\end{code}

\begin{code}
64
mkSplitUniqSupply (C# c#)
65
  = let
66
67
68
#if __GLASGOW_HASKELL__ >= 503
	mask# = (i2w (ord# c#)) `uncheckedShiftL#` (i2w_s 24#)
#else
69
	mask# = (i2w (ord# c#)) `shiftL#` (i2w_s 24#)
70
#endif
71
72
	-- here comes THE MAGIC:

73
	-- This is one of the most hammered bits in the whole compiler
74
	mk_supply#
75
	  = unsafeInterleaveIO (
76
		genSymZh    >>= \ (W# u#) ->
77
78
		mk_supply#  >>= \ s1 ->
		mk_supply#  >>= \ s2 ->
79
		return (MkSplitUniqSupply (w2i (mask# `or#` u#)) s1 s2)
80
81
	    )
    in
82
    mk_supply#
83

84
foreign import ccall unsafe "genSymZh" genSymZh :: IO Word
85

86
87
88
89
splitUniqSupply (MkSplitUniqSupply _ s1 s2) = (s1, s2)
\end{code}

\begin{code}
90
91
uniqFromSupply  (MkSplitUniqSupply n _ _)  = mkUniqueGrimily (I# n)
uniqsFromSupply (MkSplitUniqSupply n _ s2) = mkUniqueGrimily (I# n) : uniqsFromSupply s2
92
93
94
95
96
97
98
99
100
\end{code}

%************************************************************************
%*									*
\subsubsection[UniqSupply-monad]{@UniqSupply@ monad: @UniqSM@}
%*									*
%************************************************************************

\begin{code}
101
type UniqSM result = UniqSupply -> (result, UniqSupply)
102

103
104
105
-- the initUs function also returns the final UniqSupply; initUs_ drops it
initUs :: UniqSupply -> UniqSM a -> (a,UniqSupply)
initUs init_us m = case m init_us of { (r,us) -> (r,us) }
106

107
108
initUs_ :: UniqSupply -> UniqSM a -> a
initUs_ init_us m = case m init_us of { (r,us) -> r }
109
110

{-# INLINE thenUs #-}
111
{-# INLINE lazyThenUs #-}
112
113
114
115
116
117
{-# INLINE returnUs #-}
{-# INLINE splitUniqSupply #-}
\end{code}

@thenUs@ is where we split the @UniqSupply@.
\begin{code}
118
119
fixUs :: (a -> UniqSM a) -> UniqSM a
fixUs m us
120
  = (r,us')  where  (r,us') = m r us
121

122
123
thenUs :: UniqSM a -> (a -> UniqSM b) -> UniqSM b
thenUs expr cont us
124
125
  = case (expr us) of { (result, us') -> cont result us' }

126
127
128
129
lazyThenUs :: UniqSM a -> (a -> UniqSM b) -> UniqSM b
lazyThenUs expr cont us
  = let (result, us') = expr us in cont result us'

130
131
132
thenUs_ :: UniqSM a -> UniqSM b -> UniqSM b
thenUs_ expr cont us
  = case (expr us) of { (_, us') -> cont us' }
133

134

135
returnUs :: a -> UniqSM a
136
137
returnUs result us = (result, us)

138
139
140
withUs :: (UniqSupply -> (a, UniqSupply)) -> UniqSM a
withUs f us = f us	-- Ha ha!
		
141
getUs :: UniqSM UniqSupply
142
getUs us = splitUniqSupply us
143
144
145
146
147

getUniqueUs :: UniqSM Unique
getUniqueUs us = case splitUniqSupply us of
		   (us1,us2) -> (uniqFromSupply us1, us2)

148
149
150
getUniquesUs :: UniqSM [Unique]
getUniquesUs us = case splitUniqSupply us of
		      (us1,us2) -> (uniqsFromSupply us1, us2)
151
152
153
\end{code}

\begin{code}
154
155
156
157
158
159
160
mapUs :: (a -> UniqSM b) -> [a] -> UniqSM [b]
mapUs f []     = returnUs []
mapUs f (x:xs)
  = f x         `thenUs` \ r  ->
    mapUs f xs  `thenUs` \ rs ->
    returnUs (r:rs)

161
162
163
164
165
166
167
lazyMapUs :: (a -> UniqSM b) -> [a] -> UniqSM [b]
lazyMapUs f []     = returnUs []
lazyMapUs f (x:xs)
  = f x             `lazyThenUs` \ r  ->
    lazyMapUs f xs  `lazyThenUs` \ rs ->
    returnUs (r:rs)

168
mapAndUnzipUs  :: (a -> UniqSM (b,c))   -> [a] -> UniqSM ([b],[c])
169
mapAndUnzip3Us :: (a -> UniqSM (b,c,d)) -> [a] -> UniqSM ([b],[c],[d])
170
171
172
173
174
175

mapAndUnzipUs f [] = returnUs ([],[])
mapAndUnzipUs f (x:xs)
  = f x		    	`thenUs` \ (r1,  r2)  ->
    mapAndUnzipUs f xs	`thenUs` \ (rs1, rs2) ->
    returnUs (r1:rs1, r2:rs2)
176
177
178
179
180
181

mapAndUnzip3Us f [] = returnUs ([],[],[])
mapAndUnzip3Us f (x:xs)
  = f x		    	`thenUs` \ (r1,  r2,  r3)  ->
    mapAndUnzip3Us f xs	`thenUs` \ (rs1, rs2, rs3) ->
    returnUs (r1:rs1, r2:rs2, r3:rs3)
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199

thenMaybeUs :: UniqSM (Maybe a) -> (a -> UniqSM (Maybe b)) -> UniqSM (Maybe b)
thenMaybeUs m k
  = m	`thenUs` \ result ->
    case result of
      Nothing -> returnUs Nothing
      Just x  -> k x

mapAccumLUs :: (acc -> x -> UniqSM (acc, y))
	    -> acc
	    -> [x]
	    -> UniqSM (acc, [y])

mapAccumLUs f b []     = returnUs (b, [])
mapAccumLUs f b (x:xs)
  = f b x   	        	    `thenUs` \ (b__2, x__2) ->
    mapAccumLUs f b__2 xs   	    `thenUs` \ (b__3, xs__2) ->
    returnUs (b__3, x__2:xs__2)
200
\end{code}