
   0 | ||| Ordered Nat Priority Search Queue
   1 | module Data.NatPSQ
   2 | 
   3 | import public Data.NatPSQ.Internal
   4 | 
   5 | import Data.List
   6 | import Data.Maybe
   7 | 
   8 | %default total
   9 | 
  10 | --------------------------------------------------------------------------------
  11 | --          Unsafe operations
  12 | --------------------------------------------------------------------------------
  13 | 
  14 | private
  15 | link :  Key
  16 |      -> p
  17 |      -> v
  18 |      -> Key
  19 |      -> NatPSQ p v
  20 |      -> NatPSQ p v
  21 |      -> NatPSQ p v
  22 | link k p x k' k't othertree =
  23 |   let m = branchMask k k'
  24 |     in case zero m k' of
  25 |          True  =>
  26 |            Bin k p x m k't othertree
  27 |          False =>
  28 |            Bin k p x m othertree k't
  29 | 
  30 | ||| Internal function that merges two disjoint NatPSQ's that share the
  31 | ||| same prefix mask.
  32 | private
  33 | merge :  Ord p
  34 |       => Mask
  35 |       -> NatPSQ p v
  36 |       -> NatPSQ p v
  37 |       -> NatPSQ p v
  38 | merge m l r =
  39 |   case l of
  40 |     Nil          =>
  41 |       r
  42 |     Tip lk lp lx =>
  43 |       case r of
  44 |         Nil          =>
  45 |           l
  46 |         Tip rk rp rx =>
  47 |           case (lp, lk) < (rp, rk) of
  48 |             True  =>
  49 |               Bin lk lp lx m Nil r
  50 |             False =>
  51 |               Bin rk rp rx m l Nil
  52 |         Bin rk rp rx rm rl rr =>
  53 |           case (lp, lk) < (rp, rk) of
  54 |             True  =>
  55 |               Bin lk lp lx m Nil r
  56 |             False =>
  57 |               Bin rk rp rx m l (merge rm rl rr)
  58 |     Bin lk lp lx lm ll lr =>
  59 |       case r of
  60 |         Nil          =>
  61 |           l
  62 |         Tip rk rp rx =>
  63 |           case (lp, lk) < (rp, rk) of
  64 |             True  =>
  65 |               Bin lk lp lx m (merge lm ll lr) r
  66 |             False =>
  67 |               Bin rk rp rx m l Nil
  68 |         Bin rk rp rx rm rl rr =>
  69 |           case (lp, lk) < (rp, rk) of
  70 |             True  =>
  71 |               Bin lk lp lx m (merge lm ll lr) r
  72 |             False =>
  73 |               Bin rk rp rx m l (merge rm rl rr)
  74 | 
  75 | ||| Smart constructor for a Bin node.
  76 | private
  77 | bin :  Key
  78 |     -> p
  79 |     -> v
  80 |     -> Mask
  81 |     -> NatPSQ p v
  82 |     -> NatPSQ p v
  83 |     -> NatPSQ p v
  84 | bin k p x _ Nil Nil =
  85 |   Tip k p x
  86 | bin k p x m l   r   =
  87 |   Bin k p x m l r
  88 | 
  89 | ||| Internal function to insert a key that is *not* present in the priority queue.
  90 | export
  91 | unsafeInsertNew :  Ord p
  92 |                 => Key
  93 |                 -> p
  94 |                 -> v
  95 |                 -> NatPSQ p v
  96 |                 -> NatPSQ p v
  97 | unsafeInsertNew k p x t =
  98 |   case t of
  99 |     Nil          =>
 100 |       Tip k p x
 101 |     Tip k' p' x' =>
 102 |       case (p, k) < (p', k') of
 103 |         True  =>
 104 |           link k p x k' t Nil
 105 |         False =>
 106 |           link k' p' x' k (Tip k p x) Nil
 107 |     Bin k' p' x' m l r =>
 108 |       case noMatch k k' m of
 109 |         True  =>
 110 |           case (p, k) < (p', k') of
 111 |             True  =>
 112 |               link k p x k' t Nil
 113 |             False =>
 114 |               link k' p' x' k (Tip k p x) (merge m l r)
 115 |         False =>
 116 |           case (p, k) < (p', k') of
 117 |             True  =>
 118 |               case zero k' m of
 119 |                 True  =>
 120 |                   Bin k p x m (unsafeInsertNew k' p' x' l) r
 121 |                 False =>
 122 |                   Bin k p x m l (unsafeInsertNew k' p' x' r)
 123 |             False =>
 124 |               case zero k m of
 125 |                 True =>
 126 |                   Bin k' p' x' m (unsafeInsertNew k p x l) r
 127 |                 False =>
 128 |                   Bin k' p' x' m l (unsafeInsertNew k p x r)
 129 | 
 130 | --------------------------------------------------------------------------------
 131 | --          Construction
 132 | --------------------------------------------------------------------------------
 133 | 
 134 | ||| The empty queue. O(1)
 135 | export
 136 | empty : NatPSQ p v
 137 | empty =
 138 |   Nil
 139 | 
 140 | ||| Build a queue with one element. O(1)
 141 | export
 142 | singleton :  Ord p
 143 |           => Nat
 144 |           -> p
 145 |           -> v
 146 |           -> NatPSQ p v
 147 | singleton k p v =
 148 |   Tip k p v
 149 | 
 150 | --------------------------------------------------------------------------------
 151 | --          Query
 152 | --------------------------------------------------------------------------------
 153 | 
 154 | ||| Is the queue empty? O(1)
 155 | export
 156 | null :  NatPSQ p v
 157 |      -> Bool
 158 | null Nil =
 159 |   True
 160 | null _   =
 161 |   False
 162 | 
 163 | ||| The number of elements in a queue. O(1)
 164 | export
 165 | size :  NatPSQ p v
 166 |      -> Nat
 167 | size Nil               =
 168 |   0
 169 | size (Tip _ _ _)       =
 170 |   1
 171 | size (Bin _ _ _ _ l r) =
 172 |   1 + size l + size r
 173 | 
 174 | ||| The priority and value of a given key, or Nothing
 175 | ||| if the key is not bound. O(log n)
 176 | export
 177 | lookup :  Nat
 178 |        -> NatPSQ p v
 179 |        -> Maybe (p, v)
 180 | lookup key =
 181 |   go
 182 |   where
 183 |     go :  NatPSQ p v
 184 |        -> Maybe (p, v)
 185 |     go Nil            =
 186 |           Nothing
 187 |     go (Tip k' p' x') =
 188 |       case key == k' of
 189 |         True  =>
 190 |           Just (p', x')
 191 |         False =>
 192 |           Nothing
 193 |     go (Bin k' p' x' m l r) =
 194 |       case noMatch key k' m of
 195 |         True  =>
 196 |           Nothing
 197 |         False =>
 198 |           case key == k' of
 199 |             True  =>
 200 |               Just (p', x')
 201 |             False =>
 202 |               case zero key m of
 203 |                 True  =>
 204 |                   go l
 205 |                 False =>
 206 |                   go r
 207 | 
 208 | ||| Check if a key is present in the queue. O(log n)
 209 | export
 210 | member :  Nat
 211 |        -> NatPSQ p v
 212 |        -> Bool
 213 | member k =
 214 |   isJust . lookup k
 215 | 
 216 | ||| Is the key not a member of the queue? See also member. O(log n)
 217 | export
 218 | notMember :  Nat
 219 |           -> NatPSQ p v
 220 |           -> Bool
 221 | notMember k m =
 222 |   not $
 223 |     member k m
 224 | 
 225 | ||| The element with the lowest priority. O(1)
 226 | export
 227 | findMin :  NatPSQ p v
 228 |         -> Maybe (Nat, p, v)
 229 | findMin Nil               =
 230 |   Nothing
 231 | findMin (Tip k p x)       =
 232 |   Just (k, p, x)
 233 | findMin (Bin k p x _ _ _) =
 234 |   Just (k, p, x)
 235 | 
 236 | --------------------------------------------------------------------------------
 237 | --          Views
 238 | --------------------------------------------------------------------------------
 239 | 
 240 | ||| Delete a key and its priority and value from the queue. If
 241 | ||| the key was present, the associated priority and value are returned in
 242 | ||| addition to the updated queue. O(min(n, W))
 243 | export
 244 | deleteView :  Ord p
 245 |            => Nat
 246 |            -> NatPSQ p v
 247 |            -> Maybe (p, v, NatPSQ p v)
 248 | deleteView k t =
 249 |     case delFrom t of
 250 |       (_, Nothing)      =>
 251 |         Nothing
 252 |       (t', Just (p, x)) =>
 253 |         Just (p, x, t')
 254 |   where
 255 |     delFrom :  NatPSQ p v
 256 |             -> (NatPSQ p v, Maybe (p, v))
 257 |     delFrom Nil            =
 258 |       (Nil, Nothing)
 259 |     delFrom (Tip k' p' x') =
 260 |       case k == k' of
 261 |         True  =>
 262 |           (Nil, Just (p', x'))
 263 |         False =>
 264 |           (t, Nothing)
 265 |     delFrom (Bin k' p' x' m l r) =
 266 |       case noMatch k k' m of
 267 |         True  =>
 268 |           (t, Nothing)
 269 |         False =>
 270 |           case k == k' of
 271 |             True  =>
 272 |               let t' = merge m l r
 273 |                 in (t', Just (p', x'))
 274 |             False =>
 275 |               case zero k m of
 276 |                 True  =>
 277 |                   let (l', mbpx) = delFrom l
 278 |                       t'         = bin k' p' x' m l' r
 279 |                     in (t', mbpx)
 280 |                 False =>
 281 |                   let (r', mbpx) = delFrom r
 282 |                       t'         = bin k' p' x' m l  r'
 283 |                     in (t', mbpx)
 284 | 
 285 | ||| Insert a new key, priority and value into the queue. If the key
 286 | ||| is already present in the queue, then the evicted priority and value can be
 287 | ||| found the first element of the returned tuple. O(min(n, W))
 288 | export
 289 | insertView :  Ord p
 290 |            => Nat
 291 |            -> p
 292 |            -> v
 293 |            -> NatPSQ p v
 294 |            -> (Maybe (p, v), NatPSQ p v)
 295 | insertView k p x t =
 296 |   case deleteView k t of
 297 |     Nothing           =>
 298 |       (Nothing, unsafeInsertNew k p x t)
 299 |     Just (p', v', t') =>
 300 |       (Just (p', v'), unsafeInsertNew k p x t')
 301 | 
 302 | ||| Retrieve the binding with the least priority, and the
 303 | ||| rest of the queue stripped of that binding. O(min(n, W))
 304 | export
 305 | minView :  Ord p
 306 |         => NatPSQ p v
 307 |         -> Maybe (Nat, p, v, NatPSQ p v)
 308 | minView Nil               =
 309 |   Nothing
 310 | minView (Tip k p x)       =
 311 |   Just (k, p, x, Nil)
 312 | minView (Bin k p x m l r) =
 313 |   Just (k, p, x, merge m l r)
 314 | 
 315 | ||| Return a list of elements ordered by key whose priorities are at most @pt@,
 316 | ||| and the rest of the queue stripped of these elements.  The returned list of
 317 | ||| elements can be in any order: no guarantees there.
 318 | export
 319 | atMostView :  Ord p
 320 |            => p
 321 |            -> NatPSQ p v
 322 |            -> (List (Nat, p, v), NatPSQ p v)
 323 | atMostView pt t =
 324 |   go Nil t
 325 |   where
 326 |     go :  List (Key, p, v)
 327 |        -> NatPSQ p v
 328 |        -> (List (Key, p, v), NatPSQ p v)
 329 |     go acc Nil         =
 330 |       (acc, t)
 331 |     go acc (Tip k p x) =
 332 |       case p > pt of
 333 |         True  =>
 334 |           (acc, t)
 335 |         False =>
 336 |           ((k, p, x) :: acc, Nil)
 337 |     go acc (Bin k p x m l r) =
 338 |       case p > pt of
 339 |         True  =>
 340 |           (acc, t)
 341 |         False =>
 342 |           let (acc',  l') = go acc  l
 343 |               (acc'', r') = go acc' r
 344 |             in  ((k, p, x) :: acc'', merge m l' r')
 345 | 
 346 | --------------------------------------------------------------------------------
 347 | --          Delete
 348 | --------------------------------------------------------------------------------
 349 | 
 350 | ||| Delete a key and its priority and value from the queue.
 351 | ||| When the key is not a member of the queue,
 352 | ||| the original queue is returned. O(min(n, W))
 353 | export
 354 | delete :  Ord p
 355 |        => Nat
 356 |        -> NatPSQ p v
 357 |        -> NatPSQ p v
 358 | delete key =
 359 |   go
 360 |   where
 361 |     go :  NatPSQ p v
 362 |        -> NatPSQ p v
 363 |     go t =
 364 |       case t of
 365 |         Nil        =>
 366 |           Nil
 367 |         Tip k' _ _ =>
 368 |           case key == k' of
 369 |             True  =>
 370 |               Nil
 371 |             False =>
 372 |               t
 373 |         Bin k' p' x' m l r =>
 374 |           case noMatch key k' m of
 375 |             True  =>
 376 |               t
 377 |             False =>
 378 |               case key == k' of
 379 |                 True  =>
 380 |                   merge m l r
 381 |                 False =>
 382 |                   case zero key m of
 383 |                     True  =>
 384 |                       bin k' p' x' m (go l) r
 385 |                     False =>
 386 |                       bin k' p' x' m l (go r)
 387 | 
 388 | ||| Delete the binding with the least priority, and return the
 389 | ||| rest of the queue stripped of that binding.
 390 | ||| In case the queue is empty, the empty queue is returned again. O(min(n, w))
 391 | export
 392 | deleteMin :  Ord p
 393 |           => NatPSQ p v
 394 |           -> NatPSQ p v
 395 | deleteMin t =
 396 |   case minView t of
 397 |     Nothing            =>
 398 |       t
 399 |     Just (_, _, _, t') =>
 400 |       t'
 401 | 
 402 | --------------------------------------------------------------------------------
 403 | --          Insertion
 404 | --------------------------------------------------------------------------------
 405 | 
 406 | ||| Insert a new key, priority and value into the queue. If the key
 407 | ||| is already present in the queue, the associated priority and value are
 408 | ||| replaced with the supplied priority and value. O(min(n, W))
 409 | export
 410 | insert :  Ord p
 411 |        => Nat
 412 |        -> p
 413 |        -> v
 414 |        -> NatPSQ p v
 415 |        -> NatPSQ p v
 416 | insert k p x t =
 417 |   unsafeInsertNew k p x (delete k t)
 418 | 
 419 | --------------------------------------------------------------------------------
 420 | --          Alter
 421 | --------------------------------------------------------------------------------
 422 | 
 423 | ||| The expression alter f k queue alters the value x at k,
 424 | ||| or absence thereof.
 425 | ||| alter can be used to insert, delete, or update a value
 426 | ||| in a queue.
 427 | ||| It also allows you to calculate an additional value b. O(min(n, W))
 428 | export
 429 | alter :  Ord p
 430 |       => (Maybe (p, v) -> (b, Maybe (p, v)))
 431 |       -> Nat
 432 |       -> NatPSQ p v
 433 |       -> (b, NatPSQ p v)
 434 | alter f k t =
 435 |   let (t', mbx) = case deleteView k t of
 436 |                     Nothing          =>
 437 |                       (t, Nothing)
 438 |                     Just (p, v, t'') =>
 439 |                       (t'', Just (p, v))
 440 |       (b, mbx') = f mbx
 441 |     in case mbx' of
 442 |          Nothing     =>
 443 |            (b, t')
 444 |          Just (p, v) =>
 445 |            let t'' = unsafeInsertNew k p v t'
 446 |              in (b, t'')
 447 | 
 448 | ||| A variant of alter which works on the element with the
 449 | ||| minimum priority. Unlike alter,
 450 | ||| this variant also allows you to change the
 451 | ||| key of the element. O(min(n, W))
 452 | export
 453 | alterMin :  Ord p
 454 |          => (Maybe (Nat, p, v) -> (b, Maybe (Nat, p, v)))
 455 |          -> NatPSQ p v
 456 |          -> (b, NatPSQ p v)
 457 | alterMin f Nil         =
 458 |   case f Nothing of
 459 |     (b, Nothing)           =>
 460 |       (b, Nil)
 461 |     (b, Just (k', p', x')) =>
 462 |       (b, Tip k' p' x')
 463 | alterMin f (Tip k p x) =
 464 |   case f (Just (k, p, x)) of
 465 |     (b, Nothing)           =>
 466 |       (b, Nil)
 467 |     (b, Just (k', p', x')) =>
 468 |       (b, Tip k' p' x')
 469 | alterMin f (Bin k p x m l r) =
 470 |   case f (Just (k, p, x)) of
 471 |     (b, Nothing)           =>
 472 |       (b, merge m l r)
 473 |     (b, Just (k', p', x')) =>
 474 |       case k /= k' of
 475 |         True  =>
 476 |           (b, insert k' p' x' (merge m l r))
 477 |         False =>
 478 |           case p' <= p of
 479 |             True  =>
 480 |               (b, Bin k p' x' m l r)
 481 |             False =>
 482 |               (b, unsafeInsertNew k p' x' (merge m l r))
 483 | 
 484 | --------------------------------------------------------------------------------
 485 | --          Traversal
 486 | --------------------------------------------------------------------------------
 487 | 
 488 | ||| Modify every value in the queue. O(n)
 489 | export
 490 | map :  (v -> w)
 491 |     -> NatPSQ p v
 492 |     -> NatPSQ p w
 493 | map _ Nil               =
 494 |   Nil
 495 | map f (Tip k p x)       =
 496 |   Tip k p (f x)
 497 | map f (Bin k p x m l r) =
 498 |   Bin k p (f x) m (map f l) (map f r)
 499 | 
 500 | ||| Maps a function over the values and priorities of the queue.
 501 | ||| The function f must be monotonic with respect to the priorities.
 502 | ||| This means that if x < y, then fst (f k x v) < fst (f k y v).
 503 | ||| The precondition is not checked.
 504 | ||| If f is not monotonic, then the result
 505 | ||| will be invalid. O(n)
 506 | export
 507 | unsafeMapMonotonic :  (Key -> p -> v -> (q, w))
 508 |                    -> NatPSQ p v
 509 |                    -> NatPSQ q w
 510 | unsafeMapMonotonic f Nil               =
 511 |   Nil
 512 | unsafeMapMonotonic f (Tip k p x)       =
 513 |   let (p', x') = f k p x
 514 |     in Tip k p' x'
 515 | unsafeMapMonotonic f (Bin k p x m l r) =
 516 |   let (p', x') = f k p x
 517 |     in Bin k p' x' m (unsafeMapMonotonic f l) (unsafeMapMonotonic f r)
 518 | 
 519 | ||| Fold over every key, priority and value in the queue.
 520 | ||| The order in which the fold is performed is not specified. O(n)
 521 | export
 522 | fold :  (Nat -> p -> v -> a -> a)
 523 |      -> a
 524 |      -> NatPSQ p v
 525 |      -> a
 526 | fold f acc Nil                  =
 527 |   acc
 528 | fold f acc (Tip k' p' x')       =
 529 |   f k' p' x' acc
 530 | fold f acc (Bin k' p' x' m l r) =
 531 |   let acc1 = f k' p' x' acc
 532 |       acc2 = fold f acc1 l
 533 |       acc3 = fold f acc2 r
 534 |     in acc3
 535 | 
 536 | export
 537 | foldl :  (acc -> v -> acc)
 538 |       -> acc
 539 |       -> NatPSQ p v
 540 |       -> acc
 541 | foldl f acc Nil               =
 542 |   acc
 543 | foldl f acc (Tip _ _ v)       =
 544 |   f acc v
 545 | foldl f acc (Bin _ _ v _ l r) =
 546 |   foldl f (foldl f (f acc v) l) r
 547 | 
 548 | export
 549 | foldr :  (v -> acc -> acc)
 550 |       -> acc
 551 |       -> NatPSQ p v
 552 |       -> acc
 553 | foldr f acc Nil               =
 554 |   acc
 555 | foldr f acc (Tip _ _ v)       =
 556 |   f v acc
 557 | foldr f acc (Bin _ _ v _ l r) =
 558 |   foldr f (f v (foldr f acc r)) l
 559 | 
 560 | 
 561 | --------------------------------------------------------------------------------
 562 | --          Conversion
 563 | --------------------------------------------------------------------------------
 564 | 
 565 | ||| Build a queue from a list of (key, priority, value) tuples.
 566 | ||| If the list contains more than one priority and value for the same key, the
 567 | ||| last priority and value for the key is retained. O(n * min(n, W))
 568 | export
 569 | fromList :  Ord p
 570 |          => List (Nat, p, v)
 571 |          -> NatPSQ p v
 572 | fromList =
 573 |   foldl (\im, (k, p, x) => insert k p x im) empty
 574 | 
 575 | ||| Convert a queue to a list of (key, priority, value) tuples.
 576 | ||| The order of the list is not specified. O(n)
 577 | export
 578 | toList :  NatPSQ p v
 579 |        -> List (Nat, p, v)
 580 | toList Nil               =
 581 |   Nil
 582 | toList (Tip k p v)       =
 583 |   [(k, p, v)]
 584 | toList (Bin k p v _ l r) =
 585 |   (k, p, v) :: toList l ++ toList r
 586 | 
 587 | ||| Obtain the list of present keys in the queue. O(n)
 588 | export
 589 | keys :  NatPSQ p v
 590 |      -> List Nat
 591 | keys t =
 592 |   [k | (k, _, _) <- toList t]
 593 | 
 594 | --------------------------------------------------------------------------------
 595 | --          Interfaces
 596 | --------------------------------------------------------------------------------
 597 | 
 598 | export
 599 | Functor (NatPSQ p) where
 600 |   map = Data.NatPSQ.map
 601 | 
 602 | export
 603 | Foldable (NatPSQ p) where
 604 |   foldl = Data.NatPSQ.foldl
 605 |   foldr = Data.NatPSQ.foldr
 606 |   null  = Data.NatPSQ.null
 607 |