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