0 | ||| Finite Sets
   1 | module Data.Set
   2 |
   3 | import public Data.Set.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 Sets
  15 | --------------------------------------------------------------------------------
  16 |
  17 | ||| The empty set. O(1)
  18 | export
  19 | empty : Set a
  20 | empty = Tip
  21 |
  22 | --------------------------------------------------------------------------------
  23 | --          Folds
  24 | --------------------------------------------------------------------------------
  25 |
  26 | ||| Fold the values in the set using the given left-associative binary operator. O(n)
  27 | export
  28 | foldl :  (a -> b -> a)
  29 |       -> a
  30 |       -> Set b
  31 |       -> a
  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 set using the given right-associative binary operator. O(n)
  38 | export
  39 | foldr :  (a -> b -> b)
  40 |       -> b
  41 |       -> Set a
  42 |       -> b
  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 | --------------------------------------------------------------------------------
  49 | --          Insertion
  50 | --------------------------------------------------------------------------------
  51 |
  52 | ||| Insert an element in a set.
  53 | ||| If the set already contains an element equal to the given value,
  54 | ||| it is replaced with the new value. O(log n)
  55 | export
  56 | insert :  Ord a
  57 |        => a
  58 |        -> Set a
  59 |        -> Set a
  60 | insert x0 s =
  61 |   go x0 x0 s
  62 |   where
  63 |     go :  a
  64 |        -> a
  65 |        -> Set a
  66 |        -> Set a
  67 |     go orig _ Tip              =
  68 |       singleton orig
  69 |     go orig x t@(Bin sz y l r) =
  70 |       case compare x y of
  71 |         LT =>
  72 |           let l' = go orig x l
  73 |             in case l' == l of
  74 |                  True  =>
  75 |                    t
  76 |                  False =>
  77 |                    balanceL y l' r
  78 |         GT =>
  79 |           let r' = go orig x r
  80 |             in case r' == r of
  81 |                  True  =>
  82 |                    t
  83 |                  False =>
  84 |                    balanceR y l r'
  85 |         EQ =>
  86 |           case orig == y of
  87 |             True  =>
  88 |               t
  89 |             False =>
  90 |               Bin sz orig l r
  91 |
  92 | private
  93 | insertR :  Ord a
  94 |         => a
  95 |         -> Set a
  96 |         -> Set a
  97 | insertR x0 s =
  98 |   go x0 x0 s
  99 |   where
 100 |     go :  a
 101 |        -> a
 102 |        -> Set a
 103 |        -> Set a
 104 |     go orig _ Tip             =
 105 |       singleton orig
 106 |     go orig x t@(Bin _ y l r) =
 107 |       case compare x y of
 108 |         LT =>
 109 |           let l' = go orig x l
 110 |             in case l' == l of
 111 |                  True  =>
 112 |                    t
 113 |                  False =>
 114 |                    balanceL y l' r
 115 |         GT =>
 116 |           let r' = go orig x r
 117 |             in case r' == r of
 118 |                  True  =>
 119 |                    t
 120 |                  False =>
 121 |                    balanceR y l r'
 122 |         EQ =>
 123 |           t
 124 |
 125 | --------------------------------------------------------------------------------
 126 | --          Deletion
 127 | --------------------------------------------------------------------------------
 128 |
 129 | ||| Delete an element from a set. When the element is not
 130 | ||| a member of the set, the original set is returned. O(log n)
 131 | export
 132 | delete :  Ord a
 133 |        => a
 134 |        -> Set a
 135 |        -> Set a
 136 | delete =
 137 |   go
 138 |   where
 139 |     go :  a
 140 |        -> Set a
 141 |        -> Set a
 142 |     go _ Tip             =
 143 |       Tip
 144 |     go x t@(Bin _ y l r) =
 145 |       case compare x y of
 146 |         LT =>
 147 |           let l' = go x l
 148 |             in case l' == l of
 149 |                  True  =>
 150 |                    t
 151 |                  False =>
 152 |                    balanceR y l' r
 153 |         GT =>
 154 |           let r' = go x r
 155 |             in case r' == r of
 156 |                  True  =>
 157 |                    t
 158 |                  False =>
 159 |                    balanceL y l r'
 160 |         EQ =>
 161 |           glue l r
 162 |
 163 | --------------------------------------------------------------------------------
 164 | --          Query
 165 | --------------------------------------------------------------------------------
 166 |
 167 | ||| Is the element in the set? O(log n)
 168 | export
 169 | member :  Ord a
 170 |        => a
 171 |        -> Set a
 172 |        -> Bool
 173 | member _ Tip           =
 174 |   False
 175 | member x (Bin _ y l r) =
 176 |   case compare x y of
 177 |     LT =>
 178 |       member x l
 179 |     GT =>
 180 |       member x r
 181 |     EQ =>
 182 |       True
 183 |
 184 | ||| Is the element not in the set? O(log n)
 185 | export
 186 | notMember :  Ord a
 187 |           => a
 188 |           -> Set a
 189 |           -> Bool
 190 | notMember k m =
 191 |   not $ member k m
 192 |
 193 | ||| Find largest element smaller than the given one. O(log n)
 194 | export
 195 | lookupLT :  Ord a
 196 |          => a
 197 |          -> Set a
 198 |          -> Maybe a
 199 | lookupLT =
 200 |   goNothing
 201 |   where
 202 |     goJust :  a
 203 |            -> a
 204 |            -> Set a
 205 |            -> Maybe a
 206 |     goJust _ best Tip           =
 207 |       Just best
 208 |     goJust x best (Bin _ y l r) =
 209 |       case x <= y of
 210 |         True  =>
 211 |           goJust x best l
 212 |         False =>
 213 |           goJust x y r
 214 |     goNothing :  a
 215 |               -> Set a
 216 |               -> Maybe a
 217 |     goNothing _ Tip           =
 218 |       Nothing
 219 |     goNothing x (Bin _ y l r) =
 220 |       case x <= y of
 221 |         True  =>
 222 |           goNothing x l
 223 |         False =>
 224 |           goJust x y r
 225 |
 226 | ||| Find smallest element greater than the given one. O(log n)
 227 | export
 228 | lookupGT :  Ord a
 229 |          => a
 230 |          -> Set a
 231 |          -> Maybe a
 232 | lookupGT =
 233 |   goNothing
 234 |   where
 235 |     goJust :  a
 236 |            -> a
 237 |            -> Set a
 238 |            -> Maybe a
 239 |     goJust _ best Tip           =
 240 |       Just best
 241 |     goJust x best (Bin _ y l r) =
 242 |       case x < y of
 243 |         True  =>
 244 |           goJust x y l
 245 |         False =>
 246 |           goJust x best r
 247 |     goNothing :  a
 248 |               -> Set a
 249 |               -> Maybe a
 250 |     goNothing _ Tip           =
 251 |       Nothing
 252 |     goNothing x (Bin _ y l r) =
 253 |       case x < y of
 254 |         True  =>
 255 |           goJust x y l
 256 |         False =>
 257 |           goNothing x r
 258 |
 259 | ||| Find the largest element smaller or equal to the given one. O(log n)
 260 | export
 261 | lookupLE :  Ord a
 262 |          => a
 263 |          -> Set a
 264 |          -> Maybe a
 265 | lookupLE =
 266 |   goNothing
 267 |   where
 268 |     goJust :  a
 269 |            -> a
 270 |            -> Set a
 271 |            -> Maybe a
 272 |     goJust _ best Tip           =
 273 |       Just best
 274 |     goJust x best (Bin _ y l r) =
 275 |       case compare x y of
 276 |         LT =>
 277 |           goJust x best l
 278 |         EQ =>
 279 |           Just y
 280 |         GT =>
 281 |           goJust x y r
 282 |     goNothing :  a
 283 |               -> Set a
 284 |               -> Maybe a
 285 |     goNothing _ Tip           =
 286 |       Nothing
 287 |     goNothing x (Bin _ y l r) =
 288 |       case compare x y of
 289 |         LT =>
 290 |           goNothing x l
 291 |         EQ =>
 292 |           Just y
 293 |         GT =>
 294 |           goJust x y r
 295 |
 296 | ||| Find the smallest element greater or equal to the given one. O(log n)
 297 | export
 298 | lookupGE :  Ord a
 299 |          => a
 300 |          -> Set a
 301 |          -> Maybe a
 302 | lookupGE =
 303 |   goNothing
 304 |   where
 305 |     goJust :  a
 306 |            -> a
 307 |            -> Set a
 308 |            -> Maybe a
 309 |     goJust _ best Tip           =
 310 |       Just best
 311 |     goJust x best (Bin _ y l r) =
 312 |       case compare x y of
 313 |         LT =>
 314 |           goJust x y l
 315 |         EQ =>
 316 |           Just y
 317 |         GT =>
 318 |           goJust x best r
 319 |     goNothing :  a
 320 |               -> Set a
 321 |               -> Maybe a
 322 |     goNothing _ Tip           =
 323 |       Nothing
 324 |     goNothing x (Bin _ y l r) =
 325 |       case compare x y of
 326 |         LT =>
 327 |           goJust x y l
 328 |         EQ =>
 329 |           Just y
 330 |         GT =>
 331 |           goNothing x r
 332 |
 333 | ||| Is the set empty? O(1)
 334 | export
 335 | null :  Set a
 336 |      -> Bool
 337 | null Tip =
 338 |   True
 339 | null _   =
 340 |   False
 341 |
 342 | ||| Performs a split but also returns whether
 343 | ||| the pivot element was found in the original set. O(log n)
 344 | export
 345 | splitMember :  Ord a
 346 |             => a
 347 |             -> Set a
 348 |             -> (Set a, Bool, Set a)
 349 | splitMember _ Tip           =
 350 |   (Tip, False, Tip)
 351 | splitMember x (Bin _ y l r) =
 352 |   case compare x y of
 353 |     LT =>
 354 |       let (lt, found, gt) = splitMember x l
 355 |           gt'             = link y gt r
 356 |         in (lt, found, gt')
 357 |     GT =>
 358 |       let (lt, found, gt) = splitMember x r
 359 |           lt'             = link y l lt
 360 |         in (lt', found, gt)
 361 |     EQ =>
 362 |       (l, True, r)
 363 |
 364 | ||| Decompose a set into pieces based on the structure of the underlying tree.
 365 | ||| This function is useful for consuming a set in parallel. O(1)
 366 | export
 367 | splitRoot :  Set a
 368 |           -> List (Set a)
 369 | splitRoot orig =
 370 |   case orig of
 371 |     Tip         =>
 372 |       []
 373 |     Bin _ v l r =>
 374 |       [l, singleton v, r]
 375 |
 376 | ||| The expression (split x set) is a pair (set1,set2) where
 377 | ||| set1 comprises the elements of set less than x and
 378 | ||| set2 comprises the elements of set greater than x. O(log n)
 379 | export
 380 | split :  Ord a
 381 |       => a
 382 |       -> Set a
 383 |       -> (Set a, Set a)
 384 | split _ Tip           =
 385 |   (Tip, Tip)
 386 | split x (Bin _ y l r) =
 387 |   case compare x y of
 388 |     LT =>
 389 |       let (lt, gt) = split x l
 390 |         in (lt, link y gt r)
 391 |     GT =>
 392 |       let (lt, gt) = split x r
 393 |         in (link y l lt, gt)
 394 |     EQ =>
 395 |       (l, r)
 396 |
 397 | private
 398 | isSubsetOfX :  Ord a
 399 |             => Set a
 400 |             -> Set a
 401 |             -> Bool
 402 | isSubsetOfX Tip           _   =
 403 |   True
 404 | isSubsetOfX _             Tip =
 405 |   False
 406 | isSubsetOfX (Bin 1 x _ _) t   =
 407 |   member x t
 408 | isSubsetOfX (Bin _ x l r) t   =
 409 |   let (lt, found, gt) = splitMember x t
 410 |     in found             &&
 411 |        size l <= size lt &&
 412 |        size r <= size gt &&
 413 |        isSubsetOfX l lt  &&
 414 |        isSubsetOfX r gt
 415 |
 416 | ||| Indicates whether s1 is a subset of s2.
 417 | export
 418 | isSubsetOf :  Ord a
 419 |            => Set a
 420 |            -> Set a
 421 |            -> Bool
 422 | isSubsetOf t1 t2 =
 423 |   size t1 <= size t2 && isSubsetOfX t1 t2
 424 |
 425 | ||| Indicates whether s1 is a proper subset of s2.
 426 | export
 427 | isProperSubsetOf :  Ord a
 428 |                  => Set a
 429 |                  -> Set a
 430 |                  -> Bool
 431 | isProperSubsetOf s1 s2 =
 432 |   size s1 <= size s2 && isSubsetOfX s1 s2
 433 |
 434 | ||| Check whether two sets are disjoint (i.e. their intersection is empty).
 435 | export
 436 | disjoint :  Ord a
 437 |          => Set a
 438 |          -> Set a
 439 |          -> Bool
 440 | disjoint Tip           _   =
 441 |   True
 442 | disjoint _             Tip =
 443 |   True
 444 | disjoint (Bin 1 x _ _) t   =
 445 |   notMember x t
 446 | disjoint (Bin _ x l r) t   =
 447 |   let (lt, found, gt) = splitMember x t
 448 |     in not found     &&
 449 |        disjoint l lt &&
 450 |        disjoint r gt
 451 |
 452 | --------------------------------------------------------------------------------
 453 | --          Merge
 454 | --------------------------------------------------------------------------------
 455 |
 456 | ||| Merges two trees.
 457 | private
 458 | merge :  Set a
 459 |       -> Set a
 460 |       -> Set a
 461 | merge Tip                   r                     =
 462 |   r
 463 | merge l                     Tip                   =
 464 |   l
 465 | merge l@(Bin sizeL x lx rx) r@(Bin sizeR y ly ry) =
 466 |   case delta * sizeL < sizeR of
 467 |     True  =>
 468 |       balanceL y (merge l ly) ry
 469 |     False =>
 470 |       case delta*sizeR < sizeL of
 471 |         True  =>
 472 |           balanceR x lx (merge rx r)
 473 |         False =>
 474 |           glue l r
 475 |
 476 | --------------------------------------------------------------------------------
 477 | --          Filter
 478 | --------------------------------------------------------------------------------
 479 |
 480 | ||| Filter all elements that satisfy the predicate. O(n)
 481 | export
 482 | filter :  Eq (Set a)
 483 |        => (a -> Bool)
 484 |        -> Set a
 485 |        -> Set a
 486 | filter _ Tip             =
 487 |   Tip
 488 | filter p t@(Bin _ x l r) =
 489 |   case p x of
 490 |     True  =>
 491 |       case l == (filter p l) && r == (filter p r) of
 492 |         True  =>
 493 |           t
 494 |         False =>
 495 |           link x (filter p l) (filter p r)
 496 |     False =>
 497 |       merge (filter p l) (filter p r)
 498 |
 499 | ||| Partition the set into two sets, one with all elements
 500 | ||| that satisfy the predicate and one with all elements
 501 | ||| that don't satisfy the predicate.
 502 | ||| See also split.
 503 | partition :  Eq (Set a)
 504 |           => (a -> Bool)
 505 |           -> Set a
 506 |           -> (Set a, Set a)
 507 | partition _ Tip             =
 508 |   (Tip, Tip)
 509 | partition p t@(Bin _ x l r) =
 510 |   case (partition p l, partition p r) of
 511 |     ((l1, l2), (r1, r2)) =>
 512 |       case p x of
 513 |         True  =>
 514 |           case l1 == l && r1 == r of
 515 |             True  =>
 516 |               (t, merge l2 r2)
 517 |             False =>
 518 |               (link x l1 r1, merge l2 r2)
 519 |         False =>
 520 |           case l2 == l && r2 == r of
 521 |             True  =>
 522 |               (merge l1 r1, t)
 523 |             False =>
 524 |               (merge l1 r1, link x l2 r2)
 525 |
 526 | ||| Take while a predicate on the keys holds.
 527 | ||| The user is responsible for ensuring that for all elements j and k in the set,
 528 | ||| j < k ==> p j >= p k. See note at spanAntitone. O(log n)
 529 | export
 530 | takeWhileAntitone :  (a -> Bool)
 531 |                   -> Set a
 532 |                   -> Set a
 533 | takeWhileAntitone _ Tip           =
 534 |   Tip
 535 | takeWhileAntitone p (Bin _ x l r) =
 536 |   case p x of
 537 |     True  =>
 538 |       link x l (takeWhileAntitone p r)
 539 |     False =>
 540 |       takeWhileAntitone p l
 541 |
 542 | ||| Drop while a predicate on the keys holds.
 543 | ||| The user is responsible for ensuring that for all elements j and k in the map,
 544 | ||| j < k ==> p j >= p k. See note at spanAntitone. O(log n)
 545 | export
 546 | dropWhileAntitone :  (a -> Bool)
 547 |                   -> Set a
 548 |                   -> Set a
 549 | dropWhileAntitone _ Tip           =
 550 |   Tip
 551 | dropWhileAntitone p (Bin _ x l r) =
 552 |   case p x of
 553 |     True  =>
 554 |       dropWhileAntitone p r
 555 |     False =>
 556 |       link x (dropWhileAntitone p l) r
 557 |
 558 | ||| Divide a map at the point where a predicate on the keys stops holding.
 559 | ||| The user is responsible for ensuring that for all keys j and k in the map,
 560 | ||| j < k ==> p j>= p k. O(log n)
 561 | export
 562 | spanAntitone :  (a -> Bool)
 563 |              -> Set a
 564 |              -> (Set a, Set a)
 565 | spanAntitone p0 s =
 566 |   go p0 s
 567 |   where
 568 |     go :  (a -> Bool)
 569 |        -> Set a
 570 |        -> (Set a, Set a)
 571 |     go _ Tip           =
 572 |       (Tip, Tip)
 573 |     go p (Bin _ x l r) =
 574 |       case p x of
 575 |         True  =>
 576 |           let (u, v) = go p r
 577 |             in (link x l u, v)
 578 |         False =>
 579 |           let (u, v) = go p l
 580 |             in (u, link x v r)
 581 |
 582 | --------------------------------------------------------------------------------
 583 | --          Indexed
 584 | --------------------------------------------------------------------------------
 585 |
 586 | ||| Lookup the index of a element, which is its zero-based index in
 587 | ||| the sorted sequence of elements. The index is a number from 0 up to, but not
 588 | ||| including, the size of the set. O(log n)
 589 | export
 590 | lookupIndex :  Ord a
 591 |             => a
 592 |             -> Set a
 593 |             -> Maybe Nat
 594 | lookupIndex =
 595 |   go 0
 596 |   where
 597 |     go :  Nat
 598 |        -> a
 599 |        -> Set a
 600 |        -> Maybe Nat
 601 |     go _  _ Tip             =
 602 |       Nothing
 603 |     go idx x (Bin _ kx l r) =
 604 |       case compare x kx of
 605 |         LT =>
 606 |           go idx x l
 607 |         GT =>
 608 |           go (idx + size l + 1) x r
 609 |         EQ =>
 610 |           Just $ idx + size l
 611 |
 612 | ||| Return the index of an element, which is its zero-based index in
 613 | ||| the sorted sequence of elements. The index is a number from 0 up to, but not
 614 | ||| including, the size of the set. Calls idris_crash when the element is not
 615 | ||| a member of the set. O(log n)
 616 | export
 617 | findIndex :  Ord a
 618 |           => a
 619 |           -> Set a
 620 |           -> Nat
 621 | findIndex =
 622 |   go 0
 623 |   where
 624 |     go :  Nat
 625 |        -> a
 626 |        -> Set a
 627 |        -> Nat
 628 |     go _   _ Tip            =
 629 |       assert_total $ idris_crash "Data.Set.findIndex: element is not in the set"
 630 |     go idx x (Bin _ kx l r) =
 631 |       case compare x kx of
 632 |         LT =>
 633 |           go idx x l
 634 |         GT =>
 635 |           go (idx + size l + 1) x r
 636 |         EQ =>
 637 |           idx + size l
 638 |
 639 | ||| Retrieve an element by its index, i.e. by its zero-based
 640 | ||| index in the sorted sequence of elements. If the index is out of range (less
 641 | ||| than zero, greater or equal to size of the set), idris_crash is called. O(log n)
 642 | export
 643 | elemAt :  Nat
 644 |        -> Set a
 645 |        -> a
 646 | elemAt _ Tip           =
 647 |   assert_total $ idris_crash "Data.Set.elemAt: index out of range"
 648 | elemAt i (Bin _ x l r) =
 649 |   case compare i (size l) of
 650 |      LT =>
 651 |        elemAt i l
 652 |      GT =>
 653 |        elemAt ((i `minus` (size l)) `minus` 1) r
 654 |      EQ =>
 655 |        x
 656 |
 657 | ||| Delete the element at index, i.e. by its zero-based index in
 658 | ||| the sorted sequence of elements. If the index is out of range (less than zero,
 659 | ||| greater or equal to size of the set), idris_crash is called. O(log n)
 660 | export
 661 | deleteAt :  Nat
 662 |          -> Set a
 663 |          -> Set a
 664 | deleteAt i t =
 665 |   case t of
 666 |     Tip         =>
 667 |       assert_total $ idris_crash "Data.Set.deleteAt: index out of range"
 668 |     Bin _ x l r =>
 669 |       case compare i (size l) of
 670 |         LT =>
 671 |           balanceR x (deleteAt i l) r
 672 |         GT =>
 673 |           balanceL x l (deleteAt ((i `minus` (size l)) `minus` 1) r)
 674 |         EQ =>
 675 |           glue l r
 676 |
 677 | ||| Take a given number of elements in order, beginning
 678 | ||| with the smallest keys. O(log n)
 679 | export
 680 | take :  Nat
 681 |      -> Set a
 682 |      -> Set a
 683 | take i s =
 684 |   case i >= size s of
 685 |     True  =>
 686 |       s
 687 |     False =>
 688 |       go i s
 689 |   where
 690 |     go :  Nat
 691 |        -> Set a
 692 |        -> Set a
 693 |     go _ Tip           =
 694 |       Tip
 695 |     go i (Bin _ x l r) =
 696 |       case i <= 0 of
 697 |         True  =>
 698 |           Tip
 699 |         False =>
 700 |           case compare i (size l) of
 701 |             LT =>
 702 |               go i l
 703 |             GT =>
 704 |               link x l (go ((i `minus` (size l)) `minus` 1) r)
 705 |             EQ =>
 706 |               l
 707 |
 708 | ||| Drop a given number of elements in order, beginning
 709 | ||| with the smallest ones. O(log n)
 710 | export
 711 | drop :  Nat
 712 |      -> Set a
 713 |      -> Set a
 714 | drop i s =
 715 |   case i >= size s of
 716 |     True  =>
 717 |       Tip
 718 |     False =>
 719 |       go i s
 720 |   where
 721 |     go :  Nat
 722 |        -> Set a
 723 |        -> Set a
 724 |     go _ Tip             =
 725 |       Tip
 726 |     go i s@(Bin _ x l r) =
 727 |       case i <= 0 of
 728 |         True  =>
 729 |           s
 730 |         False =>
 731 |           case compare i (size l) of
 732 |             LT =>
 733 |               link x (go i l) r
 734 |             GT =>
 735 |               go ((i `minus` (size l)) `minus` 1) r
 736 |             EQ =>
 737 |               insertMin x r
 738 |
 739 | ||| Split a set at a particular index. O(log n)
 740 | export
 741 | splitAt :  Nat
 742 |         -> Set a
 743 |         -> (Set a, Set a)
 744 | splitAt i s =
 745 |   case i >= size s of
 746 |     True  =>
 747 |       (s, Tip)
 748 |     False =>
 749 |       go i s
 750 |   where
 751 |     go :  Nat
 752 |        -> Set a
 753 |        -> (Set a, Set a)
 754 |     go _ Tip             =
 755 |       (Tip, Tip)
 756 |     go i s@(Bin _ x l r) =
 757 |       case i <= 0 of
 758 |         True  =>
 759 |           (Tip, s)
 760 |         False =>
 761 |           case compare i (size l) of
 762 |             LT =>
 763 |               case go i l of
 764 |                 (ll, lr) =>
 765 |                   (ll, link x lr r)
 766 |             GT =>
 767 |               case go ((i `minus` (size l)) `minus` 1) r of
 768 |                 (rl, rr) =>
 769 |                   (link x l rl, rr)
 770 |             EQ =>
 771 |               (l, insertMin x r)
 772 |
 773 | --------------------------------------------------------------------------------
 774 | --          Min/Max
 775 | --------------------------------------------------------------------------------
 776 |
 777 | private
 778 | lookupMinSure :  a
 779 |               -> Set a
 780 |               -> a
 781 | lookupMinSure x Tip           =
 782 |   x
 783 | lookupMinSure _ (Bin _ x l _) =
 784 |   lookupMinSure x l
 785 |
 786 | ||| The minimal element of the set. Returns Nothing if the set is empty. O(log n)
 787 | export
 788 | lookupMin :  Set a
 789 |           -> Maybe a
 790 | lookupMin Tip           =
 791 |   Nothing
 792 | lookupMin (Bin _ x l _) =
 793 |   Just $ lookupMinSure x l
 794 |
 795 | private
 796 | lookupMaxSure :  a
 797 |               -> Set a
 798 |               -> a
 799 | lookupMaxSure x Tip           =
 800 |   x
 801 | lookupMaxSure _ (Bin _ x _ r) =
 802 |   lookupMaxSure x r
 803 |
 804 | ||| The maximal element of the set. Returns Nothing if the set is empty. O(log n)
 805 | export
 806 | lookupMax :  Set a
 807 |           -> Maybe a
 808 | lookupMax Tip           =
 809 |   Nothing
 810 | lookupMax (Bin _ x _ r) =
 811 |   Just $ lookupMaxSure x r
 812 |
 813 | ||| The minimal element of the set. Calls idris_crash if the set is empty. O(log n)
 814 | export
 815 | findMin :  Set a
 816 |         -> a
 817 | findMin t =
 818 |   case lookupMin t of
 819 |     Just r  =>
 820 |       r
 821 |     Nothing =>
 822 |       assert_total $ idris_crash "Data.Set.findMin: empty set has no minimal element"
 823 |
 824 | ||| The maximal element of the set. Calls idris_crash if the set is empty. O(log n)
 825 | export
 826 | findMax :  Set a
 827 |         -> a
 828 | findMax t =
 829 |   case lookupMax t of
 830 |     Just r  =>
 831 |       r
 832 |     Nothing =>
 833 |       assert_total $ idris_crash "Data.Set.findMax: empty set has no maximal element"
 834 |
 835 | ||| Delete the minimal element. Returns an empty set if the set is empty. O(log n)
 836 | export
 837 | deleteMin :  Set a
 838 |           -> Set a
 839 | deleteMin Tip              =
 840 |   Tip
 841 | deleteMin (Bin _ _ Tip r)  =
 842 |   r
 843 | deleteMin (Bin _ x l   r)  =
 844 |   balanceR x (deleteMin l) r
 845 |
 846 | ||| Delete the maximal element. Returns an empty set if the set is empty. O(log n)
 847 | export
 848 | deleteMax :  Set a
 849 |           -> Set a
 850 | deleteMax Tip             =
 851 |   Tip
 852 | deleteMax (Bin _ _ l Tip) =
 853 |   l
 854 | deleteMax (Bin _ x l r)   =
 855 |   balanceL x l (deleteMax r)
 856 |
 857 | ||| Retrieves the minimal element of the set, and
 858 | ||| the set stripped of that element, or Nothing if passed an empty set. O(log n)
 859 | export
 860 | minView :  Set a
 861 |         -> Maybe (a, Set a)
 862 | minView Tip           =
 863 |   Nothing
 864 | minView (Bin _ x l r) =
 865 |   Just $ minViewSure x l r
 866 |
 867 | ||| Delete and find the minimal element. O(log n)
 868 | export
 869 | deleteFindMin :  Set a
 870 |               -> (a, Set a)
 871 | deleteFindMin t =
 872 |   case minView t of
 873 |     Just res =>
 874 |       res
 875 |     Nothing  =>
 876 |       (assert_total $ idris_crash "Data.Set.deleteFindMin: can not return the minimal element of an empty set", Tip)
 877 |
 878 | ||| Retrieves the maximal element of the set, and
 879 | ||| the set stripped of that element, or Nothing if passed an empty set. O(log n)
 880 | export
 881 | maxView :  Set a
 882 |         -> Maybe (a, Set a)
 883 | maxView Tip           =
 884 |   Nothing
 885 | maxView (Bin _ x l r) =
 886 |   Just $ maxViewSure x l r
 887 |
 888 | ||| Delete and find the maximal element. O(log n)
 889 | export
 890 | deleteFindMax :  Set a
 891 |               -> (a, Set a)
 892 | deleteFindMax t =
 893 |   case maxView t of
 894 |     Just res =>
 895 |       res
 896 |     Nothing  =>
 897 |       (assert_total $ idris_crash "Data.Set.deleteFindMax: can not return the maximal element of an empty set", Tip)
 898 |
 899 | --------------------------------------------------------------------------------
 900 | --          Combine
 901 | --------------------------------------------------------------------------------
 902 |
 903 | ||| The expression (union t1 t2) takes the left-biased union of t1 and t2.
 904 | ||| It prefers t1 when duplicate keys are encountered.
 905 | export
 906 | union :  Eq (Set a)
 907 |       => Eq a
 908 |       => Ord a
 909 |       => Set a
 910 |       -> Set a
 911 |       -> Set a
 912 | union t1                 Tip           =
 913 |   t1
 914 | union t1                 (Bin 1 x _ _) =
 915 |   insertR x t1
 916 | union (Bin 1 x _ _)      t2            =
 917 |   insert x t2
 918 | union Tip                t2            =
 919 |   t2
 920 | union t1@(Bin _ x l1 r1) t2            =
 921 |   let (l2,r2) = split x t2
 922 |       l1l2    = union l1 l2
 923 |       r1r2    = union r1 r2
 924 |     in case l1l2 == l1 && r1r2 == r1 of
 925 |          True  =>
 926 |            t1
 927 |          False =>
 928 |            link x l1l2 r1r2
 929 |
 930 | --------------------------------------------------------------------------------
 931 | --          Difference
 932 | --------------------------------------------------------------------------------
 933 |
 934 | ||| Difference of two sets.
 935 | ||| Return elements of the first set not existing in the second set.
 936 | export
 937 | difference :  Ord a
 938 |            => Set a
 939 |            -> Set a
 940 |            -> Set a
 941 | difference Tip _               =
 942 |   Tip
 943 | difference t1  Tip             =
 944 |   t1
 945 | difference t1  (Bin _ x l2 r2) =
 946 |   let (l1, r1) = split x t1
 947 |       l1l2     = difference l1 l2
 948 |       r1r2     = difference r1 r2
 949 |     in case size l1l2 + size r1r2 == size t1 of
 950 |          True  =>
 951 |            t1
 952 |          False =>
 953 |            merge l1l2 r1r2
 954 |
 955 | ||| Same as difference.
 956 | export
 957 | (\\) :  Ord a
 958 |      => Set a
 959 |      -> Set a
 960 |      -> Set a
 961 | s1 \\ s2 =
 962 |   difference s1 s2
 963 |
 964 | --------------------------------------------------------------------------------
 965 | --          Intersection
 966 | --------------------------------------------------------------------------------
 967 |
 968 | ||| Intersection of two sets.
 969 | ||| Return data in the first set for elements existing in both sets.
 970 | export
 971 | intersection :  Eq (Set a)
 972 |              => Ord a
 973 |              => Set a
 974 |              -> Set a
 975 |              -> Set a
 976 | intersection Tip                _   =
 977 |   Tip
 978 | intersection _                  Tip =
 979 |   Tip
 980 | intersection t1@(Bin _ x l1 r1) t2  =
 981 |   let (l2,x',r2) = splitMember x t2
 982 |       l1l2       = intersection l1 l2
 983 |       r1r2       = intersection r1 r2
 984 |     in case x' of
 985 |          True  =>
 986 |            case l1l2 == l1 && r1r2 == r1 of
 987 |              True  =>
 988 |                t1
 989 |              False =>
 990 |                link x l1l2 r1r2
 991 |          False =>
 992 |            merge l1l2 r1r2
 993 |
 994 | --------------------------------------------------------------------------------
 995 | --          Lists
 996 | --------------------------------------------------------------------------------
 997 |
 998 | ||| Convert the set to a list of elements where the elements are in descending order. O(n)
 999 | export
1000 | toDescList :  Set a
1001 |            -> List a
1002 | toDescList Tip             =
1003 |   []
1004 | toDescList t@(Bin _ _ _ _) =
1005 |   foldl (\xs, x => x :: xs) [] t
1006 |
1007 | ||| Convert the set to a list of elements where the elements are in ascending order. O(n)
1008 | export
1009 | toAscList :  Set a
1010 |           -> List a
1011 | toAscList Tip             =
1012 |   []
1013 | toAscList t@(Bin _ _ _ _) =
1014 |   foldr (\x, xs => x :: xs) [] t
1015 |
1016 | ||| Convert the set to a list of elements.
1017 | export
1018 | toList :  Set a
1019 |        -> List a
1020 | toList =
1021 |   toAscList
1022 |
1023 | ||| Build a set from a list of elements.
1024 | ||| If the list contains identical element(s),
1025 | ||| the last of each identical elemen is retained.
1026 | ||| If the elements of the list are ordered, a linear-time implementation is used. O(n * log(n))
1027 | export
1028 | fromList :  Ord a
1029 |          => List a
1030 |          -> Set a
1031 | fromList []          =
1032 |   Tip
1033 | fromList (x :: [])   =
1034 |   Bin 1 x Tip Tip
1035 | fromList (x0 :: xs0) =
1036 |   case notOrdered x0 xs0 of
1037 |     True  =>
1038 |       fromList' (Bin 1 x0 Tip Tip) xs0
1039 |     False =>
1040 |       go 1 (Bin 1 x0 Tip Tip) xs0
1041 |   where
1042 |     notOrdered :  a
1043 |                -> List a
1044 |                -> Bool
1045 |     notOrdered _  []      =
1046 |       False
1047 |     notOrdered x (y :: _) =
1048 |       x >= y
1049 |     fromList' :  Set a
1050 |               -> List a
1051 |               -> Set a
1052 |     fromList' t0 xs =
1053 |       foldl (\t, x => insert x t) t0 xs
1054 |     create :  Nat
1055 |            -> List a
1056 |            -> (Set a, List a, List a)
1057 |     create _ []            =
1058 |       (Tip, [], [])
1059 |     create s xs@(x :: xss) =
1060 |       case s == 1 of
1061 |         True  =>
1062 |           case notOrdered x xss of
1063 |             True  =>
1064 |               (Bin 1 x Tip Tip, [], xss)
1065 |             False =>
1066 |               (Bin 1 x Tip Tip, xss, [])
1067 |         False =>
1068 |           let create' = assert_total $ create (integerToNat ((natToInteger s) `shiftR` 1)) xs
1069 |             in case create' of
1070 |                  res@(_, []           , _)  =>
1071 |                    res
1072 |                  (l    , [y]          , zs) =>
1073 |                    (insertMax y l, [], zs)
1074 |                  (l    , ys@(y :: yss), _)  =>
1075 |                    case notOrdered y yss of
1076 |                      True  =>
1077 |                        (l, [], ys)
1078 |                      False =>
1079 |                        let (r, zs, ws) = assert_total $ create (integerToNat ((natToInteger s) `shiftR` 1)) yss
1080 |                          in (link y l r, zs, ws)
1081 |     go :  Nat
1082 |        -> Set a
1083 |        -> List a
1084 |        -> Set a
1085 |     go _ t []            =
1086 |       t
1087 |     go _ t (x :: [])     =
1088 |       insertMax x t
1089 |     go s l xs@(x :: xss) =
1090 |       case notOrdered x xss of
1091 |         True  =>
1092 |           fromList' l xs
1093 |         False =>
1094 |           let create' = assert_total $ create s xss
1095 |             in case create' of
1096 |                  (r, ys, []) =>
1097 |                    assert_total $ go (integerToNat ((natToInteger s) `shiftL` 1)) (link x l r) ys
1098 |                  (r, _,  ys) =>
1099 |                    fromList' (link x l r) ys
1100 |
1101 | --------------------------------------------------------------------------------
1102 | --          Map
1103 | --------------------------------------------------------------------------------
1104 |
1105 | ||| Map a function over all elements in the set. O(n)
1106 | export
1107 | map :  (a -> b)
1108 |     -> Set a
1109 |     -> Set b
1110 | map _ Tip            =
1111 |   Tip
1112 | map f (Bin sz x l r) =
1113 |   Bin sz (f x) (map f l) (map f r)
1114 |
1115 | --------------------------------------------------------------------------------
1116 | --          Disjoint Union
1117 | --------------------------------------------------------------------------------
1118 |
1119 | ||| Calculate the disjoint union of two sets. O(n + m)
1120 | export
1121 | disjointUnion :  Set a
1122 |               -> Set b
1123 |               -> Set (Either a b)
1124 | disjointUnion as bs =
1125 |   merge (map Left as) (map Right bs)
1126 |
1127 | --------------------------------------------------------------------------------
1128 | --          Interfaces
1129 | --------------------------------------------------------------------------------
1130 |
1131 | export
1132 | Functor Set where
1133 |   map f s = map f s
1134 |
1135 | export
1136 | Show (List a) => Show (Set a) where
1137 |   show s = "fromList " ++ (show $ toList s)
1138 |