Data/Map.hs

   1 {-# OPTIONS_GHC -fno-bang-patterns #-}
   2
   3 -----------------------------------------------------------------------------
   4 -- |
   5 -- Module      :  Data.Map
   6 -- Copyright   :  (c) Daan Leijen 2002
   7 -- License     :  BSD-style
   8 -- Maintainer  :  libraries@haskell.org
   9 -- Stability   :  provisional
  10 -- Portability :  portable
  11 --
  12 -- An efficient implementation of maps from keys to values (dictionaries).
  13 --
  14 -- Since many function names (but not the type name) clash with
  15 -- "Prelude" names, this module is usually imported @qualified@, e.g.
  16 --
  17 -- >  import Data.Map (Map)
  18 -- >  import qualified Data.Map as Map
  19 --
  20 -- The implementation of 'Map' is based on /size balanced/ binary trees (or
  21 -- trees of /bounded balance/) as described by:
  22 --
  23 --    * Stephen Adams, \"/Efficient sets: a balancing act/\",
  24 --      Journal of Functional Programming 3(4):553-562, October 1993,
  25 --      <http://www.swiss.ai.mit.edu/~adams/BB>.
  26 --
  27 --    * J. Nievergelt and E.M. Reingold,
  28 --      \"/Binary search trees of bounded balance/\",
  29 --      SIAM journal of computing 2(1), March 1973.
  30 --
  31 -- Note that the implementation is /left-biased/ -- the elements of a
  32 -- first argument are always preferred to the second, for example in
  33 -- 'union' or 'insert'.
  34 -----------------------------------------------------------------------------
  35
  36 module Data.Map  (
  37             -- * Map type
  38               Map          -- instance Eq,Show,Read
  39
  40             -- * Operators
  41             , (!), (\\)
  42
  43
  44             -- * Query
  45             , null
  46             , size
  47             , member
  48             , notMember
  49             , lookup
  50             , findWithDefault
  51
  52             -- * Construction
  53             , empty
  54             , singleton
  55
  56             -- ** Insertion
  57             , insert
  58             , insertWith, insertWithKey, insertLookupWithKey
  59
  60             -- ** Delete\/Update
  61             , delete
  62             , adjust
  63             , adjustWithKey
  64             , update
  65             , updateWithKey
  66             , updateLookupWithKey
  67             , alter
  68
  69             -- * Combine
  70
  71             -- ** Union
  72             , union
  73             , unionWith
  74             , unionWithKey
  75             , unions
  76             , unionsWith
  77
  78             -- ** Difference
  79             , difference
  80             , differenceWith
  81             , differenceWithKey
  82
  83             -- ** Intersection
  84             , intersection
  85             , intersectionWith
  86             , intersectionWithKey
  87
  88             -- * Traversal
  89             -- ** Map
  90             , map
  91             , mapWithKey
  92             , mapAccum
  93             , mapAccumWithKey
  94             , mapKeys
  95             , mapKeysWith
  96             , mapKeysMonotonic
  97
  98             -- ** Fold
  99             , fold
 100             , foldWithKey
 101
 102             -- * Conversion
 103             , elems
 104             , keys
 105             , keysSet
 106             , assocs
 107
 108             -- ** Lists
 109             , toList
 110             , fromList
 111             , fromListWith
 112             , fromListWithKey
 113
 114             -- ** Ordered lists
 115             , toAscList
 116             , fromAscList
 117             , fromAscListWith
 118             , fromAscListWithKey
 119             , fromDistinctAscList
 120
 121             -- * Filter
 122             , filter
 123             , filterWithKey
 124             , partition
 125             , partitionWithKey
 126
 127             , mapMaybe
 128             , mapMaybeWithKey
 129             , mapEither
 130             , mapEitherWithKey
 131
 132             , split
 133             , splitLookup
 134
 135             -- * Submap
 136             , isSubmapOf, isSubmapOfBy
 137             , isProperSubmapOf, isProperSubmapOfBy
 138
 139             -- * Indexed
 140             , lookupIndex
 141             , findIndex
 142             , elemAt
 143             , updateAt
 144             , deleteAt
 145
 146             -- * Min\/Max
 147             , findMin
 148             , findMax
 149             , deleteMin
 150             , deleteMax
 151             , deleteFindMin
 152             , deleteFindMax
 153             , updateMin
 154             , updateMax
 155             , updateMinWithKey
 156             , updateMaxWithKey
 157             , minView
 158             , maxView
 159
 160             -- * Debugging
 161             , showTree
 162             , showTreeWith
 163             , valid
 164             ) where
 165
 166 import Prelude hiding (lookup,map,filter,foldr,foldl,null)
 167 import qualified Data.Set as Set
 168 import qualified Data.List as List
 169 import Data.Monoid (Monoid(..))
 170 import Data.Typeable
 171 import Control.Applicative (Applicative(..), (<$>))
 172 import Data.Traversable (Traversable(traverse))
 173 import Data.Foldable (Foldable(foldMap))
 174
 175 {-
 176 -- for quick check
 177 import qualified Prelude
 178 import qualified List
 179 import Debug.QuickCheck
 180 import List(nub,sort)
 181 -}
 182
 183 #if __GLASGOW_HASKELL__
 184 import Text.Read
 185 import Data.Generics.Basics
 186 import Data.Generics.Instances
 187 #endif
 188
 189 {--------------------------------------------------------------------
 190   Operators
 191 --------------------------------------------------------------------}
 192 infixl 9 !,\\ --
 193
 194 -- | /O(log n)/. Find the value at a key.
 195 -- Calls 'error' when the element can not be found.
 196 (!) :: Ord k => Map k a -> k -> a
 197 m ! k    = find k m
 198
 199 -- | /O(n+m)/. See 'difference'.
 200 (\\) :: Ord k => Map k a -> Map k b -> Map k a
 201 m1 \\ m2 = difference m1 m2
 202
 203 {--------------------------------------------------------------------
 204   Size balanced trees.
 205 --------------------------------------------------------------------}
 206 -- | A Map from keys @k@ to values @a@.
 207 data Map k a  = Tip
 208               | Bin {-# UNPACK #-} !Size !k a !(Map k a) !(Map k a)
 209
 210 type Size     = Int
 211
 212 instance (Ord k) => Monoid (Map k v) where
 213     mempty  = empty
 214     mappend = union
 215     mconcat = unions
 216
 217 #if __GLASGOW_HASKELL__
 218
 219 {--------------------------------------------------------------------
 220   A Data instance
 221 --------------------------------------------------------------------}
 222
 223 -- This instance preserves data abstraction at the cost of inefficiency.
 224 -- We omit reflection services for the sake of data abstraction.
 225
 226 instance (Data k, Data a, Ord k) => Data (Map k a) where
 227   gfoldl f z map = z fromList `f` (toList map)
 228   toConstr _     = error "toConstr"
 229   gunfold _ _    = error "gunfold"
 230   dataTypeOf _   = mkNorepType "Data.Map.Map"
 231   dataCast2 f    = gcast2 f
 232
 233 #endif
 234
 235 {--------------------------------------------------------------------
 236   Query
 237 --------------------------------------------------------------------}
 238 -- | /O(1)/. Is the map empty?
 239 null :: Map k a -> Bool
 240 null t
 241   = case t of
 242       Tip             -> True
 243       Bin sz k x l r  -> False
 244
 245 -- | /O(1)/. The number of elements in the map.
 246 size :: Map k a -> Int
 247 size t
 248   = case t of
 249       Tip             -> 0
 250       Bin sz k x l r  -> sz
 251
 252
 253 -- | /O(log n)/. Lookup the value at a key in the map.
 254 lookup :: (Monad m,Ord k) => k -> Map k a -> m a
 255 lookup k t = case lookup' k t of
 256     Just x -> return x
 257     Nothing -> fail "Data.Map.lookup: Key not found"
 258 lookup' :: Ord k => k -> Map k a -> Maybe a
 259 lookup' k t
 260   = case t of
 261       Tip -> Nothing
 262       Bin sz kx x l r
 263           -> case compare k kx of
 264                LT -> lookup' k l
 265                GT -> lookup' k r
 266                EQ -> Just x
 267
 268 lookupAssoc :: Ord k => k -> Map k a -> Maybe (k,a)
 269 lookupAssoc  k t
 270   = case t of
 271       Tip -> Nothing
 272       Bin sz kx x l r
 273           -> case compare k kx of
 274                LT -> lookupAssoc k l
 275                GT -> lookupAssoc k r
 276                EQ -> Just (kx,x)
 277
 278 -- | /O(log n)/. Is the key a member of the map?
 279 member :: Ord k => k -> Map k a -> Bool
 280 member k m
 281   = case lookup k m of
 282       Nothing -> False
 283       Just x  -> True
 284
 285 -- | /O(log n)/. Is the key not a member of the map?
 286 notMember :: Ord k => k -> Map k a -> Bool
 287 notMember k m = not $ member k m
 288
 289 -- | /O(log n)/. Find the value at a key.
 290 -- Calls 'error' when the element can not be found.
 291 find :: Ord k => k -> Map k a -> a
 292 find k m
 293   = case lookup k m of
 294       Nothing -> error "Map.find: element not in the map"
 295       Just x  -> x
 296
 297 -- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns
 298 -- the value at key @k@ or returns @def@ when the key is not in the map.
 299 findWithDefault :: Ord k => a -> k -> Map k a -> a
 300 findWithDefault def k m
 301   = case lookup k m of
 302       Nothing -> def
 303       Just x  -> x
 304
 305
 306
 307 {--------------------------------------------------------------------
 308   Construction
 309 --------------------------------------------------------------------}
 310 -- | /O(1)/. The empty map.
 311 empty :: Map k a
 312 empty
 313   = Tip
 314
 315 -- | /O(1)/. A map with a single element.
 316 singleton :: k -> a -> Map k a
 317 singleton k x
 318   = Bin 1 k x Tip Tip
 319
 320 {--------------------------------------------------------------------
 321   Insertion
 322 --------------------------------------------------------------------}
 323 -- | /O(log n)/. Insert a new key and value in the map.
 324 -- If the key is already present in the map, the associated value is
 325 -- replaced with the supplied value, i.e. 'insert' is equivalent to
 326 -- @'insertWith' 'const'@.
 327 insert :: Ord k => k -> a -> Map k a -> Map k a
 328 insert kx x t
 329   = case t of
 330       Tip -> singleton kx x
 331       Bin sz ky y l r
 332           -> case compare kx ky of
 333                LT -> balance ky y (insert kx x l) r
 334                GT -> balance ky y l (insert kx x r)
 335                EQ -> Bin sz kx x l r
 336
 337 -- | /O(log n)/. Insert with a combining function.
 338 -- @'insertWith' f key value mp@
 339 -- will insert the pair (key, value) into @mp@ if key does
 340 -- not exist in the map. If the key does exist, the function will
 341 -- insert the pair @(key, f new_value old_value)@.
 342 insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
 343 insertWith f k x m
 344   = insertWithKey (\k x y -> f x y) k x m
 345
 346 -- | /O(log n)/. Insert with a combining function.
 347 -- @'insertWithKey' f key value mp@
 348 -- will insert the pair (key, value) into @mp@ if key does
 349 -- not exist in the map. If the key does exist, the function will
 350 -- insert the pair @(key,f key new_value old_value)@.
 351 -- Note that the key passed to f is the same key passed to 'insertWithKey'.
 352 insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
 353 insertWithKey f kx x t
 354   = case t of
 355       Tip -> singleton kx x
 356       Bin sy ky y l r
 357           -> case compare kx ky of
 358                LT -> balance ky y (insertWithKey f kx x l) r
 359                GT -> balance ky y l (insertWithKey f kx x r)
 360                EQ -> Bin sy kx (f kx x y) l r
 361
 362 -- | /O(log n)/. The expression (@'insertLookupWithKey' f k x map@)
 363 -- is a pair where the first element is equal to (@'lookup' k map@)
 364 -- and the second element equal to (@'insertWithKey' f k x map@).
 365 insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a,Map k a)
 366 insertLookupWithKey f kx x t
 367   = case t of
 368       Tip -> (Nothing, singleton kx x)
 369       Bin sy ky y l r
 370           -> case compare kx ky of
 371                LT -> let (found,l') = insertLookupWithKey f kx x l in (found,balance ky y l' r)
 372                GT -> let (found,r') = insertLookupWithKey f kx x r in (found,balance ky y l r')
 373                EQ -> (Just y, Bin sy kx (f kx x y) l r)
 374
 375 {--------------------------------------------------------------------
 376   Deletion
 377   [delete] is the inlined version of [deleteWith (\k x -> Nothing)]
 378 --------------------------------------------------------------------}
 379 -- | /O(log n)/. Delete a key and its value from the map. When the key is not
 380 -- a member of the map, the original map is returned.
 381 delete :: Ord k => k -> Map k a -> Map k a
 382 delete k t
 383   = case t of
 384       Tip -> Tip
 385       Bin sx kx x l r
 386           -> case compare k kx of
 387                LT -> balance kx x (delete k l) r
 388                GT -> balance kx x l (delete k r)
 389                EQ -> glue l r
 390
 391 -- | /O(log n)/. Adjust a value at a specific key. When the key is not
 392 -- a member of the map, the original map is returned.
 393 adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a
 394 adjust f k m
 395   = adjustWithKey (\k x -> f x) k m
 396
 397 -- | /O(log n)/. Adjust a value at a specific key. When the key is not
 398 -- a member of the map, the original map is returned.
 399 adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
 400 adjustWithKey f k m
 401   = updateWithKey (\k x -> Just (f k x)) k m
 402
 403 -- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@
 404 -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is
 405 -- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@.
 406 update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
 407 update f k m
 408   = updateWithKey (\k x -> f x) k m
 409
 410 -- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the
 411 -- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing',
 412 -- the element is deleted. If it is (@'Just' y@), the key @k@ is bound
 413 -- to the new value @y@.
 414 updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a
 415 updateWithKey f k t
 416   = case t of
 417       Tip -> Tip
 418       Bin sx kx x l r
 419           -> case compare k kx of
 420                LT -> balance kx x (updateWithKey f k l) r
 421                GT -> balance kx x l (updateWithKey f k r)
 422                EQ -> case f kx x of
 423                        Just x' -> Bin sx kx x' l r
 424                        Nothing -> glue l r
 425
 426 -- | /O(log n)/. Lookup and update.
 427 updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a)
 428 updateLookupWithKey f k t
 429   = case t of
 430       Tip -> (Nothing,Tip)
 431       Bin sx kx x l r
 432           -> case compare k kx of
 433                LT -> let (found,l') = updateLookupWithKey f k l in (found,balance kx x l' r)
 434                GT -> let (found,r') = updateLookupWithKey f k r in (found,balance kx x l r')
 435                EQ -> case f kx x of
 436                        Just x' -> (Just x',Bin sx kx x' l r)
 437                        Nothing -> (Just x,glue l r)
 438
 439 -- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof.
 440 -- 'alter' can be used to insert, delete, or update a value in a 'Map'.
 441 -- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@
 442 alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a
 443 alter f k t
 444   = case t of
 445       Tip -> case f Nothing of
 446                Nothing -> Tip
 447                Just x -> singleton k x
 448       Bin sx kx x l r
 449           -> case compare k kx of
 450                LT -> balance kx x (alter f k l) r
 451                GT -> balance kx x l (alter f k r)
 452                EQ -> case f (Just x) of
 453                        Just x' -> Bin sx kx x' l r
 454                        Nothing -> glue l r
 455
 456 {--------------------------------------------------------------------
 457   Indexing
 458 --------------------------------------------------------------------}
 459 -- | /O(log n)/. Return the /index/ of a key. The index is a number from
 460 -- /0/ up to, but not including, the 'size' of the map. Calls 'error' when
 461 -- the key is not a 'member' of the map.
 462 findIndex :: Ord k => k -> Map k a -> Int
 463 findIndex k t
 464   = case lookupIndex k t of
 465       Nothing  -> error "Map.findIndex: element is not in the map"
 466       Just idx -> idx
 467
 468 -- | /O(log n)/. Lookup the /index/ of a key. The index is a number from
 469 -- /0/ up to, but not including, the 'size' of the map.
 470 lookupIndex :: (Monad m,Ord k) => k -> Map k a -> m Int
 471 lookupIndex k t = case lookup 0 t of
 472     Nothing -> fail "Data.Map.lookupIndex: Key not found."
 473     Just x -> return x
 474   where
 475     lookup idx Tip  = Nothing
 476     lookup idx (Bin _ kx x l r)
 477       = case compare k kx of
 478           LT -> lookup idx l
 479           GT -> lookup (idx + size l + 1) r
 480           EQ -> Just (idx + size l)
 481
 482 -- | /O(log n)/. Retrieve an element by /index/. Calls 'error' when an
 483 -- invalid index is used.
 484 elemAt :: Int -> Map k a -> (k,a)
 485 elemAt i Tip = error "Map.elemAt: index out of range"
 486 elemAt i (Bin _ kx x l r)
 487   = case compare i sizeL of
 488       LT -> elemAt i l
 489       GT -> elemAt (i-sizeL-1) r
 490       EQ -> (kx,x)
 491   where
 492     sizeL = size l
 493
 494 -- | /O(log n)/. Update the element at /index/. Calls 'error' when an
 495 -- invalid index is used.
 496 updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
 497 updateAt f i Tip  = error "Map.updateAt: index out of range"
 498 updateAt f i (Bin sx kx x l r)
 499   = case compare i sizeL of
 500       LT -> updateAt f i l
 501       GT -> updateAt f (i-sizeL-1) r
 502       EQ -> case f kx x of
 503               Just x' -> Bin sx kx x' l r
 504               Nothing -> glue l r
 505   where
 506     sizeL = size l
 507
 508 -- | /O(log n)/. Delete the element at /index/.
 509 -- Defined as (@'deleteAt' i map = 'updateAt' (\k x -> 'Nothing') i map@).
 510 deleteAt :: Int -> Map k a -> Map k a
 511 deleteAt i map
 512   = updateAt (\k x -> Nothing) i map
 513
 514
 515 {--------------------------------------------------------------------
 516   Minimal, Maximal
 517 --------------------------------------------------------------------}
 518 -- | /O(log n)/. The minimal key of the map.
 519 findMin :: Map k a -> (k,a)
 520 findMin (Bin _ kx x Tip r)  = (kx,x)
 521 findMin (Bin _ kx x l r)    = findMin l
 522 findMin Tip                 = error "Map.findMin: empty map has no minimal element"
 523
 524 -- | /O(log n)/. The maximal key of the map.
 525 findMax :: Map k a -> (k,a)
 526 findMax (Bin _ kx x l Tip)  = (kx,x)
 527 findMax (Bin _ kx x l r)    = findMax r
 528 findMax Tip                 = error "Map.findMax: empty map has no maximal element"
 529
 530 -- | /O(log n)/. Delete the minimal key.
 531 deleteMin :: Map k a -> Map k a
 532 deleteMin (Bin _ kx x Tip r)  = r
 533 deleteMin (Bin _ kx x l r)    = balance kx x (deleteMin l) r
 534 deleteMin Tip                 = Tip
 535
 536 -- | /O(log n)/. Delete the maximal key.
 537 deleteMax :: Map k a -> Map k a
 538 deleteMax (Bin _ kx x l Tip)  = l
 539 deleteMax (Bin _ kx x l r)    = balance kx x l (deleteMax r)
 540 deleteMax Tip                 = Tip
 541
 542 -- | /O(log n)/. Update the value at the minimal key.
 543 updateMin :: (a -> Maybe a) -> Map k a -> Map k a
 544 updateMin f m
 545   = updateMinWithKey (\k x -> f x) m
 546
 547 -- | /O(log n)/. Update the value at the maximal key.
 548 updateMax :: (a -> Maybe a) -> Map k a -> Map k a
 549 updateMax f m
 550   = updateMaxWithKey (\k x -> f x) m
 551
 552
 553 -- | /O(log n)/. Update the value at the minimal key.
 554 updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
 555 updateMinWithKey f t
 556   = case t of
 557       Bin sx kx x Tip r  -> case f kx x of
 558                               Nothing -> r
 559                               Just x' -> Bin sx kx x' Tip r
 560       Bin sx kx x l r    -> balance kx x (updateMinWithKey f l) r
 561       Tip                -> Tip
 562
 563 -- | /O(log n)/. Update the value at the maximal key.
 564 updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a
 565 updateMaxWithKey f t
 566   = case t of
 567       Bin sx kx x l Tip  -> case f kx x of
 568                               Nothing -> l
 569                               Just x' -> Bin sx kx x' l Tip
 570       Bin sx kx x l r    -> balance kx x l (updateMaxWithKey f r)
 571       Tip                -> Tip
 572
 573 -- | /O(log n)/. Retrieves the minimal key of the map, and the map stripped from that element
 574 -- @fail@s (in the monad) when passed an empty map.
 575 minView :: Monad m => Map k a -> m (Map k a, (k,a))
 576 minView Tip = fail "Map.minView: empty map"
 577 minView x = return (swap $ deleteFindMin x)
 578
 579 -- | /O(log n)/. Retrieves the maximal key of the map, and the map stripped from that element
 580 -- @fail@s (in the monad) when passed an empty map.
 581 maxView :: Monad m => Map k a -> m (Map k a, (k,a))
 582 maxView Tip = fail "Map.maxView: empty map"
 583 maxView x = return (swap $ deleteFindMax x)
 584
 585 swap (a,b) = (b,a)
 586
 587 {--------------------------------------------------------------------
 588   Union.
 589 --------------------------------------------------------------------}
 590 -- | The union of a list of maps:
 591 --   (@'unions' == 'Prelude.foldl' 'union' 'empty'@).
 592 unions :: Ord k => [Map k a] -> Map k a
 593 unions ts
 594   = foldlStrict union empty ts
 595
 596 -- | The union of a list of maps, with a combining operation:
 597 --   (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@).
 598 unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a
 599 unionsWith f ts
 600   = foldlStrict (unionWith f) empty ts
 601
 602 -- | /O(n+m)/.
 603 -- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@.
 604 -- It prefers @t1@ when duplicate keys are encountered,
 605 -- i.e. (@'union' == 'unionWith' 'const'@).
 606 -- The implementation uses the efficient /hedge-union/ algorithm.
 607 -- Hedge-union is more efficient on (bigset `union` smallset)
 608 union :: Ord k => Map k a -> Map k a -> Map k a
 609 union Tip t2  = t2
 610 union t1 Tip  = t1
 611 union t1 t2 = hedgeUnionL (const LT) (const GT) t1 t2
 612
 613 -- left-biased hedge union
 614 hedgeUnionL cmplo cmphi t1 Tip
 615   = t1
 616 hedgeUnionL cmplo cmphi Tip (Bin _ kx x l r)
 617   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 618 hedgeUnionL cmplo cmphi (Bin _ kx x l r) t2
 619   = join kx x (hedgeUnionL cmplo cmpkx l (trim cmplo cmpkx t2))
 620               (hedgeUnionL cmpkx cmphi r (trim cmpkx cmphi t2))
 621   where
 622     cmpkx k  = compare kx k
 623
 624 -- right-biased hedge union
 625 hedgeUnionR cmplo cmphi t1 Tip
 626   = t1
 627 hedgeUnionR cmplo cmphi Tip (Bin _ kx x l r)
 628   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 629 hedgeUnionR cmplo cmphi (Bin _ kx x l r) t2
 630   = join kx newx (hedgeUnionR cmplo cmpkx l lt)
 631                  (hedgeUnionR cmpkx cmphi r gt)
 632   where
 633     cmpkx k     = compare kx k
 634     lt          = trim cmplo cmpkx t2
 635     (found,gt)  = trimLookupLo kx cmphi t2
 636     newx        = case found of
 637                     Nothing -> x
 638                     Just (_,y) -> y
 639
 640 {--------------------------------------------------------------------
 641   Union with a combining function
 642 --------------------------------------------------------------------}
 643 -- | /O(n+m)/. Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
 644 unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
 645 unionWith f m1 m2
 646   = unionWithKey (\k x y -> f x y) m1 m2
 647
 648 -- | /O(n+m)/.
 649 -- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
 650 -- Hedge-union is more efficient on (bigset `union` smallset).
 651 unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
 652 unionWithKey f Tip t2  = t2
 653 unionWithKey f t1 Tip  = t1
 654 unionWithKey f t1 t2 = hedgeUnionWithKey f (const LT) (const GT) t1 t2
 655
 656 hedgeUnionWithKey f cmplo cmphi t1 Tip
 657   = t1
 658 hedgeUnionWithKey f cmplo cmphi Tip (Bin _ kx x l r)
 659   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 660 hedgeUnionWithKey f cmplo cmphi (Bin _ kx x l r) t2
 661   = join kx newx (hedgeUnionWithKey f cmplo cmpkx l lt)
 662                  (hedgeUnionWithKey f cmpkx cmphi r gt)
 663   where
 664     cmpkx k     = compare kx k
 665     lt          = trim cmplo cmpkx t2
 666     (found,gt)  = trimLookupLo kx cmphi t2
 667     newx        = case found of
 668                     Nothing -> x
 669                     Just (_,y) -> f kx x y
 670
 671 {--------------------------------------------------------------------
 672   Difference
 673 --------------------------------------------------------------------}
 674 -- | /O(n+m)/. Difference of two maps.
 675 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
 676 difference :: Ord k => Map k a -> Map k b -> Map k a
 677 difference Tip t2  = Tip
 678 difference t1 Tip  = t1
 679 difference t1 t2   = hedgeDiff (const LT) (const GT) t1 t2
 680
 681 hedgeDiff cmplo cmphi Tip t
 682   = Tip
 683 hedgeDiff cmplo cmphi (Bin _ kx x l r) Tip
 684   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 685 hedgeDiff cmplo cmphi t (Bin _ kx x l r)
 686   = merge (hedgeDiff cmplo cmpkx (trim cmplo cmpkx t) l)
 687           (hedgeDiff cmpkx cmphi (trim cmpkx cmphi t) r)
 688   where
 689     cmpkx k = compare kx k
 690
 691 -- | /O(n+m)/. Difference with a combining function.
 692 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
 693 differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
 694 differenceWith f m1 m2
 695   = differenceWithKey (\k x y -> f x y) m1 m2
 696
 697 -- | /O(n+m)/. Difference with a combining function. When two equal keys are
 698 -- encountered, the combining function is applied to the key and both values.
 699 -- If it returns 'Nothing', the element is discarded (proper set difference). If
 700 -- it returns (@'Just' y@), the element is updated with a new value @y@.
 701 -- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
 702 differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
 703 differenceWithKey f Tip t2  = Tip
 704 differenceWithKey f t1 Tip  = t1
 705 differenceWithKey f t1 t2   = hedgeDiffWithKey f (const LT) (const GT) t1 t2
 706
 707 hedgeDiffWithKey f cmplo cmphi Tip t
 708   = Tip
 709 hedgeDiffWithKey f cmplo cmphi (Bin _ kx x l r) Tip
 710   = join kx x (filterGt cmplo l) (filterLt cmphi r)
 711 hedgeDiffWithKey f cmplo cmphi t (Bin _ kx x l r)
 712   = case found of
 713       Nothing -> merge tl tr
 714       Just (ky,y) ->
 715           case f ky y x of
 716             Nothing -> merge tl tr
 717             Just z  -> join ky z tl tr
 718   where
 719     cmpkx k     = compare kx k
 720     lt          = trim cmplo cmpkx t
 721     (found,gt)  = trimLookupLo kx cmphi t
 722     tl          = hedgeDiffWithKey f cmplo cmpkx lt l
 723     tr          = hedgeDiffWithKey f cmpkx cmphi gt r
 724
 725
 726
 727 {--------------------------------------------------------------------
 728   Intersection
 729 --------------------------------------------------------------------}
 730 -- | /O(n+m)/. Intersection of two maps. The values in the first
 731 -- map are returned, i.e. (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@).
 732 intersection :: Ord k => Map k a -> Map k b -> Map k a
 733 intersection m1 m2
 734   = intersectionWithKey (\k x y -> x) m1 m2
 735
 736 -- | /O(n+m)/. Intersection with a combining function.
 737 intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c
 738 intersectionWith f m1 m2
 739   = intersectionWithKey (\k x y -> f x y) m1 m2
 740
 741 -- | /O(n+m)/. Intersection with a combining function.
 742 -- Intersection is more efficient on (bigset `intersection` smallset)
 743 --intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
 744 --intersectionWithKey f Tip t = Tip
 745 --intersectionWithKey f t Tip = Tip
 746 --intersectionWithKey f t1 t2 = intersectWithKey f t1 t2
 747 --
 748 --intersectWithKey f Tip t = Tip
 749 --intersectWithKey f t Tip = Tip
 750 --intersectWithKey f t (Bin _ kx x l r)
 751 --  = case found of
 752 --      Nothing -> merge tl tr
 753 --      Just y  -> join kx (f kx y x) tl tr
 754 --  where
 755 --    (lt,found,gt) = splitLookup kx t
 756 --    tl            = intersectWithKey f lt l
 757 --    tr            = intersectWithKey f gt r
 758
 759
 760 intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
 761 intersectionWithKey f Tip t = Tip
 762 intersectionWithKey f t Tip = Tip
 763 intersectionWithKey f t1@(Bin s1 k1 x1 l1 r1) t2@(Bin s2 k2 x2 l2 r2) =
 764    if s1 >= s2 then
 765       let (lt,found,gt) = splitLookupWithKey k2 t1
 766           tl            = intersectionWithKey f lt l2
 767           tr            = intersectionWithKey f gt r2
 768       in case found of
 769       Just (k,x) -> join k (f k x x2) tl tr
 770       Nothing -> merge tl tr
 771    else let (lt,found,gt) = splitLookup k1 t2
 772             tl            = intersectionWithKey f l1 lt
 773             tr            = intersectionWithKey f r1 gt
 774       in case found of
 775       Just x -> join k1 (f k1 x1 x) tl tr
 776       Nothing -> merge tl tr
 777
 778
 779
 780 {--------------------------------------------------------------------
 781   Submap
 782 --------------------------------------------------------------------}
 783 -- | /O(n+m)/.
 784 -- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@).
 785 isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
 786 isSubmapOf m1 m2
 787   = isSubmapOfBy (==) m1 m2
 788
 789 {- | /O(n+m)/.
 790  The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if
 791  all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when
 792  applied to their respective values. For example, the following
 793  expressions are all 'True':
 794
 795  > isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 796  > isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 797  > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)])
 798
 799  But the following are all 'False':
 800
 801  > isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)])
 802  > isSubmapOfBy (<)  (fromList [('a',1)]) (fromList [('a',1),('b',2)])
 803  > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)])
 804 -}
 805 isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool
 806 isSubmapOfBy f t1 t2
 807   = (size t1 <= size t2) && (submap' f t1 t2)
 808
 809 submap' f Tip t = True
 810 submap' f t Tip = False
 811 submap' f (Bin _ kx x l r) t
 812   = case found of
 813       Nothing -> False
 814       Just y  -> f x y && submap' f l lt && submap' f r gt
 815   where
 816     (lt,found,gt) = splitLookup kx t
 817
 818 -- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
 819 -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@).
 820 isProperSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool
 821 isProperSubmapOf m1 m2
 822   = isProperSubmapOfBy (==) m1 m2
 823
 824 {- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal).
 825  The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when
 826  @m1@ and @m2@ are not equal,
 827  all keys in @m1@ are in @m2@, and when @f@ returns 'True' when
 828  applied to their respective values. For example, the following
 829  expressions are all 'True':
 830
 831   > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 832   > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)])
 833
 834  But the following are all 'False':
 835
 836   > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)])
 837   > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)])
 838   > isProperSubmapOfBy (<)  (fromList [(1,1)])       (fromList [(1,1),(2,2)])
 839 -}
 840 isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool
 841 isProperSubmapOfBy f t1 t2
 842   = (size t1 < size t2) && (submap' f t1 t2)
 843
 844 {--------------------------------------------------------------------
 845   Filter and partition
 846 --------------------------------------------------------------------}
 847 -- | /O(n)/. Filter all values that satisfy the predicate.
 848 filter :: Ord k => (a -> Bool) -> Map k a -> Map k a
 849 filter p m
 850   = filterWithKey (\k x -> p x) m
 851
 852 -- | /O(n)/. Filter all keys\/values that satisfy the predicate.
 853 filterWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> Map k a
 854 filterWithKey p Tip = Tip
 855 filterWithKey p (Bin _ kx x l r)
 856   | p kx x    = join kx x (filterWithKey p l) (filterWithKey p r)
 857   | otherwise = merge (filterWithKey p l) (filterWithKey p r)
 858
 859
 860 -- | /O(n)/. partition the map according to a predicate. The first
 861 -- map contains all elements that satisfy the predicate, the second all
 862 -- elements that fail the predicate. See also 'split'.
 863 partition :: Ord k => (a -> Bool) -> Map k a -> (Map k a,Map k a)
 864 partition p m
 865   = partitionWithKey (\k x -> p x) m
 866
 867 -- | /O(n)/. partition the map according to a predicate. The first
 868 -- map contains all elements that satisfy the predicate, the second all
 869 -- elements that fail the predicate. See also 'split'.
 870 partitionWithKey :: Ord k => (k -> a -> Bool) -> Map k a -> (Map k a,Map k a)
 871 partitionWithKey p Tip = (Tip,Tip)
 872 partitionWithKey p (Bin _ kx x l r)
 873   | p kx x    = (join kx x l1 r1,merge l2 r2)
 874   | otherwise = (merge l1 r1,join kx x l2 r2)
 875   where
 876     (l1,l2) = partitionWithKey p l
 877     (r1,r2) = partitionWithKey p r
 878
 879 -- | /O(n)/. Map values and collect the 'Just' results.
 880 mapMaybe :: Ord k => (a -> Maybe b) -> Map k a -> Map k b
 881 mapMaybe f m
 882   = mapMaybeWithKey (\k x -> f x) m
 883
 884 -- | /O(n)/. Map keys\/values and collect the 'Just' results.
 885 mapMaybeWithKey :: Ord k => (k -> a -> Maybe b) -> Map k a -> Map k b
 886 mapMaybeWithKey f Tip = Tip
 887 mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of
 888   Just y  -> join kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
 889   Nothing -> merge (mapMaybeWithKey f l) (mapMaybeWithKey f r)
 890
 891 -- | /O(n)/. Map values and separate the 'Left' and 'Right' results.
 892 mapEither :: Ord k => (a -> Either b c) -> Map k a -> (Map k b, Map k c)
 893 mapEither f m
 894   = mapEitherWithKey (\k x -> f x) m
 895
 896 -- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results.
 897 mapEitherWithKey :: Ord k =>
 898   (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c)
 899 mapEitherWithKey f Tip = (Tip, Tip)
 900 mapEitherWithKey f (Bin _ kx x l r) = case f kx x of
 901   Left y  -> (join kx y l1 r1, merge l2 r2)
 902   Right z -> (merge l1 r1, join kx z l2 r2)
 903   where
 904     (l1,l2) = mapEitherWithKey f l
 905     (r1,r2) = mapEitherWithKey f r
 906
 907 {--------------------------------------------------------------------
 908   Mapping
 909 --------------------------------------------------------------------}
 910 -- | /O(n)/. Map a function over all values in the map.
 911 map :: (a -> b) -> Map k a -> Map k b
 912 map f m
 913   = mapWithKey (\k x -> f x) m
 914
 915 -- | /O(n)/. Map a function over all values in the map.
 916 mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
 917 mapWithKey f Tip = Tip
 918 mapWithKey f (Bin sx kx x l r)
 919   = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)
 920
 921 -- | /O(n)/. The function 'mapAccum' threads an accumulating
 922 -- argument through the map in ascending order of keys.
 923 mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 924 mapAccum f a m
 925   = mapAccumWithKey (\a k x -> f a x) a m
 926
 927 -- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating
 928 -- argument through the map in ascending order of keys.
 929 mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 930 mapAccumWithKey f a t
 931   = mapAccumL f a t
 932
 933 -- | /O(n)/. The function 'mapAccumL' threads an accumulating
 934 -- argument throught the map in ascending order of keys.
 935 mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 936 mapAccumL f a t
 937   = case t of
 938       Tip -> (a,Tip)
 939       Bin sx kx x l r
 940           -> let (a1,l') = mapAccumL f a l
 941                  (a2,x') = f a1 kx x
 942                  (a3,r') = mapAccumL f a2 r
 943              in (a3,Bin sx kx x' l' r')
 944
 945 -- | /O(n)/. The function 'mapAccumR' threads an accumulating
 946 -- argument throught the map in descending order of keys.
 947 mapAccumR :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c)
 948 mapAccumR f a t
 949   = case t of
 950       Tip -> (a,Tip)
 951       Bin sx kx x l r
 952           -> let (a1,r') = mapAccumR f a r
 953                  (a2,x') = f a1 kx x
 954                  (a3,l') = mapAccumR f a2 l
 955              in (a3,Bin sx kx x' l' r')
 956
 957 -- | /O(n*log n)/.
 958 -- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@.
 959 --
 960 -- The size of the result may be smaller if @f@ maps two or more distinct
 961 -- keys to the same new key.  In this case the value at the smallest of
 962 -- these keys is retained.
 963
 964 mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a
 965 mapKeys = mapKeysWith (\x y->x)
 966
 967 -- | /O(n*log n)/.
 968 -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@.
 969 --
 970 -- The size of the result may be smaller if @f@ maps two or more distinct
 971 -- keys to the same new key.  In this case the associated values will be
 972 -- combined using @c@.
 973
 974 mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a
 975 mapKeysWith c f = fromListWith c . List.map fFirst . toList
 976     where fFirst (x,y) = (f x, y)
 977
 978
 979 -- | /O(n)/.
 980 -- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@
 981 -- is strictly monotonic.
 982 -- /The precondition is not checked./
 983 -- Semi-formally, we have:
 984 --
 985 -- > and [x < y ==> f x < f y | x <- ls, y <- ls]
 986 -- >                     ==> mapKeysMonotonic f s == mapKeys f s
 987 -- >     where ls = keys s
 988
 989 mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a
 990 mapKeysMonotonic f Tip = Tip
 991 mapKeysMonotonic f (Bin sz k x l r) =
 992     Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r)
 993
 994 {--------------------------------------------------------------------
 995   Folds
 996 --------------------------------------------------------------------}
 997
 998 -- | /O(n)/. Fold the values in the map, such that
 999 -- @'fold' f z == 'Prelude.foldr' f z . 'elems'@.
1000 -- For example,
1001 --
1002 -- > elems map = fold (:) [] map
1003 --
1004 fold :: (a -> b -> b) -> b -> Map k a -> b
1005 fold f z m
1006   = foldWithKey (\k x z -> f x z) z m
1007
1008 -- | /O(n)/. Fold the keys and values in the map, such that
1009 -- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@.
1010 -- For example,
1011 --
1012 -- > keys map = foldWithKey (\k x ks -> k:ks) [] map
1013 --
1014 foldWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b
1015 foldWithKey f z t
1016   = foldr f z t
1017
1018 -- | /O(n)/. In-order fold.
1019 foldi :: (k -> a -> b -> b -> b) -> b -> Map k a -> b
1020 foldi f z Tip               = z
1021 foldi f z (Bin _ kx x l r)  = f kx x (foldi f z l) (foldi f z r)
1022
1023 -- | /O(n)/. Post-order fold.
1024 foldr :: (k -> a -> b -> b) -> b -> Map k a -> b
1025 foldr f z Tip              = z
1026 foldr f z (Bin _ kx x l r) = foldr f (f kx x (foldr f z r)) l
1027
1028 -- | /O(n)/. Pre-order fold.
1029 foldl :: (b -> k -> a -> b) -> b -> Map k a -> b
1030 foldl f z Tip              = z
1031 foldl f z (Bin _ kx x l r) = foldl f (f (foldl f z l) kx x) r
1032
1033 {--------------------------------------------------------------------
1034   List variations
1035 --------------------------------------------------------------------}
1036 -- | /O(n)/.
1037 -- Return all elements of the map in the ascending order of their keys.
1038 elems :: Map k a -> [a]
1039 elems m
1040   = [x | (k,x) <- assocs m]
1041
1042 -- | /O(n)/. Return all keys of the map in ascending order.
1043 keys  :: Map k a -> [k]
1044 keys m
1045   = [k | (k,x) <- assocs m]
1046
1047 -- | /O(n)/. The set of all keys of the map.
1048 keysSet :: Map k a -> Set.Set k
1049 keysSet m = Set.fromDistinctAscList (keys m)
1050
1051 -- | /O(n)/. Return all key\/value pairs in the map in ascending key order.
1052 assocs :: Map k a -> [(k,a)]
1053 assocs m
1054   = toList m
1055
1056 {--------------------------------------------------------------------
1057   Lists
1058   use [foldlStrict] to reduce demand on the control-stack
1059 --------------------------------------------------------------------}
1060 -- | /O(n*log n)/. Build a map from a list of key\/value pairs. See also 'fromAscList'.
1061 fromList :: Ord k => [(k,a)] -> Map k a
1062 fromList xs
1063   = foldlStrict ins empty xs
1064   where
1065     ins t (k,x) = insert k x t
1066
1067 -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'.
1068 fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a
1069 fromListWith f xs
1070   = fromListWithKey (\k x y -> f x y) xs
1071
1072 -- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'.
1073 fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
1074 fromListWithKey f xs
1075   = foldlStrict ins empty xs
1076   where
1077     ins t (k,x) = insertWithKey f k x t
1078
1079 -- | /O(n)/. Convert to a list of key\/value pairs.
1080 toList :: Map k a -> [(k,a)]
1081 toList t      = toAscList t
1082
1083 -- | /O(n)/. Convert to an ascending list.
1084 toAscList :: Map k a -> [(k,a)]
1085 toAscList t   = foldr (\k x xs -> (k,x):xs) [] t
1086
1087 -- | /O(n)/.
1088 toDescList :: Map k a -> [(k,a)]
1089 toDescList t  = foldl (\xs k x -> (k,x):xs) [] t
1090
1091
1092 {--------------------------------------------------------------------
1093   Building trees from ascending/descending lists can be done in linear time.
1094
1095   Note that if [xs] is ascending that:
1096     fromAscList xs       == fromList xs
1097     fromAscListWith f xs == fromListWith f xs
1098 --------------------------------------------------------------------}
1099 -- | /O(n)/. Build a map from an ascending list in linear time.
1100 -- /The precondition (input list is ascending) is not checked./
1101 fromAscList :: Eq k => [(k,a)] -> Map k a
1102 fromAscList xs
1103   = fromAscListWithKey (\k x y -> x) xs
1104
1105 -- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys.
1106 -- /The precondition (input list is ascending) is not checked./
1107 fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
1108 fromAscListWith f xs
1109   = fromAscListWithKey (\k x y -> f x y) xs
1110
1111 -- | /O(n)/. Build a map from an ascending list in linear time with a
1112 -- combining function for equal keys.
1113 -- /The precondition (input list is ascending) is not checked./
1114 fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
1115 fromAscListWithKey f xs
1116   = fromDistinctAscList (combineEq f xs)
1117   where
1118   -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
1119   combineEq f xs
1120     = case xs of
1121         []     -> []
1122         [x]    -> [x]
1123         (x:xx) -> combineEq' x xx
1124
1125   combineEq' z [] = [z]
1126   combineEq' z@(kz,zz) (x@(kx,xx):xs)
1127     | kx==kz    = let yy = f kx xx zz in combineEq' (kx,yy) xs
1128     | otherwise = z:combineEq' x xs
1129
1130
1131 -- | /O(n)/. Build a map from an ascending list of distinct elements in linear time.
1132 -- /The precondition is not checked./
1133 fromDistinctAscList :: [(k,a)] -> Map k a
1134 fromDistinctAscList xs
1135   = build const (length xs) xs
1136   where
1137     -- 1) use continutations so that we use heap space instead of stack space.
1138     -- 2) special case for n==5 to build bushier trees.
1139     build c 0 xs   = c Tip xs
1140     build c 5 xs   = case xs of
1141                        ((k1,x1):(k2,x2):(k3,x3):(k4,x4):(k5,x5):xx)
1142                             -> c (bin k4 x4 (bin k2 x2 (singleton k1 x1) (singleton k3 x3)) (singleton k5 x5)) xx
1143     build c n xs   = seq nr $ build (buildR nr c) nl xs
1144                    where
1145                      nl = n `div` 2
1146                      nr = n - nl - 1
1147
1148     buildR n c l ((k,x):ys) = build (buildB l k x c) n ys
1149     buildB l k x c r zs     = c (bin k x l r) zs
1150
1151
1152
1153 {--------------------------------------------------------------------
1154   Utility functions that return sub-ranges of the original
1155   tree. Some functions take a comparison function as argument to
1156   allow comparisons against infinite values. A function [cmplo k]
1157   should be read as [compare lo k].
1158
1159   [trim cmplo cmphi t]  A tree that is either empty or where [cmplo k == LT]
1160                         and [cmphi k == GT] for the key [k] of the root.
1161   [filterGt cmp t]      A tree where for all keys [k]. [cmp k == LT]
1162   [filterLt cmp t]      A tree where for all keys [k]. [cmp k == GT]
1163
1164   [split k t]           Returns two trees [l] and [r] where all keys
1165                         in [l] are <[k] and all keys in [r] are >[k].
1166   [splitLookup k t]     Just like [split] but also returns whether [k]
1167                         was found in the tree.
1168 --------------------------------------------------------------------}
1169
1170 {--------------------------------------------------------------------
1171   [trim lo hi t] trims away all subtrees that surely contain no
1172   values between the range [lo] to [hi]. The returned tree is either
1173   empty or the key of the root is between @lo@ and @hi@.
1174 --------------------------------------------------------------------}
1175 trim :: (k -> Ordering) -> (k -> Ordering) -> Map k a -> Map k a
1176 trim cmplo cmphi Tip = Tip
1177 trim cmplo cmphi t@(Bin sx kx x l r)
1178   = case cmplo kx of
1179       LT -> case cmphi kx of
1180               GT -> t
1181               le -> trim cmplo cmphi l
1182       ge -> trim cmplo cmphi r
1183
1184 trimLookupLo :: Ord k => k -> (k -> Ordering) -> Map k a -> (Maybe (k,a), Map k a)
1185 trimLookupLo lo cmphi Tip = (Nothing,Tip)
1186 trimLookupLo lo cmphi t@(Bin sx kx x l r)
1187   = case compare lo kx of
1188       LT -> case cmphi kx of
1189               GT -> (lookupAssoc lo t, t)
1190               le -> trimLookupLo lo cmphi l
1191       GT -> trimLookupLo lo cmphi r
1192       EQ -> (Just (kx,x),trim (compare lo) cmphi r)
1193
1194
1195 {--------------------------------------------------------------------
1196   [filterGt k t] filter all keys >[k] from tree [t]
1197   [filterLt k t] filter all keys <[k] from tree [t]
1198 --------------------------------------------------------------------}
1199 filterGt :: Ord k => (k -> Ordering) -> Map k a -> Map k a
1200 filterGt cmp Tip = Tip
1201 filterGt cmp (Bin sx kx x l r)
1202   = case cmp kx of
1203       LT -> join kx x (filterGt cmp l) r
1204       GT -> filterGt cmp r
1205       EQ -> r
1206
1207 filterLt :: Ord k => (k -> Ordering) -> Map k a -> Map k a
1208 filterLt cmp Tip = Tip
1209 filterLt cmp (Bin sx kx x l r)
1210   = case cmp kx of
1211       LT -> filterLt cmp l
1212       GT -> join kx x l (filterLt cmp r)
1213       EQ -> l
1214
1215 {--------------------------------------------------------------------
1216   Split
1217 --------------------------------------------------------------------}
1218 -- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@ where
1219 -- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@.
1220 split :: Ord k => k -> Map k a -> (Map k a,Map k a)
1221 split k Tip = (Tip,Tip)
1222 split k (Bin sx kx x l r)
1223   = case compare k kx of
1224       LT -> let (lt,gt) = split k l in (lt,join kx x gt r)
1225       GT -> let (lt,gt) = split k r in (join kx x l lt,gt)
1226       EQ -> (l,r)
1227
1228 -- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just
1229 -- like 'split' but also returns @'lookup' k map@.
1230 splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a)
1231 splitLookup k Tip = (Tip,Nothing,Tip)
1232 splitLookup k (Bin sx kx x l r)
1233   = case compare k kx of
1234       LT -> let (lt,z,gt) = splitLookup k l in (lt,z,join kx x gt r)
1235       GT -> let (lt,z,gt) = splitLookup k r in (join kx x l lt,z,gt)
1236       EQ -> (l,Just x,r)
1237
1238 -- | /O(log n)/.
1239 splitLookupWithKey :: Ord k => k -> Map k a -> (Map k a,Maybe (k,a),Map k a)
1240 splitLookupWithKey k Tip = (Tip,Nothing,Tip)
1241 splitLookupWithKey k (Bin sx kx x l r)
1242   = case compare k kx of
1243       LT -> let (lt,z,gt) = splitLookupWithKey k l in (lt,z,join kx x gt r)
1244       GT -> let (lt,z,gt) = splitLookupWithKey k r in (join kx x l lt,z,gt)
1245       EQ -> (l,Just (kx, x),r)
1246
1247 -- | /O(log n)/. Performs a 'split' but also returns whether the pivot
1248 -- element was found in the original set.
1249 splitMember :: Ord k => k -> Map k a -> (Map k a,Bool,Map k a)
1250 splitMember x t = let (l,m,r) = splitLookup x t in
1251      (l,maybe False (const True) m,r)
1252
1253
1254 {--------------------------------------------------------------------
1255   Utility functions that maintain the balance properties of the tree.
1256   All constructors assume that all values in [l] < [k] and all values
1257   in [r] > [k], and that [l] and [r] are valid trees.
1258
1259   In order of sophistication:
1260     [Bin sz k x l r]  The type constructor.
1261     [bin k x l r]     Maintains the correct size, assumes that both [l]
1262                       and [r] are balanced with respect to each other.
1263     [balance k x l r] Restores the balance and size.
1264                       Assumes that the original tree was balanced and
1265                       that [l] or [r] has changed by at most one element.
1266     [join k x l r]    Restores balance and size.
1267
1268   Furthermore, we can construct a new tree from two trees. Both operations
1269   assume that all values in [l] < all values in [r] and that [l] and [r]
1270   are valid:
1271     [glue l r]        Glues [l] and [r] together. Assumes that [l] and
1272                       [r] are already balanced with respect to each other.
1273     [merge l r]       Merges two trees and restores balance.
1274
1275   Note: in contrast to Adam's paper, we use (<=) comparisons instead
1276   of (<) comparisons in [join], [merge] and [balance].
1277   Quickcheck (on [difference]) showed that this was necessary in order
1278   to maintain the invariants. It is quite unsatisfactory that I haven't
1279   been able to find out why this is actually the case! Fortunately, it
1280   doesn't hurt to be a bit more conservative.
1281 --------------------------------------------------------------------}
1282
1283 {--------------------------------------------------------------------
1284   Join
1285 --------------------------------------------------------------------}
1286 join :: Ord k => k -> a -> Map k a -> Map k a -> Map k a
1287 join kx x Tip r  = insertMin kx x r
1288 join kx x l Tip  = insertMax kx x l
1289 join kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz)
1290   | delta*sizeL <= sizeR  = balance kz z (join kx x l lz) rz
1291   | delta*sizeR <= sizeL  = balance ky y ly (join kx x ry r)
1292   | otherwise             = bin kx x l r
1293
1294
1295 -- insertMin and insertMax don't perform potentially expensive comparisons.
1296 insertMax,insertMin :: k -> a -> Map k a -> Map k a
1297 insertMax kx x t
1298   = case t of
1299       Tip -> singleton kx x
1300       Bin sz ky y l r
1301           -> balance ky y l (insertMax kx x r)
1302
1303 insertMin kx x t
1304   = case t of
1305       Tip -> singleton kx x
1306       Bin sz ky y l r
1307           -> balance ky y (insertMin kx x l) r
1308
1309 {--------------------------------------------------------------------
1310   [merge l r]: merges two trees.
1311 --------------------------------------------------------------------}
1312 merge :: Map k a -> Map k a -> Map k a
1313 merge Tip r   = r
1314 merge l Tip   = l
1315 merge l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry)
1316   | delta*sizeL <= sizeR = balance ky y (merge l ly) ry
1317   | delta*sizeR <= sizeL = balance kx x lx (merge rx r)
1318   | otherwise            = glue l r
1319
1320 {--------------------------------------------------------------------
1321   [glue l r]: glues two trees together.
1322   Assumes that [l] and [r] are already balanced with respect to each other.
1323 --------------------------------------------------------------------}
1324 glue :: Map k a -> Map k a -> Map k a
1325 glue Tip r = r
1326 glue l Tip = l
1327 glue l r
1328   | size l > size r = let ((km,m),l') = deleteFindMax l in balance km m l' r
1329   | otherwise       = let ((km,m),r') = deleteFindMin r in balance km m l r'
1330
1331
1332 -- | /O(log n)/. Delete and find the minimal element.
1333 deleteFindMin :: Map k a -> ((k,a),Map k a)
1334 deleteFindMin t
1335   = case t of
1336       Bin _ k x Tip r -> ((k,x),r)
1337       Bin _ k x l r   -> let (km,l') = deleteFindMin l in (km,balance k x l' r)
1338       Tip             -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip)
1339
1340 -- | /O(log n)/. Delete and find the maximal element.
1341 deleteFindMax :: Map k a -> ((k,a),Map k a)
1342 deleteFindMax t
1343   = case t of
1344       Bin _ k x l Tip -> ((k,x),l)
1345       Bin _ k x l r   -> let (km,r') = deleteFindMax r in (km,balance k x l r')
1346       Tip             -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip)
1347
1348
1349 {--------------------------------------------------------------------
1350   [balance l x r] balances two trees with value x.
1351   The sizes of the trees should balance after decreasing the
1352   size of one of them. (a rotation).
1353
1354   [delta] is the maximal relative difference between the sizes of
1355           two trees, it corresponds with the [w] in Adams' paper.
1356   [ratio] is the ratio between an outer and inner sibling of the
1357           heavier subtree in an unbalanced setting. It determines
1358           whether a double or single rotation should be performed
1359           to restore balance. It is correspondes with the inverse
1360           of $\alpha$ in Adam's article.
1361
1362   Note that:
1363   - [delta] should be larger than 4.646 with a [ratio] of 2.
1364   - [delta] should be larger than 3.745 with a [ratio] of 1.534.
1365
1366   - A lower [delta] leads to a more 'perfectly' balanced tree.
1367   - A higher [delta] performs less rebalancing.
1368
1369   - Balancing is automatic for random data and a balancing
1370     scheme is only necessary to avoid pathological worst cases.
1371     Almost any choice will do, and in practice, a rather large
1372     [delta] may perform better than smaller one.
1373
1374   Note: in contrast to Adam's paper, we use a ratio of (at least) [2]
1375   to decide whether a single or double rotation is needed. Allthough
1376   he actually proves that this ratio is needed to maintain the
1377   invariants, his implementation uses an invalid ratio of [1].
1378 --------------------------------------------------------------------}
1379 delta,ratio :: Int
1380 delta = 5
1381 ratio = 2
1382
1383 balance :: k -> a -> Map k a -> Map k a -> Map k a
1384 balance k x l r
1385   | sizeL + sizeR <= 1    = Bin sizeX k x l r
1386   | sizeR >= delta*sizeL  = rotateL k x l r
1387   | sizeL >= delta*sizeR  = rotateR k x l r
1388   | otherwise             = Bin sizeX k x l r
1389   where
1390     sizeL = size l
1391     sizeR = size r
1392     sizeX = sizeL + sizeR + 1
1393
1394 -- rotate
1395 rotateL k x l r@(Bin _ _ _ ly ry)
1396   | size ly < ratio*size ry = singleL k x l r
1397   | otherwise               = doubleL k x l r
1398
1399 rotateR k x l@(Bin _ _ _ ly ry) r
1400   | size ry < ratio*size ly = singleR k x l r
1401   | otherwise               = doubleR k x l r
1402
1403 -- basic rotations
1404 singleL k1 x1 t1 (Bin _ k2 x2 t2 t3)  = bin k2 x2 (bin k1 x1 t1 t2) t3
1405 singleR k1 x1 (Bin _ k2 x2 t1 t2) t3  = bin k2 x2 t1 (bin k1 x1 t2 t3)
1406
1407 doubleL k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4)
1408 doubleR k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4)
1409
1410
1411 {--------------------------------------------------------------------
1412   The bin constructor maintains the size of the tree
1413 --------------------------------------------------------------------}
1414 bin :: k -> a -> Map k a -> Map k a -> Map k a
1415 bin k x l r
1416   = Bin (size l + size r + 1) k x l r
1417
1418
1419 {--------------------------------------------------------------------
1420   Eq converts the tree to a list. In a lazy setting, this
1421   actually seems one of the faster methods to compare two trees
1422   and it is certainly the simplest :-)
1423 --------------------------------------------------------------------}
1424 instance (Eq k,Eq a) => Eq (Map k a) where
1425   t1 == t2  = (size t1 == size t2) && (toAscList t1 == toAscList t2)
1426
1427 {--------------------------------------------------------------------
1428   Ord
1429 --------------------------------------------------------------------}
1430
1431 instance (Ord k, Ord v) => Ord (Map k v) where
1432     compare m1 m2 = compare (toAscList m1) (toAscList m2)
1433
1434 {--------------------------------------------------------------------
1435   Functor
1436 --------------------------------------------------------------------}
1437 instance Functor (Map k) where
1438   fmap f m  = map f m
1439
1440 instance Traversable (Map k) where
1441   traverse f Tip = pure Tip
1442   traverse f (Bin s k v l r)
1443     = flip (Bin s k) <$> traverse f l <*> f v <*> traverse f r
1444
1445 instance Foldable (Map k) where
1446   foldMap _f Tip = mempty
1447   foldMap f (Bin _s _k v l r)
1448     = foldMap f l `mappend` f v `mappend` foldMap f r
1449
1450 {--------------------------------------------------------------------
1451   Read
1452 --------------------------------------------------------------------}
1453 instance (Ord k, Read k, Read e) => Read (Map k e) where
1454 #ifdef __GLASGOW_HASKELL__
1455   readPrec = parens $ prec 10 $ do
1456     Ident "fromList" <- lexP
1457     xs <- readPrec
1458     return (fromList xs)
1459
1460   readListPrec = readListPrecDefault
1461 #else
1462   readsPrec p = readParen (p > 10) $ \ r -> do
1463     ("fromList",s) <- lex r
1464     (xs,t) <- reads s
1465     return (fromList xs,t)
1466 #endif
1467
1468 -- parses a pair of things with the syntax a:=b
1469 readPair :: (Read a, Read b) => ReadS (a,b)
1470 readPair s = do (a, ct1)    <- reads s
1471                 (":=", ct2) <- lex ct1
1472                 (b, ct3)    <- reads ct2
1473                 return ((a,b), ct3)
1474
1475 {--------------------------------------------------------------------
1476   Show
1477 --------------------------------------------------------------------}
1478 instance (Show k, Show a) => Show (Map k a) where
1479   showsPrec d m  = showParen (d > 10) $
1480     showString "fromList " . shows (toList m)
1481
1482 showMap :: (Show k,Show a) => [(k,a)] -> ShowS
1483 showMap []
1484   = showString "{}"
1485 showMap (x:xs)
1486   = showChar '{' . showElem x . showTail xs
1487   where
1488     showTail []     = showChar '}'
1489     showTail (x:xs) = showString ", " . showElem x . showTail xs
1490
1491     showElem (k,x)  = shows k . showString " := " . shows x
1492
1493
1494 -- | /O(n)/. Show the tree that implements the map. The tree is shown
1495 -- in a compressed, hanging format.
1496 showTree :: (Show k,Show a) => Map k a -> String
1497 showTree m
1498   = showTreeWith showElem True False m
1499   where
1500     showElem k x  = show k ++ ":=" ++ show x
1501
1502
1503 {- | /O(n)/. The expression (@'showTreeWith' showelem hang wide map@) shows
1504  the tree that implements the map. Elements are shown using the @showElem@ function. If @hang@ is
1505  'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If
1506  @wide@ is 'True', an extra wide version is shown.
1507
1508 >  Map> let t = fromDistinctAscList [(x,()) | x <- [1..5]]
1509 >  Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True False t
1510 >  (4,())
1511 >  +--(2,())
1512 >  |  +--(1,())
1513 >  |  +--(3,())
1514 >  +--(5,())
1515 >
1516 >  Map> putStrLn $ showTreeWith (\k x -> show (k,x)) True True t
1517 >  (4,())
1518 >  |
1519 >  +--(2,())
1520 >  |  |
1521 >  |  +--(1,())
1522 >  |  |
1523 >  |  +--(3,())
1524 >  |
1525 >  +--(5,())
1526 >
1527 >  Map> putStrLn $ showTreeWith (\k x -> show (k,x)) False True t
1528 >  +--(5,())
1529 >  |
1530 >  (4,())
1531 >  |
1532 >  |  +--(3,())
1533 >  |  |
1534 >  +--(2,())
1535 >     |
1536 >     +--(1,())
1537
1538 -}
1539 showTreeWith :: (k -> a -> String) -> Bool -> Bool -> Map k a -> String
1540 showTreeWith showelem hang wide t
1541   | hang      = (showsTreeHang showelem wide [] t) ""
1542   | otherwise = (showsTree showelem wide [] [] t) ""
1543
1544 showsTree :: (k -> a -> String) -> Bool -> [String] -> [String] -> Map k a -> ShowS
1545 showsTree showelem wide lbars rbars t
1546   = case t of
1547       Tip -> showsBars lbars . showString "|\n"
1548       Bin sz kx x Tip Tip
1549           -> showsBars lbars . showString (showelem kx x) . showString "\n"
1550       Bin sz kx x l r
1551           -> showsTree showelem wide (withBar rbars) (withEmpty rbars) r .
1552              showWide wide rbars .
1553              showsBars lbars . showString (showelem kx x) . showString "\n" .
1554              showWide wide lbars .
1555              showsTree showelem wide (withEmpty lbars) (withBar lbars) l
1556
1557 showsTreeHang :: (k -> a -> String) -> Bool -> [String] -> Map k a -> ShowS
1558 showsTreeHang showelem wide bars t
1559   = case t of
1560       Tip -> showsBars bars . showString "|\n"
1561       Bin sz kx x Tip Tip
1562           -> showsBars bars . showString (showelem kx x) . showString "\n"
1563       Bin sz kx x l r
1564           -> showsBars bars . showString (showelem kx x) . showString "\n" .
1565              showWide wide bars .
1566              showsTreeHang showelem wide (withBar bars) l .
1567              showWide wide bars .
1568              showsTreeHang showelem wide (withEmpty bars) r
1569
1570
1571 showWide wide bars
1572   | wide      = showString (concat (reverse bars)) . showString "|\n"
1573   | otherwise = id
1574
1575 showsBars :: [String] -> ShowS
1576 showsBars bars
1577   = case bars of
1578       [] -> id
1579       _  -> showString (concat (reverse (tail bars))) . showString node
1580
1581 node           = "+--"
1582 withBar bars   = "|  ":bars
1583 withEmpty bars = "   ":bars
1584
1585 {--------------------------------------------------------------------
1586   Typeable
1587 --------------------------------------------------------------------}
1588
1589 #include "Typeable.h"
1590 INSTANCE_TYPEABLE2(Map,mapTc,"Map")
1591
1592 {--------------------------------------------------------------------
1593   Assertions
1594 --------------------------------------------------------------------}
1595 -- | /O(n)/. Test if the internal map structure is valid.
1596 valid :: Ord k => Map k a -> Bool
1597 valid t
1598   = balanced t && ordered t && validsize t
1599
1600 ordered t
1601   = bounded (const True) (const True) t
1602   where
1603     bounded lo hi t
1604       = case t of
1605           Tip              -> True
1606           Bin sz kx x l r  -> (lo kx) && (hi kx) && bounded lo (<kx) l && bounded (>kx) hi r
1607
1608 -- | Exported only for "Debug.QuickCheck"
1609 balanced :: Map k a -> Bool
1610 balanced t
1611   = case t of
1612       Tip              -> True
1613       Bin sz kx x l r  -> (size l + size r <= 1 || (size l <= delta*size r && size r <= delta*size l)) &&
1614                           balanced l && balanced r
1615
1616
1617 validsize t
1618   = (realsize t == Just (size t))
1619   where
1620     realsize t
1621       = case t of
1622           Tip             -> Just 0
1623           Bin sz kx x l r -> case (realsize l,realsize r) of
1624                               (Just n,Just m)  | n+m+1 == sz  -> Just sz
1625                               other            -> Nothing
1626
1627 {--------------------------------------------------------------------
1628   Utilities
1629 --------------------------------------------------------------------}
1630 foldlStrict f z xs
1631   = case xs of
1632       []     -> z
1633       (x:xx) -> let z' = f z x in seq z' (foldlStrict f z' xx)
1634
1635
1636 {-
1637 {--------------------------------------------------------------------
1638   Testing
1639 --------------------------------------------------------------------}
1640 testTree xs   = fromList [(x,"*") | x <- xs]
1641 test1 = testTree [1..20]
1642 test2 = testTree [30,29..10]
1643 test3 = testTree [1,4,6,89,2323,53,43,234,5,79,12,9,24,9,8,423,8,42,4,8,9,3]
1644
1645 {--------------------------------------------------------------------
1646   QuickCheck
1647 --------------------------------------------------------------------}
1648 qcheck prop
1649   = check config prop
1650   where
1651     config = Config
1652       { configMaxTest = 500
1653       , configMaxFail = 5000
1654       , configSize    = \n -> (div n 2 + 3)
1655       , configEvery   = \n args -> let s = show n in s ++ [ '\b' | _ <- s ]
1656       }
1657
1658
1659 {--------------------------------------------------------------------
1660   Arbitrary, reasonably balanced trees
1661 --------------------------------------------------------------------}
1662 instance (Enum k,Arbitrary a) => Arbitrary (Map k a) where
1663   arbitrary = sized (arbtree 0 maxkey)
1664             where maxkey  = 10000
1665
1666 arbtree :: (Enum k,Arbitrary a) => Int -> Int -> Int -> Gen (Map k a)
1667 arbtree lo hi n
1668   | n <= 0        = return Tip
1669   | lo >= hi      = return Tip
1670   | otherwise     = do{ x  <- arbitrary
1671                       ; i  <- choose (lo,hi)
1672                       ; m  <- choose (1,30)
1673                       ; let (ml,mr)  | m==(1::Int)= (1,2)
1674                                      | m==2       = (2,1)
1675                                      | m==3       = (1,1)
1676                                      | otherwise  = (2,2)
1677                       ; l  <- arbtree lo (i-1) (n `div` ml)
1678                       ; r  <- arbtree (i+1) hi (n `div` mr)
1679                       ; return (bin (toEnum i) x l r)
1680                       }
1681
1682
1683 {--------------------------------------------------------------------
1684   Valid tree's
1685 --------------------------------------------------------------------}
1686 forValid :: (Show k,Enum k,Show a,Arbitrary a,Testable b) => (Map k a -> b) -> Property
1687 forValid f
1688   = forAll arbitrary $ \t ->
1689 --    classify (balanced t) "balanced" $
1690     classify (size t == 0) "empty" $
1691     classify (size t > 0  && size t <= 10) "small" $
1692     classify (size t > 10 && size t <= 64) "medium" $
1693     classify (size t > 64) "large" $
1694     balanced t ==> f t
1695
1696 forValidIntTree :: Testable a => (Map Int Int -> a) -> Property
1697 forValidIntTree f
1698   = forValid f
1699
1700 forValidUnitTree :: Testable a => (Map Int () -> a) -> Property
1701 forValidUnitTree f
1702   = forValid f
1703
1704
1705 prop_Valid
1706   = forValidUnitTree $ \t -> valid t
1707
1708 {--------------------------------------------------------------------
1709   Single, Insert, Delete
1710 --------------------------------------------------------------------}
1711 prop_Single :: Int -> Int -> Bool
1712 prop_Single k x
1713   = (insert k x empty == singleton k x)
1714
1715 prop_InsertValid :: Int -> Property
1716 prop_InsertValid k
1717   = forValidUnitTree $ \t -> valid (insert k () t)
1718
1719 prop_InsertDelete :: Int -> Map Int () -> Property
1720 prop_InsertDelete k t
1721   = (lookup k t == Nothing) ==> delete k (insert k () t) == t
1722
1723 prop_DeleteValid :: Int -> Property
1724 prop_DeleteValid k
1725   = forValidUnitTree $ \t ->
1726     valid (delete k (insert k () t))
1727
1728 {--------------------------------------------------------------------
1729   Balance
1730 --------------------------------------------------------------------}
1731 prop_Join :: Int -> Property
1732 prop_Join k
1733   = forValidUnitTree $ \t ->
1734     let (l,r) = split k t
1735     in valid (join k () l r)
1736
1737 prop_Merge :: Int -> Property
1738 prop_Merge k
1739   = forValidUnitTree $ \t ->
1740     let (l,r) = split k t
1741     in valid (merge l r)
1742
1743
1744 {--------------------------------------------------------------------
1745   Union
1746 --------------------------------------------------------------------}
1747 prop_UnionValid :: Property
1748 prop_UnionValid
1749   = forValidUnitTree $ \t1 ->
1750     forValidUnitTree $ \t2 ->
1751     valid (union t1 t2)
1752
1753 prop_UnionInsert :: Int -> Int -> Map Int Int -> Bool
1754 prop_UnionInsert k x t
1755   = union (singleton k x) t == insert k x t
1756
1757 prop_UnionAssoc :: Map Int Int -> Map Int Int -> Map Int Int -> Bool
1758 prop_UnionAssoc t1 t2 t3
1759   = union t1 (union t2 t3) == union (union t1 t2) t3
1760
1761 prop_UnionComm :: Map Int Int -> Map Int Int -> Bool
1762 prop_UnionComm t1 t2
1763   = (union t1 t2 == unionWith (\x y -> y) t2 t1)
1764
1765 prop_UnionWithValid
1766   = forValidIntTree $ \t1 ->
1767     forValidIntTree $ \t2 ->
1768     valid (unionWithKey (\k x y -> x+y) t1 t2)
1769
1770 prop_UnionWith :: [(Int,Int)] -> [(Int,Int)] -> Bool
1771 prop_UnionWith xs ys
1772   = sum (elems (unionWith (+) (fromListWith (+) xs) (fromListWith (+) ys)))
1773     == (sum (Prelude.map snd xs) + sum (Prelude.map snd ys))
1774
1775 prop_DiffValid
1776   = forValidUnitTree $ \t1 ->
1777     forValidUnitTree $ \t2 ->
1778     valid (difference t1 t2)
1779
1780 prop_Diff :: [(Int,Int)] -> [(Int,Int)] -> Bool
1781 prop_Diff xs ys
1782   =  List.sort (keys (difference (fromListWith (+) xs) (fromListWith (+) ys)))
1783     == List.sort ((List.\\) (nub (Prelude.map fst xs))  (nub (Prelude.map fst ys)))
1784
1785 prop_IntValid
1786   = forValidUnitTree $ \t1 ->
1787     forValidUnitTree $ \t2 ->
1788     valid (intersection t1 t2)
1789
1790 prop_Int :: [(Int,Int)] -> [(Int,Int)] -> Bool
1791 prop_Int xs ys
1792   =  List.sort (keys (intersection (fromListWith (+) xs) (fromListWith (+) ys)))
1793     == List.sort (nub ((List.intersect) (Prelude.map fst xs)  (Prelude.map fst ys)))
1794
1795 {--------------------------------------------------------------------
1796   Lists
1797 --------------------------------------------------------------------}
1798 prop_Ordered
1799   = forAll (choose (5,100)) $ \n ->
1800     let xs = [(x,()) | x <- [0..n::Int]]
1801     in fromAscList xs == fromList xs
1802
1803 prop_List :: [Int] -> Bool
1804 prop_List xs
1805   = (sort (nub xs) == [x | (x,()) <- toList (fromList [(x,()) | x <- xs])])
1806 -}