1 -- this is basically from ghc-2.01 2 -- with some specialisations added 3 4 5 6 7 8 9 10 11 12 13 14 15 16 module FiniteMap ( 17 FiniteMap, -- abstract type 18 19 emptyFM, unitFM, listToFM, 20 21 addToFM, 22 addToFM_C, 23 addListToFM, 24 addListToFM_C, 25 delFromFM , 26 delListFromFM, 27 28 plusFM, 29 plusFM_C, 30 minusFM, 31 foldFM, 32 33 intersectFM , 34 intersectFM_C , 35 mapFM , filterFM , 36 37 sizeFM, isEmptyFM, elemFM, lookupFM, lookupWithDefaultFM, 38 39 fmToList, keysFM, eltsFM 40 41 42 ) where 43 44 import Maybes 45 46 ------------------------------------------ 47 -- import TA 48 -- import Set 49 50 {- # SPECIALIZE lookupFM :: FiniteMap Int (Set (STerm Int)) -> Int -> Maybe (Set (STerm Int)) #-} 51 {- # SPECIALIZE lookupFM :: FiniteMap (STerm Int) (Set Int) -> (STerm Int) -> Maybe (Set Int) #-} 52 53 {- # SPECIALIZE lookupWithDefaultFM :: FiniteMap Int (Set (STerm Int)) -> Set (STerm Int) -> Int -> (Set (STerm Int)) #-} 54 {- # SPECIALIZE lookupWithDefaultFM :: FiniteMap (STerm Int) (Set Int) -> (Set Int) -> (STerm Int) -> (Set Int) #-} 55 56 57 58 59 60 61 -- SIGH: but we use unboxed "sizes"... 62 63 64 65 66 -- BUILDING 67 emptyFM :: FiniteMap key elt 68 unitFM :: key -> elt -> FiniteMap key elt 69 listToFM :: (Ord key {--}) => [(key,elt)] -> FiniteMap key elt 70 -- In the case of duplicates, the last is taken 71 72 73 -- ADDING AND DELETING 74 -- Throws away any previous binding 75 -- In the list case, the items are added starting with the 76 -- first one in the list 77 addToFM :: (Ord key {--}) => FiniteMap key elt -> key -> elt -> FiniteMap key elt 78 addListToFM :: (Ord key {--}) => FiniteMap key elt -> [(key,elt)] -> FiniteMap key elt 79 80 -- Combines with previous binding 81 addToFM_C :: (Ord key {--}) => (elt -> elt -> elt) 82 -> FiniteMap key elt -> key -> elt 83 -> FiniteMap key elt 84 addListToFM_C :: (Ord key {--}) => (elt -> elt -> elt) 85 -> FiniteMap key elt -> [(key,elt)] 86 -> FiniteMap key elt 87 88 -- Deletion doesn't complain if you try to delete something 89 -- which isn't there 90 delFromFM :: (Ord key {--}) => FiniteMap key elt -> key -> FiniteMap key elt 91 delListFromFM :: (Ord key {--}) => FiniteMap key elt -> [key] -> FiniteMap key elt 92 93 -- COMBINING 94 -- Bindings in right argument shadow those in the left 95 plusFM :: (Ord key {--}) => FiniteMap key elt -> FiniteMap key elt 96 -> FiniteMap key elt 97 98 -- Combines bindings for the same thing with the given function 99 plusFM_C :: (Ord key {--}) => (elt -> elt -> elt) 100 -> FiniteMap key elt -> FiniteMap key elt -> FiniteMap key elt 101 102 minusFM :: (Ord key {--}) => FiniteMap key elt -> FiniteMap key elt -> FiniteMap key elt 103 -- (minusFM a1 a2) deletes from a1 any bindings which are bound in a2 104 105 intersectFM :: (Ord key {--}) => FiniteMap key elt -> FiniteMap key elt -> FiniteMap key elt 106 intersectFM_C :: (Ord key {--}) => (elt -> elt -> elt) 107 -> FiniteMap key elt -> FiniteMap key elt -> FiniteMap key elt 108 109 -- MAPPING, FOLDING, FILTERING 110 foldFM :: (key -> elt -> a -> a) -> a -> FiniteMap key elt -> a 111 mapFM :: (key -> elt1 -> elt2) -> FiniteMap key elt1 -> FiniteMap key elt2 112 filterFM :: (Ord key {--}) => (key -> elt -> Bool) 113 -> FiniteMap key elt -> FiniteMap key elt 114 115 -- INTERROGATING 116 sizeFM :: FiniteMap key elt -> Int 117 isEmptyFM :: FiniteMap key elt -> Bool 118 119 elemFM :: (Ord key {--}) => key -> FiniteMap key elt -> Bool 120 lookupFM :: (Ord key {--}) => FiniteMap key elt -> key -> Maybe elt 121 lookupWithDefaultFM 122 :: (Ord key {--}) => FiniteMap key elt -> elt -> key -> elt 123 -- lookupWithDefaultFM supplies a "default" elt 124 -- to return for an unmapped key 125 126 -- LISTIFYING 127 fmToList :: FiniteMap key elt -> [(key,elt)] 128 keysFM :: FiniteMap key elt -> [key] 129 eltsFM :: FiniteMap key elt -> [elt] 130 131 data FiniteMap key elt 132 = EmptyFM 133 | Branch key elt -- Key and elt stored here 134 Int{-STRICT-} -- Size >= 1 135 (FiniteMap key elt) -- Children 136 (FiniteMap key elt) 137 138 emptyFM = EmptyFM 139 {- 140 emptyFM 141 = Branch bottom bottom 0 bottom bottom 142 where 143 bottom = panic "emptyFM" 144 -} 145 146 -- #define EmptyFM (Branch _ _ 0 _ _) 147 148 unitFM key elt = Branch key elt 1 emptyFM emptyFM 149 150 listToFM = addListToFM emptyFM 151 152 153 154 addToFM fm key elt = addToFM_C (\ old new -> new) fm key elt 155 156 addToFM_C combiner EmptyFM key elt = unitFM key elt 157 addToFM_C combiner (Branch key elt size fm_l fm_r) new_key new_elt 158 159 | new_key < key = mkBalBranch key elt (addToFM_C combiner fm_l new_key new_elt) fm_r 160 | new_key > key = mkBalBranch key elt fm_l (addToFM_C combiner fm_r new_key new_elt) 161 | otherwise = Branch new_key (combiner elt new_elt) size fm_l fm_r 162 163 164 addListToFM fm key_elt_pairs = addListToFM_C (\ old new -> new) fm key_elt_pairs 165 166 addListToFM_C combiner fm key_elt_pairs 167 = foldl add fm key_elt_pairs -- foldl adds from the left 168 where 169 add fmap (key,elt) = addToFM_C combiner fmap key elt 170 171 delFromFM EmptyFM del_key = emptyFM 172 delFromFM (Branch key elt size fm_l fm_r) del_key 173 174 | del_key > key 175 = mkBalBranch key elt fm_l (delFromFM fm_r del_key) 176 177 | del_key < key 178 = mkBalBranch key elt (delFromFM fm_l del_key) fm_r 179 180 | key == del_key 181 = glueBal fm_l fm_r 182 183 184 delListFromFM fm keys = foldl delFromFM fm keys 185 186 plusFM_C combiner EmptyFM fm2 = fm2 187 plusFM_C combiner fm1 EmptyFM = fm1 188 plusFM_C combiner fm1 (Branch split_key elt2 _ left right) 189 = mkVBalBranch split_key new_elt 190 (plusFM_C combiner lts left) 191 (plusFM_C combiner gts right) 192 where 193 lts = splitLT fm1 split_key 194 gts = splitGT fm1 split_key 195 new_elt = case lookupFM fm1 split_key of 196 Nothing -> elt2 197 Just elt1 -> combiner elt1 elt2 198 199 -- It's worth doing plusFM specially, because we don't need 200 -- to do the lookup in fm1. 201 202 plusFM EmptyFM fm2 = fm2 203 plusFM fm1 EmptyFM = fm1 204 plusFM fm1 (Branch split_key elt1 _ left right) 205 = mkVBalBranch split_key elt1 (plusFM lts left) (plusFM gts right) 206 where 207 lts = splitLT fm1 split_key 208 gts = splitGT fm1 split_key 209 210 minusFM EmptyFM fm2 = emptyFM 211 minusFM fm1 EmptyFM = fm1 212 minusFM fm1 (Branch split_key elt _ left right) 213 = glueVBal (minusFM lts left) (minusFM gts right) 214 -- The two can be way different, so we need glueVBal 215 where 216 lts = splitLT fm1 split_key -- NB gt and lt, so the equal ones 217 gts = splitGT fm1 split_key -- are not in either. 218 219 intersectFM fm1 fm2 = intersectFM_C (\ left right -> right) fm1 fm2 220 221 intersectFM_C combiner fm1 EmptyFM = emptyFM 222 intersectFM_C combiner EmptyFM fm2 = emptyFM 223 intersectFM_C combiner fm1 (Branch split_key elt2 _ left right) 224 225 | maybeToBool maybe_elt1 -- split_elt *is* in intersection 226 = mkVBalBranch split_key (combiner elt1 elt2) (intersectFM_C combiner lts left) 227 (intersectFM_C combiner gts right) 228 229 | otherwise -- split_elt is *not* in intersection 230 = glueVBal (intersectFM_C combiner lts left) (intersectFM_C combiner gts right) 231 232 where 233 lts = splitLT fm1 split_key -- NB gt and lt, so the equal ones 234 gts = splitGT fm1 split_key -- are not in either. 235 236 maybe_elt1 = lookupFM fm1 split_key 237 Just elt1 = maybe_elt1 238 239 foldFM k z EmptyFM = z 240 foldFM k z (Branch key elt _ fm_l fm_r) 241 = foldFM k (k key elt (foldFM k z fm_r)) fm_l 242 243 mapFM f EmptyFM = emptyFM 244 mapFM f (Branch key elt size fm_l fm_r) 245 = Branch key (f key elt) size (mapFM f fm_l) (mapFM f fm_r) 246 247 filterFM p EmptyFM = emptyFM 248 filterFM p (Branch key elt _ fm_l fm_r) 249 | p key elt -- Keep the item 250 = mkVBalBranch key elt (filterFM p fm_l) (filterFM p fm_r) 251 252 | otherwise -- Drop the item 253 = glueVBal (filterFM p fm_l) (filterFM p fm_r) 254 255 --{-# INLINE sizeFM #-} 256 sizeFM EmptyFM = 0 257 sizeFM (Branch _ _ size _ _) = size 258 259 isEmptyFM fm = sizeFM fm == 0 260 261 lookupFM EmptyFM key = Nothing 262 lookupFM (Branch key elt _ fm_l fm_r) key_to_find 263 264 | key_to_find < key = lookupFM fm_l key_to_find 265 | key_to_find > key = lookupFM fm_r key_to_find 266 | otherwise = Just elt 267 268 269 key `elemFM` fm 270 = case (lookupFM fm key) of { Nothing -> False; Just elt -> True } 271 272 lookupWithDefaultFM fm deflt key 273 = case (lookupFM fm key) of { Nothing -> deflt; Just elt -> elt } 274 275 fmToList fm = foldFM (\ key elt rest -> (key,elt) : rest) [] fm 276 keysFM fm = foldFM (\ key elt rest -> key : rest) [] fm 277 eltsFM fm = foldFM (\ key elt rest -> elt : rest) [] fm 278 279 sIZE_RATIO :: Int 280 sIZE_RATIO = 5 281 282 mkBranch :: (Ord key {--}) -- Used for the assertion checking only 283 => Int 284 -> key -> elt 285 -> FiniteMap key elt -> FiniteMap key elt 286 -> FiniteMap key elt 287 288 mkBranch which key elt fm_l fm_r 289 = --{--} 290 291 let 292 result = Branch key elt (unbox (1 + left_size + right_size)) fm_l fm_r 293 in 294 -- if sizeFM result <= 8 then 295 result 296 -- else 297 -- pprTrace ("mkBranch:"++(show which)) (ppr PprDebug result) ( 298 -- result 299 -- ) 300 where 301 left_ok = case fm_l of 302 EmptyFM -> True 303 Branch left_key _ _ _ _ -> let 304 biggest_left_key = fst (findMax fm_l) 305 in 306 biggest_left_key < key 307 right_ok = case fm_r of 308 EmptyFM -> True 309 Branch right_key _ _ _ _ -> let 310 smallest_right_key = fst (findMin fm_r) 311 in 312 key < smallest_right_key 313 balance_ok = True -- sigh 314 {- LATER: 315 balance_ok 316 = -- Both subtrees have one or no elements... 317 (left_size + right_size <= 1) 318 -- NO || left_size == 0 -- ??? 319 -- NO || right_size == 0 -- ??? 320 -- ... or the number of elements in a subtree does not exceed 321 -- sIZE_RATIO times the number of elements in the other subtree 322 || (left_size * sIZE_RATIO >= right_size && 323 right_size * sIZE_RATIO >= left_size) 324 -} 325 326 left_size = sizeFM fm_l 327 right_size = sizeFM fm_r 328 329 330 unbox :: Int -> Int 331 unbox x = x 332 333 334 mkBalBranch :: (Ord key {--}) 335 => key -> elt 336 -> FiniteMap key elt -> FiniteMap key elt 337 -> FiniteMap key elt 338 339 mkBalBranch key elt fm_L fm_R 340 341 | size_l + size_r < 2 342 = mkBranch 1{-which-} key elt fm_L fm_R 343 344 | size_r > sIZE_RATIO * size_l -- Right tree too big 345 = case fm_R of 346 Branch _ _ _ fm_rl fm_rr 347 | sizeFM fm_rl < 2 * sizeFM fm_rr -> single_L fm_L fm_R 348 | otherwise -> double_L fm_L fm_R 349 -- Other case impossible 350 351 | size_l > sIZE_RATIO * size_r -- Left tree too big 352 = case fm_L of 353 Branch _ _ _ fm_ll fm_lr 354 | sizeFM fm_lr < 2 * sizeFM fm_ll -> single_R fm_L fm_R 355 | otherwise -> double_R fm_L fm_R 356 -- Other case impossible 357 358 | otherwise -- No imbalance 359 = mkBranch 2{-which-} key elt fm_L fm_R 360 361 where 362 size_l = sizeFM fm_L 363 size_r = sizeFM fm_R 364 365 single_L fm_l (Branch key_r elt_r _ fm_rl fm_rr) 366 = mkBranch 3{-which-} key_r elt_r (mkBranch 4{-which-} key elt fm_l fm_rl) fm_rr 367 368 double_L fm_l (Branch key_r elt_r _ (Branch key_rl elt_rl _ fm_rll fm_rlr) fm_rr) 369 = mkBranch 5{-which-} key_rl elt_rl (mkBranch 6{-which-} key elt fm_l fm_rll) 370 (mkBranch 7{-which-} key_r elt_r fm_rlr fm_rr) 371 372 single_R (Branch key_l elt_l _ fm_ll fm_lr) fm_r 373 = mkBranch 8{-which-} key_l elt_l fm_ll (mkBranch 9{-which-} key elt fm_lr fm_r) 374 375 double_R (Branch key_l elt_l _ fm_ll (Branch key_lr elt_lr _ fm_lrl fm_lrr)) fm_r 376 = mkBranch 10{-which-} key_lr elt_lr (mkBranch 11{-which-} key_l elt_l fm_ll fm_lrl) 377 (mkBranch 12{-which-} key elt fm_lrr fm_r) 378 379 mkVBalBranch :: (Ord key {--}) 380 => key -> elt 381 -> FiniteMap key elt -> FiniteMap key elt 382 -> FiniteMap key elt 383 384 -- Assert: in any call to (mkVBalBranch_C comb key elt l r), 385 -- (a) all keys in l are < all keys in r 386 -- (b) all keys in l are < key 387 -- (c) all keys in r are > key 388 389 mkVBalBranch key elt EmptyFM fm_r = addToFM fm_r key elt 390 mkVBalBranch key elt fm_l EmptyFM = addToFM fm_l key elt 391 392 mkVBalBranch key elt fm_l@(Branch key_l elt_l _ fm_ll fm_lr) 393 fm_r@(Branch key_r elt_r _ fm_rl fm_rr) 394 | sIZE_RATIO * size_l < size_r 395 = mkBalBranch key_r elt_r (mkVBalBranch key elt fm_l fm_rl) fm_rr 396 397 | sIZE_RATIO * size_r < size_l 398 = mkBalBranch key_l elt_l fm_ll (mkVBalBranch key elt fm_lr fm_r) 399 400 | otherwise 401 = mkBranch 13{-which-} key elt fm_l fm_r 402 403 where 404 size_l = sizeFM fm_l 405 size_r = sizeFM fm_r 406 407 glueBal :: (Ord key {--}) 408 => FiniteMap key elt -> FiniteMap key elt 409 -> FiniteMap key elt 410 411 glueBal EmptyFM fm2 = fm2 412 glueBal fm1 EmptyFM = fm1 413 glueBal fm1 fm2 414 -- The case analysis here (absent in Adams' program) is really to deal 415 -- with the case where fm2 is a singleton. Then deleting the minimum means 416 -- we pass an empty tree to mkBalBranch, which breaks its invariant. 417 | sizeFM fm2 > sizeFM fm1 418 = mkBalBranch mid_key2 mid_elt2 fm1 (deleteMin fm2) 419 420 | otherwise 421 = mkBalBranch mid_key1 mid_elt1 (deleteMax fm1) fm2 422 where 423 (mid_key1, mid_elt1) = findMax fm1 424 (mid_key2, mid_elt2) = findMin fm2 425 426 glueVBal :: (Ord key {--}) 427 => FiniteMap key elt -> FiniteMap key elt 428 -> FiniteMap key elt 429 430 glueVBal EmptyFM fm2 = fm2 431 glueVBal fm1 EmptyFM = fm1 432 glueVBal fm_l@(Branch key_l elt_l _ fm_ll fm_lr) 433 fm_r@(Branch key_r elt_r _ fm_rl fm_rr) 434 | sIZE_RATIO * size_l < size_r 435 = mkBalBranch key_r elt_r (glueVBal fm_l fm_rl) fm_rr 436 437 | sIZE_RATIO * size_r < size_l 438 = mkBalBranch key_l elt_l fm_ll (glueVBal fm_lr fm_r) 439 440 | otherwise -- We now need the same two cases as in glueBal above. 441 = glueBal fm_l fm_r 442 where 443 (mid_key_l,mid_elt_l) = findMax fm_l 444 (mid_key_r,mid_elt_r) = findMin fm_r 445 size_l = sizeFM fm_l 446 size_r = sizeFM fm_r 447 448 splitLT, splitGT :: (Ord key {--}) => FiniteMap key elt -> key -> FiniteMap key elt 449 450 -- splitLT fm split_key = fm restricted to keys < split_key 451 -- splitGT fm split_key = fm restricted to keys > split_key 452 453 splitLT EmptyFM split_key = emptyFM 454 splitLT (Branch key elt _ fm_l fm_r) split_key 455 456 | split_key < key = splitLT fm_l split_key 457 | split_key > key = mkVBalBranch key elt fm_l (splitLT fm_r split_key) 458 | otherwise = fm_l 459 460 461 splitGT EmptyFM split_key = emptyFM 462 splitGT (Branch key elt _ fm_l fm_r) split_key 463 464 | split_key > key = splitGT fm_r split_key 465 | split_key < key = mkVBalBranch key elt (splitGT fm_l split_key) fm_r 466 | otherwise = fm_r 467 468 469 findMin :: FiniteMap key elt -> (key,elt) 470 findMin (Branch key elt _ EmptyFM _) = (key,elt) 471 findMin (Branch key elt _ fm_l _) = findMin fm_l 472 473 deleteMin :: (Ord key {--}) => FiniteMap key elt -> FiniteMap key elt 474 deleteMin (Branch key elt _ EmptyFM fm_r) = fm_r 475 deleteMin (Branch key elt _ fm_l fm_r) = mkBalBranch key elt (deleteMin fm_l) fm_r 476 477 findMax :: FiniteMap key elt -> (key,elt) 478 findMax (Branch key elt _ _ EmptyFM) = (key,elt) 479 findMax (Branch key elt _ _ fm_r) = findMax fm_r 480 481 deleteMax :: (Ord key {--}) => FiniteMap key elt -> FiniteMap key elt 482 deleteMax (Branch key elt _ fm_l EmptyFM) = fm_l 483 deleteMax (Branch key elt _ fm_l fm_r) = mkBalBranch key elt fm_l (deleteMax fm_r) 484 485 486 487 488 instance (Eq key, Eq elt) => Eq (FiniteMap key elt) where 489 fm_1 == fm_2 = (sizeFM fm_1 == sizeFM fm_2) && -- quick test 490 (fmToList fm_1 == fmToList fm_2) 491 492 {- NO: not clear what The Right Thing to do is: 493 instance (Ord key, Ord elt) => Ord (FiniteMap key elt) where 494 fm_1 <= fm_2 = (sizeFM fm_1 <= sizeFM fm_2) && -- quick test 495 (fmToList fm_1 <= fmToList fm_2) 496 -} 497 498 499 500 501 502 503 instance (Show a, Show b) => Show (FiniteMap a b) where 504 showsPrec p fm = 505 showsPrec p (fmToList fm)