0 | ||| Finite Maps
   1 | module Data.Map
   2 |
   3 | import public Data.Map.Internal
   4 |
   5 | import Data.Bits
   6 | import Data.List
   7 |
   8 | %hide Prelude.null
   9 | %hide Prelude.toList
  10 |
  11 | %default total
  12 |
  13 | --------------------------------------------------------------------------------
  14 | --          Creating Maps
  15 | --------------------------------------------------------------------------------
  16 |
  17 | ||| The empty map. O(1)
  18 | export
  19 | empty : Map k v
  20 | empty = Tip
  21 |
  22 | --------------------------------------------------------------------------------
  23 | --          Folds
  24 | --------------------------------------------------------------------------------
  25 |
  26 | ||| Fold the values in the map using the given left-associative binary operator. O(n)
  27 | export
  28 | foldl :  (v -> w -> v)
  29 |       -> v
  30 |       -> Map k w
  31 |       -> v
  32 | foldl f z Tip             =
  33 |   z
  34 | foldl f z (Bin _ _ x l r) =
  35 |   foldl f (f (foldl f z l) x) r
  36 |
  37 | ||| Fold the values in the map using the given right-associative binary operator. O(n)
  38 | export
  39 | foldr :  (v -> w -> w)
  40 |       -> w
  41 |       -> Map k v
  42 |       -> w
  43 | foldr f z Tip             =
  44 |   z
  45 | foldr f z (Bin _ _ x l r) =
  46 |   foldr f (f x (foldr f z r)) l
  47 |
  48 | ||| Fold the keys and values in the map using the given left-associative binary operator. O(n)
  49 | export
  50 | foldlWithKey :  (v -> k -> w -> v)
  51 |              -> v
  52 |              -> Map k w
  53 |              -> v
  54 | foldlWithKey f z Tip              =
  55 |   z
  56 | foldlWithKey f z (Bin _ kx x l r) =
  57 |   foldlWithKey f (f (foldlWithKey f z l) kx x) r
  58 |
  59 | ||| Fold the keys and values in the map using the given right-associative binary operator. O(n)
  60 | export
  61 | foldrWithKey :  (k -> v -> w -> w)
  62 |              -> w
  63 |              -> Map k v
  64 |              -> w
  65 | foldrWithKey f z Tip              =
  66 |   z
  67 | foldrWithKey f z (Bin _ kx x l r) =
  68 |   foldrWithKey f (f kx x (foldrWithKey f z r)) l
  69 |
  70 | --------------------------------------------------------------------------------
  71 | --          Insertion
  72 | --------------------------------------------------------------------------------
  73 |
  74 | ||| Insert a new key and value in the map.
  75 | ||| If the key is already present in the map, the associated value is
  76 | ||| replaced with the supplied value. O(log n)
  77 | export
  78 | insert :  Eq (Map k v)
  79 |        => Eq v
  80 |        => Ord k
  81 |        => k
  82 |        -> v
  83 |        -> Map k v
  84 |        -> Map k v
  85 | insert kx0 vx0 m =
  86 |   go kx0 kx0 vx0 m
  87 |   where
  88 |     go :  k
  89 |        -> k
  90 |        -> v
  91 |        -> Map k v
  92 |        -> Map k v
  93 |     go orig _  x Tip                 =
  94 |       singleton orig x
  95 |     go orig kx x t@(Bin sz ky y l r) =
  96 |       case compare kx ky of
  97 |         LT =>
  98 |           let l' = go orig kx x l
  99 |             in case l' == l of
 100 |                  True  =>
 101 |                    t
 102 |                  False =>
 103 |                    balanceL ky y l' r
 104 |         GT =>
 105 |           let r' = go orig kx x r
 106 |             in case r' == r of
 107 |                  True  =>
 108 |                    t
 109 |                  False =>
 110 |                    balanceR ky y l r'
 111 |         EQ =>
 112 |           case (x == y && orig == ky) of
 113 |             True  =>
 114 |               t
 115 |             False =>
 116 |               Bin sz orig x l r
 117 |
 118 | private
 119 | insertR :  Eq (Map k v)
 120 |         => Eq v
 121 |         => Ord k
 122 |         => k
 123 |         -> v
 124 |         -> Map k v
 125 |         -> Map k v
 126 | insertR kx0 =
 127 |   go kx0 kx0
 128 |   where
 129 |     go :  k
 130 |        -> k
 131 |        -> v
 132 |        -> Map k v
 133 |        -> Map k v
 134 |     go orig _  x Tip                =
 135 |       singleton orig x
 136 |     go orig kx x t@(Bin _ ky y l r) =
 137 |       case compare kx ky of
 138 |         LT =>
 139 |           let l' = go orig kx x l
 140 |             in case l' == l of
 141 |                  True  =>
 142 |                    t
 143 |                  False =>
 144 |                    balanceL ky y l' r
 145 |         GT =>
 146 |           let r' = go orig kx x r
 147 |             in case r' == r of
 148 |                  True  =>
 149 |                    t
 150 |                  False =>
 151 |                    balanceR ky y l r'
 152 |         EQ =>
 153 |           t
 154 |
 155 | ||| Insert with a function, combining new value and old value.
 156 | ||| insertWith f key value mp
 157 | ||| will insert the pair (key, value) into mp if key does
 158 | ||| not exist in the map. If the key does exist, the function will
 159 | ||| insert the pair (key, f new_value old_value). O(log n)
 160 | export
 161 | insertWith :  Ord k
 162 |            => (v -> v -> v)
 163 |            -> k
 164 |            -> v
 165 |            -> Map k v
 166 |            -> Map k v
 167 | insertWith =
 168 |   go
 169 |   where
 170 |     go :  (a -> a -> a)
 171 |        -> k
 172 |        -> a
 173 |        -> Map k a
 174 |        -> Map k a
 175 |     go _ kx x Tip               =
 176 |       singleton kx x
 177 |     go f kx x (Bin sy ky y l r) =
 178 |       case compare kx ky of
 179 |         LT =>
 180 |           balanceL ky y (go f kx x l) r
 181 |         GT =>
 182 |           balanceR ky y l (go f kx x r)
 183 |         EQ =>
 184 |           Bin sy kx (f x y) l r
 185 |
 186 | private
 187 | insertWithR :  Ord k
 188 |             => (v -> v -> v)
 189 |             -> k
 190 |             -> v
 191 |             -> Map k v
 192 |             -> Map k v
 193 | insertWithR =
 194 |   go
 195 |   where
 196 |     go :  (v -> v -> v)
 197 |        -> k
 198 |        -> v
 199 |        -> Map k v
 200 |        -> Map k v
 201 |     go _ kx x Tip               =
 202 |       singleton kx x
 203 |     go f kx x (Bin sy ky y l r) =
 204 |       case compare kx ky of
 205 |         LT =>
 206 |           balanceL ky y (go f kx x l) r
 207 |         GT =>
 208 |           balanceR ky y l (go f kx x r)
 209 |         EQ =>
 210 |           Bin sy ky (f y x) l r
 211 |
 212 | private
 213 | insertWithKeyR :  Ord k
 214 |                => (k -> v -> v -> v)
 215 |                -> k
 216 |                -> v
 217 |                -> Map k v
 218 |                -> Map k v
 219 | insertWithKeyR =
 220 |   go
 221 |   where
 222 |     go :  (k -> v -> v -> v)
 223 |        -> k
 224 |        -> v
 225 |        -> Map k v
 226 |        -> Map k v
 227 |     go _ kx x Tip               =
 228 |       singleton kx x
 229 |     go f kx x (Bin sy ky y l r) =
 230 |       case compare kx ky of
 231 |         LT =>
 232 |           balanceL ky y (go f kx x l) r
 233 |         GT =>
 234 |           balanceR ky y l (go f kx x r)
 235 |         EQ =>
 236 |           Bin sy ky (f ky y x) l r
 237 |
 238 | ||| Insert with a function, combining key, new value and old value.
 239 | ||| insertWithKey f key value mp
 240 | ||| will insert the pair (key, value) into map if key does
 241 | ||| not exist in the map. If the key does exist, the function will
 242 | ||| insert the pair (key,f key new_value old_value). O(log n)
 243 | export
 244 | insertWithKey :  Ord k
 245 |               => (k -> v -> v -> v)
 246 |               -> k
 247 |               -> v
 248 |               -> Map k v
 249 |               -> Map k v
 250 | insertWithKey =
 251 |   go
 252 |   where
 253 |     go :  (k -> v -> v -> v)
 254 |        -> k
 255 |        -> v
 256 |        -> Map k v
 257 |        -> Map k v
 258 |     go _ kx x Tip               =
 259 |       singleton kx x
 260 |     go f kx x (Bin sy ky y l r) =
 261 |       case compare kx ky of
 262 |         LT =>
 263 |           balanceL ky y (go f kx x l) r
 264 |         GT =>
 265 |           balanceR ky y l (go f kx x r)
 266 |         EQ =>
 267 |           Bin sy kx (f kx x y) l r
 268 |
 269 | ||| Combines insert operation with old value retrieval.
 270 | ||| The expression (insertLookupWithKey f k x map)
 271 | ||| is a pair where the first element is equal to (lookup k map)
 272 | ||| and the second element equal to (insertWithKey f k x map). O(log n)
 273 | export
 274 | insertLookupWithKey :  Ord k
 275 |                     => (k -> v -> v -> v)
 276 |                     -> k
 277 |                     -> v
 278 |                     -> Map k v
 279 |                     -> (Maybe v, Map k v)
 280 | insertLookupWithKey f0 k0 x0 =
 281 |   go f0 k0 x0
 282 |   where
 283 |     go :  (k -> a -> a -> a)
 284 |        -> k
 285 |        -> a
 286 |        -> Map k a
 287 |        -> (Maybe a, Map k a)
 288 |     go _ kx x Tip               =
 289 |       (Nothing, singleton kx x)
 290 |     go f kx x (Bin sy ky y l r) =
 291 |       case compare kx ky of
 292 |         LT =>
 293 |           let (found, l') = go f kx x l
 294 |               t'          = balanceL ky y l' r
 295 |             in (found, t')
 296 |         GT =>
 297 |           let (found, r') = go f kx x r
 298 |               t'          = balanceR ky y l r'
 299 |             in (found, t')
 300 |         EQ =>
 301 |           (Just y, Bin sy kx (f kx x y) l r)
 302 |
 303 | --------------------------------------------------------------------------------
 304 | --          Deletion/Update
 305 | --------------------------------------------------------------------------------
 306 |
 307 | ||| Delete a key and its value from the map. When the key is not
 308 | ||| a member of the map, the original map is returned. O(log n)
 309 | export
 310 | delete :  Eq (Map k v)
 311 |        => Eq k
 312 |        => Ord k
 313 |        => k
 314 |        -> Map k v
 315 |        -> Map k v
 316 | delete =
 317 |   go
 318 |   where
 319 |     go :  k
 320 |        -> Map k v
 321 |        -> Map k v
 322 |     go _ Tip                =
 323 |       Tip
 324 |     go k t@(Bin _ kx x l r) =
 325 |       case compare k kx of
 326 |         LT =>
 327 |           let l' = go k l
 328 |             in case l' == l of
 329 |                  True  =>
 330 |                    t
 331 |                  False =>
 332 |                    balanceR kx x l' r
 333 |         GT =>
 334 |           let r' = go k r
 335 |             in case r' == r of
 336 |                  True  =>
 337 |                    t
 338 |                  False =>
 339 |                    balanceL kx x l r'
 340 |         EQ =>
 341 |           glue l r
 342 |
 343 | ||| Adjust a value at a specific key. When the key is not
 344 | ||| a member of the map, the original map is returned. O(log n)
 345 | export
 346 | adjustWithKey :  Ord k
 347 |               => (k -> v -> v)
 348 |               -> k
 349 |               -> Map k v
 350 |               -> Map k v
 351 | adjustWithKey =
 352 |   go
 353 |   where
 354 |     go :  (k -> v -> v)
 355 |        -> k
 356 |        -> Map k v
 357 |        -> Map k v
 358 |     go _ _ Tip               =
 359 |       Tip
 360 |     go f k (Bin sx kx x l r) =
 361 |       case compare k kx of
 362 |         LT =>
 363 |           Bin sx kx x (go f k l) r
 364 |         GT =>
 365 |           Bin sx kx x l (go f k r)
 366 |         EQ =>
 367 |           Bin sx kx (f kx x) l r
 368 |
 369 | ||| Update a value at a specific key with the result of the provided function.
 370 | ||| When the key is not a member of the map, the original map is returned. O(log n)
 371 | export
 372 | adjust :  Ord k
 373 |        => (v -> v)
 374 |        -> k
 375 |        -> Map k v
 376 |        -> Map k v
 377 | adjust f =
 378 |   adjustWithKey (\_, x => f x)
 379 |
 380 | ||| The expression (updateWithKey f k map) updates the
 381 | ||| value x at k (if it is in the map). If (f k x) is Nothing,
 382 | ||| the element is deleted. If it is (Just y), the key k is bound
 383 | ||| to the new value y. O(log n)
 384 | export
 385 | updateWithKey :  Ord k
 386 |               => (k -> v -> Maybe v)
 387 |               -> k
 388 |               -> Map k v
 389 |               -> Map k v
 390 | updateWithKey =
 391 |   go
 392 |   where
 393 |     go :  (k -> v -> Maybe v)
 394 |        -> k
 395 |        -> Map k v
 396 |        -> Map k v
 397 |     go _ _ Tip               =
 398 |       Tip
 399 |     go f k (Bin sx kx x l r) =
 400 |       case compare k kx of
 401 |         LT =>
 402 |           balanceR kx x (go f k l) r
 403 |         GT =>
 404 |           balanceL kx x l (go f k r)
 405 |         EQ =>
 406 |           case f kx x of
 407 |             Just x' =>
 408 |               Bin sx kx x' l r
 409 |             Nothing =>
 410 |               glue l r
 411 |
 412 | ||| The expression (update f k map) updates the value x
 413 | ||| at k (if it is in the map). If (f x) is Nothing, the element is
 414 | ||| deleted. If it is (Just y), the key k is bound to the new value y. O(log n)
 415 | export
 416 | update :  Ord k
 417 |        => (v -> Maybe v)
 418 |        -> k
 419 |        -> Map k v
 420 |        -> Map k v
 421 | update f =
 422 |   updateWithKey (\_, x => f x)
 423 |
 424 | ||| Lookup and update. See also updateWithKey.
 425 | ||| The function returns changed value, if it is updated.
 426 | ||| Returns the original key value if the map entry is deleted. O(log n)
 427 | export
 428 | updateLookupWithKey :  Ord k
 429 |                     => (k -> v -> Maybe v)
 430 |                     -> k
 431 |                     -> Map k v
 432 |                     -> (Maybe v, Map k v)
 433 | updateLookupWithKey f0 k0 =
 434 |   go f0 k0
 435 |   where
 436 |     go :  (k -> v -> Maybe v)
 437 |        -> k
 438 |        -> Map k v
 439 |        -> (Maybe v, Map k v)
 440 |     go _ _ Tip               =
 441 |       (Nothing, Tip)
 442 |     go f k (Bin sx kx x l r) =
 443 |       case compare k kx of
 444 |         LT =>
 445 |           let (found, l') = go f k l
 446 |               t'          = balanceR kx x l' r
 447 |             in (found, t')
 448 |         GT =>
 449 |           let (found, r') = go f k r
 450 |               t'          = balanceL kx x l r'
 451 |             in (found, t')
 452 |         EQ =>
 453 |           case f kx x of
 454 |             Just x' =>
 455 |               (Just x', Bin sx kx x' l r)
 456 |             Nothing =>
 457 |               let glued = glue l r
 458 |                 in (Just x, glued)
 459 |
 460 | ||| The expression (alter f k map) alters the value x at k, or absence thereof.
 461 | ||| alter can be used to insert, delete, or update a value in a Map. O(log n)
 462 | export
 463 | alter :  Ord k
 464 |       => (Maybe v -> Maybe v)
 465 |       -> k
 466 |       -> Map k v
 467 |       -> Map k v
 468 | alter =
 469 |   go
 470 |   where
 471 |     go :  (Maybe v -> Maybe v)
 472 |        -> k
 473 |        -> Map k v
 474 |        -> Map k v
 475 |     go f k Tip               =
 476 |       case f Nothing of
 477 |         Nothing =>
 478 |           Tip
 479 |         Just x  =>
 480 |           singleton k x
 481 |     go f k (Bin sx kx x l r) =
 482 |       case compare k kx of
 483 |         LT =>
 484 |           balance kx x (go f k l) r
 485 |         GT =>
 486 |           balance kx x l (go f k r)
 487 |         EQ =>
 488 |           case f (Just x) of
 489 |             Just x' =>
 490 |               Bin sx kx x' l r
 491 |             Nothing =>
 492 |               glue l r
 493 |
 494 | --------------------------------------------------------------------------------
 495 | --          Query
 496 | --------------------------------------------------------------------------------
 497 |
 498 | ||| Lookup the value at a key in the map.
 499 | ||| The function will return the corresponding value as (Just value),
 500 | ||| or Nothing if the key isn't in the map. O(log n)
 501 | export
 502 | lookup :  Ord k
 503 |        => k
 504 |        -> Map k v
 505 |        -> Maybe v
 506 | lookup =
 507 |   go
 508 |   where
 509 |     go :  k
 510 |        -> Map k v
 511 |        -> Maybe v
 512 |     go _ Tip              =
 513 |       Nothing
 514 |     go k (Bin _ kx x l r) =
 515 |       case compare k kx of
 516 |         LT =>
 517 |           go k l
 518 |         GT =>
 519 |           go k r
 520 |         EQ =>
 521 |           Just x
 522 |
 523 | ||| Find the value at a key.
 524 | ||| Returns Nothing when the element can not be found. O(log n)
 525 | export
 526 | (!?) :  Ord k
 527 |      => Map k v
 528 |      -> k
 529 |      -> Maybe v
 530 | (!?) m k =
 531 |   lookup k m
 532 |
 533 | ||| Is the key a member of the map? See also notMember. O(log n)
 534 | export
 535 | member :  Ord k
 536 |        => k
 537 |        -> Map k v
 538 |        -> Bool
 539 | member _ Tip              =
 540 |   False
 541 | member k (Bin _ kx _ l r) =
 542 |   case compare k kx of
 543 |     LT =>
 544 |       member k l
 545 |     GT =>
 546 |       member k r
 547 |     EQ =>
 548 |       True
 549 |
 550 | ||| Is the key not a member of the map? See also member. O(log n)
 551 | export
 552 | notMember :  Ord k
 553 |           => k
 554 |           -> Map k v
 555 |           -> Bool
 556 | notMember k m =
 557 |   not $
 558 |     member k m
 559 |
 560 | ||| Find the value at a key.
 561 | ||| Calls idris_crash when the element can not be found. O(log n)
 562 | export
 563 | find :  Ord k
 564 |      => k
 565 |      -> Map k v
 566 |      -> v
 567 | find _ Tip              =
 568 |   assert_total $ idris_crash "Data.Map.find: given key is not an element in the map"
 569 | find k (Bin _ kx x l r) =
 570 |   case compare k kx of
 571 |     LT =>
 572 |       find k l
 573 |     GT =>
 574 |       find k r
 575 |     EQ =>
 576 |       x
 577 |
 578 | ||| Find the value at a key.
 579 | ||| Calls idris_crash when the element can not be found. O(log n)
 580 | export
 581 | (!!) :  Ord k
 582 |      => Map k v
 583 |      -> k
 584 |      -> v
 585 | (!!) m k =
 586 |   find k m
 587 |
 588 | ||| The expression (findWithDefault def k map) returns
 589 | ||| the value at key k or returns default value def
 590 | ||| when the key is not in the map. O(log n)
 591 | export
 592 | findWithDefault :  Ord k
 593 |                 => v
 594 |                 -> k
 595 |                 -> Map k v
 596 |                 -> v
 597 | findWithDefault def _ Tip              =
 598 |   def
 599 | findWithDefault def k (Bin _ kx x l r) =
 600 |   case compare k kx of
 601 |     LT =>
 602 |       findWithDefault def k l
 603 |     GT =>
 604 |       findWithDefault def k r
 605 |     EQ =>
 606 |       x
 607 |
 608 | ||| Find largest key smaller than the given one and return the
 609 | ||| corresponding (key, value) pair. O(log n)
 610 | export
 611 | lookupLT :  Ord k
 612 |          => k
 613 |          -> Map k v
 614 |          -> Maybe (k, v)
 615 | lookupLT =
 616 |   goNothing
 617 |   where
 618 |     goJust :  k
 619 |            -> k
 620 |            -> v
 621 |            -> Map k v
 622 |            -> Maybe (k, v)
 623 |     goJust _ kx' x' Tip              =
 624 |       Just (kx', x')
 625 |     goJust k kx' x' (Bin _ kx x l r) =
 626 |       case k <= kx of
 627 |         True  =>
 628 |           goJust k kx' x' l
 629 |         False =>
 630 |           goJust k kx x r
 631 |     goNothing :  k
 632 |               -> Map k v
 633 |               -> Maybe (k, v)
 634 |     goNothing _ Tip              =
 635 |       Nothing
 636 |     goNothing k (Bin _ kx x l r) =
 637 |       case k <= kx of
 638 |         True  =>
 639 |           goNothing k l
 640 |         False =>
 641 |           goJust k kx x r
 642 |
 643 | ||| Find smallest key greater than the given one and return the
 644 | ||| corresponding (key, value) pair. O(log n)
 645 | export
 646 | lookupGT :  Ord k
 647 |          => k
 648 |          -> Map k v
 649 |          -> Maybe (k, v)
 650 | lookupGT =
 651 |   goNothing
 652 |   where
 653 |     goJust :  k
 654 |            -> k
 655 |            -> v
 656 |            -> Map k v
 657 |            -> Maybe (k, v)
 658 |     goJust _ kx' x' Tip              =
 659 |       Just (kx',x')
 660 |     goJust k kx' x' (Bin _ kx x l r) =
 661 |       case k < kx of
 662 |         True  =>
 663 |           goJust k kx x l
 664 |         False =>
 665 |           goJust k kx' x' r
 666 |     goNothing :  k
 667 |               -> Map k v
 668 |               -> Maybe (k, v)
 669 |     goNothing _ Tip              =
 670 |       Nothing
 671 |     goNothing k (Bin _ kx x l r) =
 672 |       case k < kx of
 673 |         True  =>
 674 |           goJust k kx x l
 675 |         False =>
 676 |           goNothing k r
 677 |
 678 | ||| Find largest key smaller or equal to the given one and return
 679 | ||| the corresponding (key, value) pair. O(log n)
 680 | export
 681 | lookupLE :  Ord k
 682 |          => k
 683 |          -> Map k v
 684 |          -> Maybe (k, v)
 685 | lookupLE =
 686 |   goNothing
 687 |   where
 688 |     goJust :  k
 689 |            -> k
 690 |            -> v
 691 |            -> Map k v
 692 |            -> Maybe (k, v)
 693 |     goJust _ kx' x' Tip              =
 694 |       Just (kx', x')
 695 |     goJust k kx' x' (Bin _ kx x l r) =
 696 |       case compare k kx of
 697 |         LT =>
 698 |           goJust k kx' x' l
 699 |         EQ =>
 700 |           Just (kx,x)
 701 |         GT =>
 702 |           goJust k kx x r
 703 |     goNothing :  k
 704 |               -> Map k v
 705 |               -> Maybe (k, v)
 706 |     goNothing _ Tip              =
 707 |       Nothing
 708 |     goNothing k (Bin _ kx x l r) =
 709 |       case compare k kx of
 710 |         LT =>
 711 |           goNothing k l
 712 |         EQ =>
 713 |           Just (kx,x)
 714 |         GT =>
 715 |           goJust k kx x r
 716 |
 717 | ||| Find smallest key greater or equal to the given one and return
 718 | ||| the corresponding (key, value) pair. O(log n)
 719 | export
 720 | lookupGE :  Ord k
 721 |          => k
 722 |          -> Map k v
 723 |          -> Maybe (k, v)
 724 | lookupGE =
 725 |   goNothing
 726 |   where
 727 |     goJust :  k
 728 |            -> k
 729 |            -> v
 730 |            -> Map k v
 731 |            -> Maybe (k, v)
 732 |     goJust _ kx' x' Tip              =
 733 |       Just (kx', x')
 734 |     goJust k kx' x' (Bin _ kx x l r) =
 735 |       case compare k kx of
 736 |         LT =>
 737 |           goJust k kx x l
 738 |         EQ =>
 739 |           Just (kx, x)
 740 |         GT =>
 741 |           goJust k kx' x' r
 742 |     goNothing :  k
 743 |               -> Map k v
 744 |               -> Maybe (k, v)
 745 |     goNothing _ Tip              =
 746 |       Nothing
 747 |     goNothing k (Bin _ kx x l r) =
 748 |       case compare k kx of
 749 |         LT =>
 750 |           goJust k kx x l
 751 |         EQ =>
 752 |           Just (kx, x)
 753 |         GT =>
 754 |           goNothing k r
 755 |
 756 | --------------------------------------------------------------------------------
 757 | --          Size
 758 | --------------------------------------------------------------------------------
 759 |
 760 | ||| Is the map empty? O(1)
 761 | export
 762 | null :  Map k v
 763 |      -> Bool
 764 | null Tip =
 765 |   True
 766 | null _   =
 767 |   False
 768 |
 769 | --------------------------------------------------------------------------------
 770 | --          Filter
 771 | --------------------------------------------------------------------------------
 772 |
 773 | private
 774 | splitMember :  Ord k
 775 |             => k
 776 |             -> Map k v
 777 |             -> (Map k v, Bool, Map k v)
 778 | splitMember k0 m =
 779 |   go k0 m
 780 |   where
 781 |     go :  k
 782 |        -> Map k v
 783 |        -> (Map k v, Bool, Map k v)
 784 |     go k t =
 785 |       case t of
 786 |         Tip            =>
 787 |           (Tip, False, Tip)
 788 |         Bin _ kx x l r =>
 789 |           case compare k kx of
 790 |             LT =>
 791 |               let (lt, z, gt) = go k l
 792 |                 in (lt, z, link kx x gt r)
 793 |             GT =>
 794 |               let (lt, z, gt) = go k r
 795 |                 in (link kx x l lt, z, gt)
 796 |             EQ =>
 797 |               (l, True, r)
 798 |
 799 | ||| Decompose a map into pieces based on the structure of the underlying tree.
 800 | ||| This function is useful for consuming a map in parallel. O(1)
 801 | export
 802 | splitRoot :  Map k v
 803 |           -> List (Map k v)
 804 | splitRoot orig =
 805 |   case orig of
 806 |     Tip           =>
 807 |       []
 808 |     Bin _ k v l r =>
 809 |       [l, singleton k v, r]
 810 |
 811 | ||| The expression (splitLookup k map) splits a map just
 812 | ||| like split but also returns lookup k map. O(log n)
 813 | export
 814 | splitLookup :  Ord k
 815 |             => k
 816 |             -> Map k v
 817 |             -> (Map k v, Maybe v, Map k v)
 818 | splitLookup k0 m =
 819 |   go k0 m
 820 |   where
 821 |     go :  k
 822 |        -> Map k v
 823 |        -> (Map k v, Maybe v, Map k v)
 824 |     go k t =
 825 |       case t of
 826 |         Tip            =>
 827 |           (Tip, Nothing, Tip)
 828 |         Bin _ kx x l r =>
 829 |           case compare k kx of
 830 |           LT =>
 831 |             let (lt,z,gt) = go k l
 832 |               in (lt,z,link kx x gt r)
 833 |           GT =>
 834 |             let (lt,z,gt) = go k r
 835 |               in (link kx x l lt, z, gt)
 836 |           EQ =>
 837 |             (l, Just x, r)
 838 |
 839 | ||| The expression (split k map) is a pair (map1,map2) where
 840 | ||| the keys in map1 are smaller than k and the keys in map2 larger than k.
 841 | ||| Any key equal to k is found in neither map1 nor map2. O(log n)
 842 | export
 843 | split :  Ord k
 844 |       => k
 845 |       -> Map k v
 846 |       -> (Map k v, Map k v)
 847 | split k0 t0 =
 848 |   go k0 t0
 849 |   where
 850 |     go :  k
 851 |        -> Map k v
 852 |        -> (Map k v, Map k v)
 853 |     go k t =
 854 |       case t of
 855 |         Tip            =>
 856 |           (Tip, Tip)
 857 |         Bin _ kx x l r =>
 858 |           case compare k kx of
 859 |             LT =>
 860 |               let (lt, gt) = go k l
 861 |                 in (lt, link kx x gt r)
 862 |             GT =>
 863 |               let (lt, gt) = go k r
 864 |                 in (link kx x l lt, gt)
 865 |             EQ =>
 866 |               (l, r)
 867 |
 868 | ||| Filter all keys/values that satisfy the predicate. O(n)
 869 | export
 870 | filterWithKey :  Eq (Map k v)
 871 |               => (k -> v -> Bool)
 872 |               -> Map k v
 873 |               -> Map k v
 874 | filterWithKey _ Tip                =
 875 |   Tip
 876 | filterWithKey p t@(Bin _ kx x l r) =
 877 |   case p kx x of
 878 |     True  =>
 879 |       case (filterWithKey p l) == l && (filterWithKey p r) == r of
 880 |         True  =>
 881 |           t
 882 |         False =>
 883 |           link kx x (filterWithKey p l) (filterWithKey p r)
 884 |     False =>
 885 |       link2 (filterWithKey p l) (filterWithKey p r)
 886 |
 887 | ||| Filter all values that satisfy the predicate. O(n)
 888 | export
 889 | filter :  Eq (Map k v)
 890 |        => (v -> Bool)
 891 |        -> Map k v
 892 |        -> Map k v
 893 | filter p m =
 894 |   filterWithKey (\_, x => p x) m
 895 |
 896 | ||| Partition the map according to a predicate. The first
 897 | ||| map contains all elements that satisfy the predicate, the second all
 898 | ||| elements that fail the predicate. See also split. O(n)
 899 | export
 900 | partitionWithKey :  Eq (Map k v)
 901 |                  => (k -> v -> Bool)
 902 |                  -> Map k v
 903 |                  -> (Map k v, Map k v)
 904 | partitionWithKey p0 t0 =
 905 |   go p0 t0
 906 |   where
 907 |     go :  (k -> v -> Bool)
 908 |        -> Map k v
 909 |        -> (Map k v, Map k v)
 910 |     go _ Tip                =
 911 |       (Tip, Tip)
 912 |     go p t@(Bin _ kx x l r) =
 913 |       case p kx x of
 914 |         True  =>
 915 |           ( case (fst $ go p l) == l && (fst $ go p r) == r of
 916 |              True  =>
 917 |                t
 918 |              False =>
 919 |                link kx x (fst $ go p l) (fst $ go p r)
 920 |           , link2 (snd $ go p l) (snd $ go p r)
 921 |           )
 922 |         False =>
 923 |           ( link2 (fst $ go p l) (fst $ go p r)
 924 |           , case (snd $ go p l) == l && (snd $ go p r) == r of
 925 |               True  =>
 926 |                 t
 927 |               False =>
 928 |                 link kx x (snd $ go p l) (snd $ go p r)
 929 |           )
 930 |
 931 | ||| Take while a predicate on the keys holds.
 932 | ||| The user is responsible for ensuring that for all keys j and k in the map,
 933 | ||| j < k ==> p j >= p k. See note at spanAntitone. O(log n)
 934 | export
 935 | takeWhileAntitone :  (k -> Bool)
 936 |                   -> Map k v
 937 |                   -> Map k v
 938 | takeWhileAntitone _ Tip              =
 939 |   Tip
 940 | takeWhileAntitone p (Bin _ kx x l r) =
 941 |   case p kx of
 942 |     True  =>
 943 |       link kx x l (takeWhileAntitone p r)
 944 |     False =>
 945 |       takeWhileAntitone p l
 946 |
 947 | ||| Drop while a predicate on the keys holds.
 948 | ||| The user is responsible for ensuring that for all keys j and k in the map,
 949 | ||| j < k ==> p j >= p k. See note at spanAntitone. O(log n)
 950 | export
 951 | dropWhileAntitone :  (k -> Bool)
 952 |                   -> Map k v
 953 |                   -> Map k v
 954 | dropWhileAntitone _ Tip              =
 955 |   Tip
 956 | dropWhileAntitone p (Bin _ kx x l r) =
 957 |   case p kx of
 958 |     True  =>
 959 |       dropWhileAntitone p r
 960 |     False =>
 961 |       link kx x (dropWhileAntitone p l) r
 962 |
 963 | ||| Divide a map at the point where a predicate on the keys stops holding.
 964 | ||| The user is responsible for ensuring that for all keys j and k in the map,
 965 | ||| j < k ==> p j>= p k. O(log n)
 966 | export
 967 | spanAntitone :  (k -> Bool)
 968 |              -> Map k v
 969 |              -> (Map k v, Map k v)
 970 | spanAntitone p0 m =
 971 |   go p0 m
 972 |   where
 973 |     go :  (k -> Bool)
 974 |        -> Map k v
 975 |        -> (Map k v, Map k v)
 976 |     go _ Tip              =
 977 |       (Tip, Tip)
 978 |     go p (Bin _ kx x l r) =
 979 |       case p kx of
 980 |         True  =>
 981 |           let (u, v) = go p r
 982 |             in (link kx x l u, v)
 983 |         False =>
 984 |           let (u, v) = go p l
 985 |             in (u, link kx x v r)
 986 |
 987 | ||| Map keys/values and collect the Just results. O(n)
 988 | export
 989 | mapMaybeWithKey :  (k -> a -> Maybe b)
 990 |                 -> Map k a
 991 |                 -> Map k b
 992 | mapMaybeWithKey _ Tip              =
 993 |   Tip
 994 | mapMaybeWithKey f (Bin _ kx x l r) =
 995 |   case f kx x of
 996 |     Just y  =>
 997 |       link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
 998 |     Nothing =>
 999 |       link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r)
1000 |
1001 | ||| Map values and collect the Just results. O(n)
1002 | export
1003 | mapMaybe :  (a -> Maybe b)
1004 |          -> Map k a
1005 |          -> Map k b
1006 | mapMaybe f =
1007 |   mapMaybeWithKey (\_, x => f x)
1008 |
1009 | ||| Map keys/values and separate the Left and Right results. O(n)
1010 | export
1011 | mapEitherWithKey :  (k -> a -> Either b c)
1012 |                  -> Map k a
1013 |                  -> (Map k b, Map k c)
1014 | mapEitherWithKey f0 t0 =
1015 |   go f0 t0
1016 |   where
1017 |     go :  (k -> a -> Either b c)
1018 |        -> Map k a
1019 |        -> (Map k b, Map k c)
1020 |     go _ Tip              =
1021 |       (Tip, Tip)
1022 |     go f (Bin _ kx x l r) =
1023 |       case f kx x of
1024 |         Left  y =>
1025 |           (link kx y (fst $ go f l) (fst $ go f r), link2 (snd $ go f l) (snd $ go f r))
1026 |         Right z =>
1027 |           (link2 (fst $ go f l) (fst $ go f r), link kx z (snd $ go f l) (snd $ go f r))
1028 |
1029 | ||| Map values and separate the Left and Right results. O(n)
1030 | export
1031 | mapEither :  (a -> Either b c)
1032 |           -> Map k a
1033 |           -> (Map k b, Map k c)
1034 | mapEither f m =
1035 |   mapEitherWithKey (\_, x => f x) m
1036 |
1037 | --------------------------------------------------------------------------------
1038 | --          Submap
1039 | --------------------------------------------------------------------------------
1040 |
1041 | private
1042 | submap' :  Ord a
1043 |         => (b -> c -> Bool)
1044 |         -> Map a b
1045 |         -> Map a c
1046 |         -> Bool
1047 | submap' _ Tip _              =
1048 |   True
1049 | submap' _ _   Tip            =
1050 |   False
1051 | submap' f (Bin 1 kx x _ _) t =
1052 |   case lookup kx t of
1053 |     Just y  =>
1054 |       f x y
1055 |     Nothing =>
1056 |       False
1057 | submap' f (Bin _ kx x l r) t =
1058 |   let (lt,found,gt) = splitLookup kx t
1059 |     in case found of
1060 |          Nothing =>
1061 |            False
1062 |          Just y  =>
1063 |            f x y && size l <= size lt
1064 |                  && size r <= size gt
1065 |                  && submap' f l lt
1066 |                  && submap' f r gt
1067 |
1068 | ||| The expression (isSubmapOfBy f t1 t2) returns True if
1069 | ||| all keys in t1 are in tree t2, and when f returns True when
1070 | ||| applied to their respective values.
1071 | export
1072 | isSubmapOfBy :  Ord k
1073 |              => (a -> b -> Bool)
1074 |              -> Map k a
1075 |              -> Map k b
1076 |              -> Bool
1077 | isSubmapOfBy f t1 t2 =
1078 |   size t1 <= size t2 && submap' f t1 t2
1079 |
1080 | ||| This function is defined as (isSubmapOf = isSubmapOfBy (==)).
1081 | export
1082 | isSubmapOf :  Eq v
1083 |            => Ord k
1084 |            => Map k v
1085 |            -> Map k v
1086 |            -> Bool
1087 | isSubmapOf m1 m2 =
1088 |   isSubmapOfBy (==) m1 m2
1089 |
1090 | ||| Is this a proper submap? (ie. a submap but not equal).
1091 | ||| The expression (isProperSubmapOfBy f m1 m2) returns True when
1092 | ||| keys m1 and keys m2 are not equal,
1093 | ||| all keys in m1 are in m2, and when f returns True when
1094 | ||| applied to their respective values.
1095 | export
1096 | isProperSubmapOfBy :  Ord k
1097 |                    => (a -> b -> Bool)
1098 |                    -> Map k a
1099 |                    -> Map k b
1100 |                    -> Bool
1101 | isProperSubmapOfBy f t1 t2 =
1102 |   size t1 < size t2 && submap' f t1 t2
1103 |
1104 | ||| Is this a proper submap? (ie. a submap but not equal).
1105 | ||| Defined as (isProperSubmapOf = isProperSubmapOfBy (==)).
1106 | export
1107 | isProperSubmapOf :  Eq v
1108 |                  => Ord k
1109 |                  => Map k v
1110 |                  -> Map k v
1111 |                  -> Bool
1112 | isProperSubmapOf m1 m2 =
1113 |   isProperSubmapOfBy (==) m1 m2
1114 |
1115 | --------------------------------------------------------------------------------
1116 | --          Indexed
1117 | --------------------------------------------------------------------------------
1118 |
1119 | ||| Lookup the index of a key, which is its zero-based index in
1120 | ||| the sequence sorted by keys. The index is a number from 0 up to, but not
1121 | ||| including, the size of the map. O(log n)
1122 | export
1123 | lookupIndex :  Ord k
1124 |             => k
1125 |             -> Map k v
1126 |             -> Maybe Nat
1127 | lookupIndex =
1128 |   go 0
1129 |   where
1130 |     go :  Nat
1131 |        -> k
1132 |        -> Map k a
1133 |        -> Maybe Nat
1134 |     go _  _ Tip               =
1135 |       Nothing
1136 |     go idx k (Bin _ kx _ l r) =
1137 |       case compare k kx of
1138 |         LT =>
1139 |           go idx k l
1140 |         GT =>
1141 |           go (idx + size l + 1) k r
1142 |         EQ =>
1143 |           Just $ idx + size l
1144 |
1145 | ||| Return the index of a key, which is its zero-based index in
1146 | ||| the sequence sorted by keys. The index is a number from 0 up to, but not
1147 | ||| including, the size of the map. Calls idris_crash when the key is not
1148 | ||| a member of the map. O(log n)
1149 | export
1150 | findIndex :  Ord k
1151 |           => k
1152 |           -> Map k v
1153 |           -> Nat
1154 | findIndex =
1155 |   go 0
1156 |   where
1157 |     go :  Nat
1158 |        -> k
1159 |        -> Map k a
1160 |        -> Nat
1161 |     go _   _ Tip              =
1162 |       assert_total $ idris_crash "Data.Map.findIndex: element is not in the map"
1163 |     go idx k (Bin _ kx _ l r) =
1164 |       case compare k kx of
1165 |         LT =>
1166 |           go idx k l
1167 |         GT =>
1168 |           go (idx + size l + 1) k r
1169 |         EQ =>
1170 |           idx + size l
1171 |
1172 | ||| Retrieve an element by its index, i.e. by its zero-based
1173 | ||| index in the sequence sorted by keys. If the index is out of range (less
1174 | ||| than zero, greater or equal to size of the map), idris_crash is called. O(log n)
1175 | export
1176 | elemAt :  Nat
1177 |        -> Map k v
1178 |        -> (k, v)
1179 | elemAt _ Tip              =
1180 |   assert_total $ idris_crash "Data.Map.elemAt: index out of range"
1181 | elemAt i (Bin _ kx x l r) =
1182 |   case compare i (size l) of
1183 |      LT =>
1184 |        elemAt i l
1185 |      GT =>
1186 |        elemAt ((i `minus` (size l)) `minus` 1) r
1187 |      EQ =>
1188 |        (kx,x)
1189 |
1190 | ||| Update the element at index, i.e. by its zero-based index in
1191 | ||| the sequence sorted by keys. If the index is out of range (less than zero,
1192 | ||| greater or equal to size of the map), idris_crash is called. O(log n)
1193 | export
1194 | updateAt :  (k -> v -> Maybe v)
1195 |          -> Nat
1196 |          -> Map k v
1197 |          -> Map k v
1198 | updateAt f i t =
1199 |   case t of
1200 |     Tip             =>
1201 |       assert_total $ idris_crash "Data.Map.updateAt: index out of range"
1202 |     Bin sx kx x l r =>
1203 |       case compare i (size l) of
1204 |         LT =>
1205 |           balanceR kx x (updateAt f i l) r
1206 |         GT =>
1207 |           balanceL kx x l (updateAt f ((i `minus` (size l)) `minus` 1) r)
1208 |         EQ =>
1209 |           case f kx x of
1210 |             Just x' =>
1211 |               Bin sx kx x' l r
1212 |             Nothing =>
1213 |               glue l r
1214 |
1215 | ||| Delete the element at index, i.e. by its zero-based index in
1216 | ||| the sequence sorted by keys. If the index is out of range (less than zero,
1217 | ||| greater or equal to size of the map), idris_crash is called. O(log n)
1218 | export
1219 | deleteAt :  Nat
1220 |          -> Map k v
1221 |          -> Map k v
1222 | deleteAt i t =
1223 |   case t of
1224 |     Tip            =>
1225 |       assert_total $ idris_crash "Data.Map.deleteAt: index out of range"
1226 |     Bin _ kx x l r =>
1227 |       case compare i (size l) of
1228 |         LT =>
1229 |           balanceR kx x (deleteAt i l) r
1230 |         GT =>
1231 |           balanceL kx x l (deleteAt ((i `minus` (size l)) `minus` 1) r)
1232 |         EQ =>
1233 |           glue l r
1234 |
1235 | ||| Take a given number of entries in key order, beginning
1236 | ||| with the smallest keys. O(log n)
1237 | export
1238 | take :  Nat
1239 |      -> Map k v
1240 |      -> Map k v
1241 | take i m =
1242 |   case i >= size m of
1243 |     True  =>
1244 |       m
1245 |     False =>
1246 |       go i m
1247 |   where
1248 |     go :  Nat
1249 |        -> Map k v
1250 |        -> Map k v
1251 |     go _ Tip              =
1252 |       Tip
1253 |     go i (Bin _ kx x l r) =
1254 |       case i <= 0 of
1255 |         True  =>
1256 |           Tip
1257 |         False =>
1258 |           case compare i (size l) of
1259 |             LT =>
1260 |               go i l
1261 |             GT =>
1262 |               link kx x l (go ((i `minus` (size l)) `minus` 1) r)
1263 |             EQ =>
1264 |               l
1265 |
1266 | ||| Drop a given number of entries in key order, beginning
1267 | ||| with the smallest keys. O(log n)
1268 | export
1269 | drop :  Nat
1270 |      -> Map k v
1271 |      -> Map k v
1272 | drop i m =
1273 |   case i >= size m of
1274 |     True  =>
1275 |       Tip
1276 |     False =>
1277 |       go i m
1278 |   where
1279 |     go :  Nat
1280 |        -> Map k v
1281 |        -> Map k v
1282 |     go _ Tip              =
1283 |       Tip
1284 |     go i (Bin _ kx x l r) =
1285 |       case i <= 0 of
1286 |         True  =>
1287 |           m
1288 |         False =>
1289 |           case compare i (size l) of
1290 |             LT =>
1291 |               link kx x (go i l) r
1292 |             GT =>
1293 |               go ((i `minus` (size l)) `minus` 1) r
1294 |             EQ =>
1295 |               insertMin kx x r
1296 |
1297 | ||| Split a map at a particular index. O(log n)
1298 | export
1299 | splitAt :  Nat
1300 |         -> Map k v
1301 |         -> (Map k v, Map k v)
1302 | splitAt i m =
1303 |   case i >= size m of
1304 |     True  =>
1305 |       (m,Tip)
1306 |     False =>
1307 |       go i m
1308 |   where
1309 |     go :  Nat
1310 |        -> Map k v
1311 |        -> (Map k v, Map k v)
1312 |     go _ Tip              =
1313 |       (Tip, Tip)
1314 |     go i (Bin _ kx x l r) =
1315 |       case i <= 0 of
1316 |         True  =>
1317 |           (Tip, m)
1318 |         False =>
1319 |           case compare i (size l) of
1320 |             LT =>
1321 |               case go i l of
1322 |                 (ll, lr) =>
1323 |                   (ll, link kx x lr r)
1324 |             GT =>
1325 |               case go ((i `minus` (size l)) `minus` 1) r of
1326 |                 (rl, rr) =>
1327 |                   (link kx x l rl, rr)
1328 |             EQ =>
1329 |               (l, insertMin kx x r)
1330 |
1331 | --------------------------------------------------------------------------------
1332 | --          Min/Max
1333 | --------------------------------------------------------------------------------
1334 |
1335 | private
1336 | lookupMinSure :  k
1337 |               -> v
1338 |               -> Map k v
1339 |               -> (k, v)
1340 | lookupMinSure k v Tip             =
1341 |   (k, v)
1342 | lookupMinSure _ _ (Bin _ k v l _) =
1343 |   lookupMinSure k v l
1344 |
1345 | ||| The minimal key of the map. Returns Nothing if the map is empty. O(log n)
1346 | export
1347 | lookupMin :  Map k v
1348 |           -> Maybe (k, v)
1349 | lookupMin Tip             =
1350 |   Nothing
1351 | lookupMin (Bin _ k v l _) =
1352 |   Just $ lookupMinSure k v l
1353 |
1354 | private
1355 | lookupMaxSure :  k
1356 |               -> v
1357 |               -> Map k v
1358 |               -> (k, v)
1359 | lookupMaxSure k v Tip             =
1360 |   (k, v)
1361 | lookupMaxSure _ _ (Bin _ k v _ r) =
1362 |   lookupMaxSure k v r
1363 |
1364 | ||| The maximal key of the map. Returns Nothing if the map is empty. O(log n)
1365 | export
1366 | lookupMax :  Map k v
1367 |           -> Maybe (k, v)
1368 | lookupMax Tip             =
1369 |   Nothing
1370 | lookupMax (Bin _ k v _ r) =
1371 |   Just $ lookupMaxSure k v r
1372 |
1373 | ||| The minimal key of the map. Calls idris_crash if the map is empty. O(log n)
1374 | export
1375 | findMin :  Map k v
1376 |         -> (k, v)
1377 | findMin t =
1378 |   case lookupMin t of
1379 |     Just r  =>
1380 |       r
1381 |     Nothing =>
1382 |       assert_total $ idris_crash "Data.Map.findMin: empty map has no minimal element"
1383 |
1384 | ||| The maximal key of the map. Calls idris_crash if the map is empty. O(log n)
1385 | export
1386 | findMax :  Map k v
1387 |         -> (k, v)
1388 | findMax t =
1389 |   case lookupMax t of
1390 |     Just r  =>
1391 |       r
1392 |     Nothing =>
1393 |       assert_total $ idris_crash "Data.Map.findMax: empty map has no maximal element"
1394 |
1395 | ||| Delete the minimal key. Returns an empty map if the map is empty. O(log n)
1396 | export
1397 | deleteMin :  Map k v
1398 |           -> Map k v
1399 | deleteMin Tip                 =
1400 |   Tip
1401 | deleteMin (Bin _ _  _ Tip r)  =
1402 |   r
1403 | deleteMin (Bin _ kx x l   r)  =
1404 |   balanceR kx x (deleteMin l) r
1405 |
1406 | ||| Delete the maximal key. Returns an empty map if the map is empty. O(log n)
1407 | export
1408 | deleteMax :  Map k v
1409 |           -> Map k v
1410 | deleteMax Tip                 =
1411 |   Tip
1412 | deleteMax (Bin _ _  _ l Tip)  =
1413 |   l
1414 | deleteMax (Bin _ kx x l r)    =
1415 |   balanceL kx x l (deleteMax r)
1416 |
1417 | ||| Retrieves the minimal (key,value) pair of the map, and
1418 | ||| the map stripped of that element, or Nothing if passed an empty map. O(log n)
1419 | export
1420 | minViewWithKey :  Map k v
1421 |                -> Maybe ((k, v), Map k v)
1422 | minViewWithKey Tip             =
1423 |   Nothing
1424 | minViewWithKey (Bin _ k x l r) =
1425 |   Just $
1426 |     case minViewSure k x l r of
1427 |       MinView' km xm t =>
1428 |         ((km, xm), t)
1429 |
1430 | ||| Delete and find the minimal element. O(log n)
1431 | export
1432 | deleteFindMin :  Map k v
1433 |               -> ((k, v), Map k v)
1434 | deleteFindMin t =
1435 |   case minViewWithKey t of
1436 |     Just res =>
1437 |       res
1438 |     Nothing  =>
1439 |       (assert_total $ idris_crash "Data.Map.deleteFindMin: can not return the minimal element of an empty map", Tip)
1440 |
1441 | ||| Retrieves the maximal (key,value) pair of the map, and
1442 | ||| the map stripped of that element, or Nothing if passed an empty map. O(log n)
1443 | export
1444 | maxViewWithKey :  Map k v
1445 |                -> Maybe ((k, v), Map k v)
1446 | maxViewWithKey Tip             =
1447 |   Nothing
1448 | maxViewWithKey (Bin _ k x l r) =
1449 |   Just $
1450 |     case maxViewSure k x l r of
1451 |       MaxView' km xm t =>
1452 |         ((km, xm), t)
1453 |
1454 | ||| Delete and find the maximal element. O(log n)
1455 | export
1456 | deleteFindMax :  Map k v
1457 |               -> ((k, v), Map k v)
1458 | deleteFindMax t =
1459 |   case maxViewWithKey t of
1460 |     Just res =>
1461 |       res
1462 |     Nothing  =>
1463 |       (assert_total $ idris_crash "Data.Map.deleteFindMax: can not return the maximal element of an empty map", Tip)
1464 |
1465 | ||| Update the value at the minimal key. O(log n)
1466 | export
1467 | updateMinWithKey :  (k -> v -> Maybe v)
1468 |                  -> Map k v
1469 |                  -> Map k v
1470 | updateMinWithKey _ Tip                 =
1471 |   Tip
1472 | updateMinWithKey f (Bin sx kx x Tip r) =
1473 |   case f kx x of
1474 |     Nothing => r
1475 |     Just x' => Bin sx kx x' Tip r
1476 | updateMinWithKey f (Bin _ kx x l r)    =
1477 |   balanceR kx x (updateMinWithKey f l) r
1478 |
1479 | ||| Update the value at the minimal key. O(log n)
1480 | export
1481 | updateMin :  (v -> Maybe v)
1482 |           -> Map k v
1483 |           -> Map k v
1484 | updateMin f m =
1485 |   updateMinWithKey (\_, x => f x) m
1486 |
1487 | ||| Update the value at the maximal key. O(log n)
1488 | export
1489 | updateMaxWithKey :  (k -> v -> Maybe v)
1490 |                  -> Map k v
1491 |                  -> Map k v
1492 | updateMaxWithKey _ Tip                 =
1493 |   Tip
1494 | updateMaxWithKey f (Bin sx kx x l Tip) =
1495 |   case f kx x of
1496 |     Nothing =>
1497 |       l
1498 |     Just x' =>
1499 |       Bin sx kx x' l Tip
1500 | updateMaxWithKey f (Bin _ kx x l r)    =
1501 |   balanceL kx x l (updateMaxWithKey f r)
1502 |
1503 | ||| Update the value at the maximal key. O(log n)
1504 | export
1505 | updateMax :  (v -> Maybe v)
1506 |           -> Map k v
1507 |           -> Map k v
1508 | updateMax f m =
1509 |   updateMaxWithKey (\_, x => f x) m
1510 |
1511 | ||| Retrieves the value associated with minimal key of the
1512 | ||| map, and the map stripped of that element, or Nothing if passed an empty map. O(log n)
1513 | export
1514 | minView :  Map k v
1515 |         -> Maybe (v, Map k v)
1516 | minView t =
1517 |   case minViewWithKey t of
1518 |     Nothing           =>
1519 |       Nothing
1520 |     Just ((_, x), t') =>
1521 |       Just (x, t')
1522 |
1523 | ||| Retrieves the value associated with maximal key of the
1524 | ||| map, and the map stripped of that element, or Nothing if passed an empty map. O(log n)
1525 | export
1526 | maxView :  Map k v
1527 |         -> Maybe (v, Map k v)
1528 | maxView t =
1529 |   case maxViewWithKey t of
1530 |     Nothing           =>
1531 |       Nothing
1532 |     Just ((_, x), t') =>
1533 |       Just (x, t')
1534 |
1535 | --------------------------------------------------------------------------------
1536 | --          Combine
1537 | --------------------------------------------------------------------------------
1538 |
1539 | ||| The expression (union t1 t2) takes the left-biased union of t1 and t2.
1540 | ||| It prefers t1 when duplicate keys are encountered.
1541 | export
1542 | union :  Eq (Map k v)
1543 |       => Eq v
1544 |       => Ord k
1545 |       => Map k v
1546 |       -> Map k v
1547 |       -> Map k v
1548 | union t1                     Tip                 =
1549 |   t1
1550 | union t1                     (Bin _ k x Tip Tip) =
1551 |   insertR k x t1
1552 | union (Bin _ k x Tip Tip)    t2                  =
1553 |   insert k x t2
1554 | union Tip                    t2                  =
1555 |   t2
1556 | union t1@(Bin _ k1 x1 l1 r1) t2                  =
1557 |   let (l2, r2) = split k1 t2
1558 |     in case (union l1 l2) == l1 && (union r1 r2) == r1 of
1559 |          True  =>
1560 |            t1
1561 |          False =>
1562 |            link k1 x1 (union l1 l2) (union r1 r2)
1563 |
1564 | ||| Union with a combining function.
1565 | export
1566 | unionWith :  Ord k
1567 |           => (v -> v -> v)
1568 |           -> Map k v
1569 |           -> Map k v
1570 |           -> Map k v
1571 | unionWith _ t1                  Tip                 =
1572 |   t1
1573 | unionWith f t1                  (Bin _ k x Tip Tip) =
1574 |   insertWithR f k x t1
1575 | unionWith f (Bin _ k x Tip Tip) t2                  =
1576 |   insertWith f k x t2
1577 | unionWith _ Tip                 t2                  =
1578 |   t2
1579 | unionWith f (Bin _ k1 x1 l1 r1) t2                  =
1580 |   let (l2, mb, r2) = splitLookup k1 t2
1581 |     in case mb of
1582 |          Nothing =>
1583 |            link k1 x1 (unionWith f l1 l2) (unionWith f r1 r2)
1584 |          Just x2 =>
1585 |            link k1 (f x1 x2) (unionWith f l1 l2) (unionWith f r1 r2)
1586 |
1587 | ||| Union with a combining function.
1588 | export
1589 | unionWithKey :  Ord k
1590 |              => (k -> v -> v -> v)
1591 |              -> Map k v
1592 |              -> Map k v
1593 |              -> Map k v
1594 | unionWithKey _ t1                  Tip                 =
1595 |   t1
1596 | unionWithKey f t1                  (Bin _ k x Tip Tip) =
1597 |   insertWithKeyR f k x t1
1598 | unionWithKey f (Bin _ k x Tip Tip) t2                  =
1599 |   insertWithKey f k x t2
1600 | unionWithKey _ Tip                 t2                  =
1601 |   t2
1602 | unionWithKey f (Bin _ k1 x1 l1 r1) t2                  =
1603 |   let (l2, mb, r2) = splitLookup k1 t2
1604 |     in case mb of
1605 |          Nothing =>
1606 |            link k1 x1 (unionWithKey f l1 l2) (unionWithKey f r1 r2)
1607 |          Just x2 =>
1608 |            link k1 (f k1 x1 x2) (unionWithKey f l1 l2) (unionWithKey f r1 r2)
1609 |
1610 | ||| The union of a list of maps.
1611 | export
1612 | unions :  Eq (Map k v)
1613 |        => Eq v
1614 |        => Foldable f
1615 |        => Ord k
1616 |        => f (Map k v)
1617 |        -> Map k v
1618 | unions ts =
1619 |   foldl union empty ts
1620 |
1621 | ||| The union of a list of maps, with a combining operation.
1622 | export
1623 | unionsWith :  Foldable f
1624 |            => Ord k
1625 |            => (v -> v -> v)
1626 |            -> f (Map k v)
1627 |            -> Map k v
1628 | unionsWith f ts =
1629 |   foldl (unionWith f) empty ts
1630 |
1631 | --------------------------------------------------------------------------------
1632 | --          Difference
1633 | --------------------------------------------------------------------------------
1634 |
1635 | ||| Difference of two maps.
1636 | ||| Return elements of the first map not existing in the second map.
1637 | export
1638 | difference :  Ord k
1639 |            => Map k a
1640 |            -> Map k b
1641 |            -> Map k a
1642 | difference Tip _                =
1643 |   Tip
1644 | difference t1 Tip               =
1645 |   t1
1646 | difference t1 (Bin _ k _ l2 r2) =
1647 |   let (l1, r1) = split k t1
1648 |     in case size (difference l1 l2) + size (difference r1 r2) == size t1 of
1649 |          True  =>
1650 |            t1
1651 |          False =>
1652 |            link2 (difference l1 l2) (difference r1 r2)
1653 |
1654 | ||| Same as difference.
1655 | export
1656 | (\\) :  Ord k
1657 |      => Map k a
1658 |      -> Map k b
1659 |      -> Map k a
1660 | m1 \\ m2 =
1661 |   difference m1 m2
1662 |
1663 | --------------------------------------------------------------------------------
1664 | --          Intersection
1665 | --------------------------------------------------------------------------------
1666 |
1667 | ||| Intersection of two maps.
1668 | ||| Return data in the first map for the keys existing in both maps.
1669 | export
1670 | intersection :  Eq (Map k a)
1671 |              => Ord k
1672 |              => Map k a
1673 |              -> Map k b
1674 |              -> Map k a
1675 | intersection Tip                  _   =
1676 |   Tip
1677 | intersection _                    Tip =
1678 |   Tip
1679 | intersection t1@(Bin _ k x l1 r1) t2  =
1680 |   case splitMember k t2 of
1681 |     (l2, True, r2)  =>
1682 |       case (intersection l1 l2) == l1 && (intersection r1 r2) == r1 of
1683 |         True  =>
1684 |           t1
1685 |         False =>
1686 |           link k x (intersection l1 l2) (intersection r1 r2)
1687 |     (l2, False, r2) =>
1688 |       link2 (intersection l1 l2) (intersection r1 r2)
1689 |
1690 | ||| Intersection with a combining function.
1691 | export
1692 | intersectionWith :  Ord k
1693 |                  => (a -> b -> c)
1694 |                  -> Map k a
1695 |                  -> Map k b
1696 |                  -> Map k c
1697 | intersectionWith f Tip                _   =
1698 |   Tip
1699 | intersectionWith f _                  Tip =
1700 |   Tip
1701 | intersectionWith f (Bin _ k x1 l1 r1) t2  =
1702 |   case splitLookup k t2 of
1703 |     (l2, Just x2, r2) =>
1704 |       link k (f x1 x2) (intersectionWith f l1 l2) (intersectionWith f r1 r2)
1705 |     (l2, Nothing, r2) =>
1706 |       link2 (intersectionWith f l1 l2) (intersectionWith f r1 r2)
1707 |
1708 | ||| Intersection with a combining function.
1709 | export
1710 | intersectionWithKey :  Ord k
1711 |                     => (k -> a -> b -> c)
1712 |                     -> Map k a
1713 |                     -> Map k b
1714 |                     -> Map k c
1715 | intersectionWithKey f Tip                _   =
1716 |   Tip
1717 | intersectionWithKey f _                  Tip =
1718 |   Tip
1719 | intersectionWithKey f (Bin _ k x1 l1 r1) t2  =
1720 |   case splitLookup k t2 of
1721 |     (l2, Just x2, r2) =>
1722 |       link k (f k x1 x2) (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
1723 |     (l2, Nothing, r2) =>
1724 |       link2 (intersectionWithKey f l1 l2) (intersectionWithKey f r1 r2)
1725 |
1726 | --------------------------------------------------------------------------------
1727 | --          Disjoint
1728 | --------------------------------------------------------------------------------
1729 |
1730 | ||| Check whether the key sets of two
1731 | ||| maps are disjoint (i.e., their intersection is empty).
1732 | export
1733 | disjoint :  Ord k
1734 |          => Map k a
1735 |          -> Map k b
1736 |          -> Bool
1737 | disjoint Tip             _   =
1738 |   True
1739 | disjoint _               Tip =
1740 |   True
1741 | disjoint (Bin 1 k _ _ _) t   =
1742 |   k `notMember` t
1743 | disjoint (Bin _ k _ l r) t   =
1744 |   let (lt,found,gt) = splitMember k t
1745 |     in not found && disjoint l lt && disjoint r gt
1746 |
1747 | --------------------------------------------------------------------------------
1748 | --          Compose
1749 | --------------------------------------------------------------------------------
1750 |
1751 | ||| Relate the keys of one map to the values of
1752 | ||| the other, by using the values of the former as keys for lookups in the latter.
1753 | ||| O(n * log(m)), where m is the size of the first argument.
1754 | export
1755 | compose :  Ord b
1756 |         => Map b c
1757 |         -> Map a b
1758 |         -> Map a c
1759 | compose bc ab =
1760 |   case null bc of
1761 |     True  =>
1762 |       empty
1763 |     False =>
1764 |       mapMaybe ((!?) bc) ab
1765 |
1766 | --------------------------------------------------------------------------------
1767 | --          Traversal
1768 | --------------------------------------------------------------------------------
1769 |
1770 | ||| Map a function over all values in the map. O(n)
1771 | export
1772 | map :  (v -> w)
1773 |     -> Map k v
1774 |     -> Map k w
1775 | map _ Tip               =
1776 |   Tip
1777 | map f (Bin sx kx x l r) =
1778 |   Bin sx kx (f x) (map f l) (map f r)
1779 |
1780 | ||| Map a function over all values in the map. O(n)
1781 | export
1782 | mapWithKey :  (k -> v -> w)
1783 |            -> Map k v
1784 |            -> Map k w
1785 | mapWithKey _ Tip               =
1786 |   Tip
1787 | mapWithKey f (Bin sx kx x l r) =
1788 |   Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)
1789 |
1790 | ||| The function mapAccumL threads an accumulating
1791 | ||| argument through the map in ascending order of keys. O(n)
1792 | export
1793 | mapAccumL :  (v -> k -> w -> (v, c))
1794 |           -> v
1795 |           -> Map k w
1796 |           -> (v, Map k c)
1797 | mapAccumL _ a Tip               =
1798 |   (a, Tip)
1799 | mapAccumL f a (Bin sx kx x l r) =
1800 |   let (a1, l') = mapAccumL f a l
1801 |       (a2, x') = f a1 kx x
1802 |       (a3, r') = mapAccumL f a2 r
1803 |     in (a3, Bin sx kx x' l' r')
1804 |
1805 | ||| The function mapAccumRWithKey threads an accumulating
1806 | ||| argument through the map in descending order of keys. O(n)
1807 | export
1808 | mapAccumRWithKey :  (v -> k -> w -> (v, c))
1809 |                  -> v
1810 |                  -> Map k w
1811 |                  -> (v, Map k c)
1812 | mapAccumRWithKey _ a Tip               =
1813 |   (a, Tip)
1814 | mapAccumRWithKey f a (Bin sx kx x l r) =
1815 |   let (a1, r') = mapAccumRWithKey f a r
1816 |       (a2, x') = f a1 kx x
1817 |       (a3, l') = mapAccumRWithKey f a2 l
1818 |     in (a3, Bin sx kx x' l' r')
1819 |
1820 | ||| The function mapAccumWithKey threads an accumulating
1821 | ||| argument through the map in ascending order of keys. O(n)
1822 | export
1823 | mapAccumWithKey :  (v -> k -> w -> (v, c))
1824 |                 -> v
1825 |                 -> Map k w
1826 |                 -> (v, Map k c)
1827 | mapAccumWithKey f a t =
1828 |   mapAccumL f a t
1829 |
1830 | ||| The function mapAccum threads an accumulating
1831 | ||| argument through the map in ascending order of keys. O(n)
1832 | export
1833 | mapAccum :  (v -> w -> (v, c))
1834 |          -> v
1835 |          -> Map k w
1836 |          -> (v, Map k c)
1837 | mapAccum f a m =
1838 |   mapAccumWithKey (\a', _, x' => f a' x') a m
1839 |
1840 | --------------------------------------------------------------------------------
1841 | --          Lists
1842 | --------------------------------------------------------------------------------
1843 |
1844 | ||| Convert the map to a list of key/value pairs where the keys are in descending order. O(n)
1845 | export
1846 | toDescList :  Map k v
1847 |            -> List (k, v)
1848 | toDescList Tip               =
1849 |   []
1850 | toDescList t@(Bin _ _ _ _ _) =
1851 |   foldlWithKey (\xs, k, x => (k,x) :: xs) [] t
1852 |
1853 | ||| Convert the map to a list of key/value pairs where the keys are in ascending order. O(n)
1854 | export
1855 | toAscList :  Map k v
1856 |           -> List (k, v)
1857 | toAscList Tip               =
1858 |   []
1859 | toAscList t@(Bin _ _ _ _ _) =
1860 |   foldrWithKey (\k, x, xs => (k,x) :: xs) [] t
1861 |
1862 | ||| Convert the map to a list of key/value pairs.
1863 | export
1864 | toList :  Map k v
1865 |        -> List (k, v)
1866 | toList =
1867 |   toAscList
1868 |
1869 | ||| Build a map from a list of key/value pairs.
1870 | ||| If the list contains more than one value for the same key, the last value
1871 | ||| for the key is retained.
1872 | ||| If the keys of the list are ordered, a linear-time implementation is used. O(n * log(n))
1873 | export
1874 | fromList :  Ord k
1875 |          => Eq v
1876 |          => Eq (Map k v)
1877 |          => List (k, v)
1878 |          -> Map k v
1879 | fromList Nil                =
1880 |   Tip
1881 | fromList ((kx, x) :: [])    =
1882 |   Bin 1 kx x Tip Tip
1883 | fromList ((kx0, x0) :: xs0) =
1884 |   case notOrdered kx0 xs0 of
1885 |     True  =>
1886 |       fromList' (Bin 1 kx0 x0 Tip Tip) xs0
1887 |     False =>
1888 |       go 1 (Bin 1 kx0 x0 Tip Tip) xs0
1889 |   where
1890 |     notOrdered :  Ord a
1891 |                => a
1892 |                -> List (a, b)
1893 |                -> Bool
1894 |     notOrdered _  []             =
1895 |       False
1896 |     notOrdered kx ((ky, _) :: _) =
1897 |       kx >= ky
1898 |     fromList' :  Map k v
1899 |               -> List (k, v)
1900 |               -> Map k v
1901 |     fromList' t0 xs =
1902 |       foldl (\t, (k, x) => insert k x t) t0 xs
1903 |     create :  Nat
1904 |            -> List (k, v)
1905 |            -> (Map k v, List (k, v), List (k, v))
1906 |     create _ []                  =
1907 |       (Tip, [], [])
1908 |     create s xs@((kx, x) :: xss) =
1909 |       case s == 1 of
1910 |         True  =>
1911 |           case notOrdered kx xss of
1912 |             True  =>
1913 |               (Bin 1 kx x Tip Tip, [], xss)
1914 |             False =>
1915 |               (Bin 1 kx x Tip Tip, xss, [])
1916 |         False =>
1917 |           let create' = assert_total $ create (integerToNat ((natToInteger s) `shiftR` 1)) xs
1918 |             in case create' of
1919 |                  res@(_, []                 , _)  =>
1920 |                    res
1921 |                  (l    , [(ky, y)]          , zs) =>
1922 |                    (insertMax ky y l, [], zs)
1923 |                  (l    , ys@((ky, y) :: yss), _)  =>
1924 |                    case notOrdered ky yss of
1925 |                      True  =>
1926 |                        (l, [], ys)
1927 |                      False =>
1928 |                        let (r, zs, ws) = assert_total $ create (integerToNat ((natToInteger s) `shiftR` 1)) yss
1929 |                          in (link ky y l r, zs, ws)
1930 |     go :  Nat
1931 |        -> Map k v
1932 |        -> List (k, v)
1933 |        -> Map k v
1934 |     go _ t []                  =
1935 |       t
1936 |     go _ t ((kx, x) :: [])     =
1937 |       insertMax kx x t
1938 |     go s l xs@((kx, x) :: xss) =
1939 |       case notOrdered kx xss of
1940 |         True  =>
1941 |           fromList' l xs
1942 |         False =>
1943 |           case create s xss of
1944 |             (r, ys, []) =>
1945 |               assert_total $ go (integerToNat ((natToInteger s) `shiftL` 1)) (link kx x l r) ys
1946 |             (r, _,  ys) =>
1947 |               fromList' (link kx x l r) ys
1948 |
1949 | --------------------------------------------------------------------------------
1950 | --          Keys
1951 | --------------------------------------------------------------------------------
1952 |
1953 | ||| Gets the keys of the map.
1954 | export
1955 | keys :  Map k v
1956 |      -> List k
1957 | keys =
1958 |   map fst . toList
1959 |
1960 | --------------------------------------------------------------------------------
1961 | --          Values
1962 | --------------------------------------------------------------------------------
1963 |
1964 | ||| Gets the values of the map.
1965 | ||| Could contain duplicates.
1966 | export
1967 | values :  Map k v
1968 |        -> List v
1969 | values =
1970 |   map snd . toList
1971 |
1972 | --------------------------------------------------------------------------------
1973 | --          Interfaces
1974 | --------------------------------------------------------------------------------
1975 |
1976 | export
1977 | Functor (Map k) where
1978 |   map f m = map f m
1979 |
1980 | export
1981 | Foldable (Map k) where
1982 |   foldl f z = foldl f z . values
1983 |   foldr f z = foldr f z . values
1984 |   null      = null
1985 |
1986 | private
1987 | Show k => Show v => Show (List (k, v)) where
1988 |   show xs = "[" ++ show' xs ++ "]"
1989 |     where
1990 |       show'' : (k, v) -> String
1991 |       show'' (k, v) = "(" ++ show k ++ ", " ++ show v ++ ")"
1992 |       show' : List (k, v) -> String
1993 |       show' Nil        = ""
1994 |       show' (x :: Nil) = show'' x
1995 |       show' (x :: xs)  = show'' x ++ ", " ++ show' xs
1996 |
1997 | export
1998 | Show (List (k, v)) => Show (Map k v) where
1999 |   show m = "fromList " ++ (show $ toList m)
2000 |