79a7c7cef68b8492116e5be2741cd6d6e8daeaf5
[packages/hoopl.git] / src / Compiler / Hoopl / Unique.hs
1 {-# LANGUAGE CPP, TypeFamilies #-}
2 #if __GLASGOW_HASKELL__ >= 709
3 {-# LANGUAGE Safe #-}
4 #elif __GLASGOW_HASKELL__ >= 701
5 {-# LANGUAGE Trustworthy #-}
6 #endif
7
8
9 module Compiler.Hoopl.Unique
10 ( Unique, intToUnique
11 , UniqueSet, UniqueMap
12 , UniqueMonad(..)
13 , SimpleUniqueMonad, runSimpleUniqueMonad
14 , UniqueMonadT, runUniqueMonadT
15
16 , uniqueToInt -- exposed through GHC module only!
17 )
18
19 where
20
21 import Compiler.Hoopl.Checkpoint
22 import Compiler.Hoopl.Collections
23
24 import qualified Data.IntMap as M
25 import qualified Data.IntSet as S
26
27 import Control.Applicative as AP
28 import Control.Monad (ap,liftM)
29
30 -----------------------------------------------------------------------------
31 -- Unique
32 -----------------------------------------------------------------------------
33
34 type Unique = Int
35
36 uniqueToInt :: Unique -> Int
37 uniqueToInt = id
38
39 intToUnique :: Int -> Unique
40 intToUnique = id
41
42 -----------------------------------------------------------------------------
43 -- UniqueSet
44
45 newtype UniqueSet = US S.IntSet deriving (Eq, Ord, Show)
46
47 instance IsSet UniqueSet where
48 type ElemOf UniqueSet = Unique
49
50 setNull (US s) = S.null s
51 setSize (US s) = S.size s
52 setMember k (US s) = S.member k s
53
54 setEmpty = US S.empty
55 setSingleton k = US (S.singleton k)
56 setInsert k (US s) = US (S.insert k s)
57 setDelete k (US s) = US (S.delete k s)
58
59 setUnion (US x) (US y) = US (S.union x y)
60 setDifference (US x) (US y) = US (S.difference x y)
61 setIntersection (US x) (US y) = US (S.intersection x y)
62 setIsSubsetOf (US x) (US y) = S.isSubsetOf x y
63
64 setFold k z (US s) = S.fold k z s
65
66 setElems (US s) = S.elems s
67 setFromList ks = US (S.fromList ks)
68
69 -----------------------------------------------------------------------------
70 -- UniqueMap
71
72 newtype UniqueMap v = UM (M.IntMap v) deriving (Eq, Ord, Show)
73
74 instance IsMap UniqueMap where
75 type KeyOf UniqueMap = Unique
76
77 mapNull (UM m) = M.null m
78 mapSize (UM m) = M.size m
79 mapMember k (UM m) = M.member k m
80 mapLookup k (UM m) = M.lookup k m
81 mapFindWithDefault def k (UM m) = M.findWithDefault def k m
82
83 mapEmpty = UM M.empty
84 mapSingleton k v = UM (M.singleton k v)
85 mapInsert k v (UM m) = UM (M.insert k v m)
86 mapInsertWith f k v (UM m) = UM (M.insertWith f k v m)
87 mapDelete k (UM m) = UM (M.delete k m)
88
89 mapUnion (UM x) (UM y) = UM (M.union x y)
90 mapUnionWithKey f (UM x) (UM y) = UM (M.unionWithKey (f . intToUnique) x y)
91 mapDifference (UM x) (UM y) = UM (M.difference x y)
92 mapIntersection (UM x) (UM y) = UM (M.intersection x y)
93 mapIsSubmapOf (UM x) (UM y) = M.isSubmapOf x y
94
95 mapMap f (UM m) = UM (M.map f m)
96 mapMapWithKey f (UM m) = UM (M.mapWithKey (f . intToUnique) m)
97 mapFold k z (UM m) = M.fold k z m
98 mapFoldWithKey k z (UM m) = M.foldWithKey (k . intToUnique) z m
99 mapFilter f (UM m) = UM (M.filter f m)
100
101 mapElems (UM m) = M.elems m
102 mapKeys (UM m) = M.keys m
103 mapToList (UM m) = M.toList m
104 mapFromList assocs = UM (M.fromList assocs)
105 mapFromListWith f assocs = UM (M.fromListWith f assocs)
106
107 ----------------------------------------------------------------
108 -- Monads
109
110 class Monad m => UniqueMonad m where
111 freshUnique :: m Unique
112
113 newtype SimpleUniqueMonad a = SUM { unSUM :: [Unique] -> (a, [Unique]) }
114
115 instance Functor SimpleUniqueMonad where
116 fmap = liftM
117
118 instance Applicative SimpleUniqueMonad where
119 pure a = SUM $ \us -> (a, us)
120 (<*>) = ap
121
122 instance Monad SimpleUniqueMonad where
123 return = AP.pure
124 m >>= k = SUM $ \us -> let (a, us') = unSUM m us in
125 unSUM (k a) us'
126
127 instance UniqueMonad SimpleUniqueMonad where
128 freshUnique = SUM $ f
129 where f (u:us) = (u, us)
130 f _ = error "Unique.freshUnique(SimpleUniqueMonad): empty list"
131
132 instance CheckpointMonad SimpleUniqueMonad where
133 type Checkpoint SimpleUniqueMonad = [Unique]
134 checkpoint = SUM $ \us -> (us, us)
135 restart us = SUM $ \_ -> ((), us)
136
137 runSimpleUniqueMonad :: SimpleUniqueMonad a -> a
138 runSimpleUniqueMonad m = fst (unSUM m allUniques)
139
140 ----------------------------------------------------------------
141
142 newtype UniqueMonadT m a = UMT { unUMT :: [Unique] -> m (a, [Unique]) }
143
144 instance Monad m => Functor (UniqueMonadT m) where
145 fmap = liftM
146
147 instance Monad m => Applicative (UniqueMonadT m) where
148 pure a = UMT $ \us -> return (a, us)
149 (<*>) = ap
150
151 instance Monad m => Monad (UniqueMonadT m) where
152 return = pure
153 m >>= k = UMT $ \us -> do { (a, us') <- unUMT m us; unUMT (k a) us' }
154
155 instance Monad m => UniqueMonad (UniqueMonadT m) where
156 freshUnique = UMT $ f
157 where f (u:us) = return (u, us)
158 f _ = error "Unique.freshUnique(UniqueMonadT): empty list"
159
160 runUniqueMonadT :: Monad m => UniqueMonadT m a -> m a
161 runUniqueMonadT m = do { (a, _) <- unUMT m allUniques; return a }
162
163 allUniques :: [Unique]
164 allUniques = [1..]