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