
   0 | ||| Hash Priority Search Queue
   1 | module Data.HashPSQ
   2 | 
   3 | import public Data.HashPSQ.Internal
   4 | import public Data.NatPSQ
   5 | import public Data.NatPSQ.Internal
   6 | import public Data.OrdPSQ
   7 | 
   8 | import Data.List
   9 | import Data.Hashable
  10 | import Data.Maybe
  11 | 
  12 | %default total
  13 | 
  14 | --------------------------------------------------------------------------------
  15 | --          Insertion
  16 | --------------------------------------------------------------------------------
  17 | 
  18 | private
  19 | ins :  Eq k
  20 |     => Ord k
  21 |     => Ord p
  22 |     => k
  23 |     -> p
  24 |     -> v
  25 |     -> Maybe (p, Bucket k p v)
  26 |     -> Maybe (p, Bucket k p v)
  27 | ins k p v Nothing                        =
  28 |   Just (p, MkBucket k v (OrdPSQ.empty))
  29 | ins k p v (Just (p', MkBucket k' v' os)) =
  30 |   case k' == k of
  31 |     True  =>
  32 |       -- Tricky: p might have less priority than an item in os.
  33 |       Just (mkBucket k p v os)
  34 |     False =>
  35 |       case p' < p || (p == p' && k' < k) of
  36 |         True  =>
  37 |           Just (p', MkBucket k' v' (OrdPSQ.insert k p v os))
  38 |         False =>
  39 |           case OrdPSQ.member k os of
  40 |             True  =>
  41 |               -- This is a bit tricky: k might already be present in 'os' and we
  42 |               -- don't want to end up with duplicate keys.
  43 |               Just (p, MkBucket k v (OrdPSQ.insert k' p' v' (OrdPSQ.delete k os)))
  44 |             False =>
  45 |               Just (p, MkBucket k v (OrdPSQ.insert k' p' v' os))
  46 | 
  47 | ||| Insert a new key, priority and value into the queue.
  48 | ||| If the key is already present in the queue,
  49 | ||| the associated priority and value are
  50 | ||| replaced with the supplied priority and value. O(min(n, W))
  51 | export
  52 | insert :  Hashable k
  53 |        => Ord k
  54 |        => Ord p
  55 |        => k
  56 |        -> p
  57 |        -> v
  58 |        -> HashPSQ k p v
  59 |        -> HashPSQ k p v
  60 | insert k p v (MkHashPSQ npsq) =
  61 |   MkHashPSQ $
  62 |     snd     $
  63 |       NatPSQ.alter (\x => ((), ins k p v x)) (cast $ hash k) npsq
  64 | 
  65 | --------------------------------------------------------------------------------
  66 | --          Construction
  67 | --------------------------------------------------------------------------------
  68 | 
  69 | ||| The empty queue. O(1)
  70 | export
  71 | empty : HashPSQ k p v
  72 | empty =
  73 |   MkHashPSQ NatPSQ.empty
  74 | 
  75 | ||| Build a queue with one element. O(1)
  76 | export
  77 | singleton :  Hashable k
  78 |           => Ord k
  79 |           => Ord p
  80 |           => k
  81 |           -> p
  82 |           -> v
  83 |           -> HashPSQ k p v
  84 | singleton k p v =
  85 |   Data.HashPSQ.insert k p v empty
  86 | 
  87 | --------------------------------------------------------------------------------
  88 | --          Query
  89 | --------------------------------------------------------------------------------
  90 | 
  91 | ||| Is the queue empty? O(1)
  92 | export
  93 | null :  HashPSQ k p v
  94 |      -> Bool
  95 | null (MkHashPSQ npsq) =
  96 |   NatPSQ.null npsq
  97 | 
  98 | ||| The number of elements in a queue. O(n)
  99 | export
 100 | size :  HashPSQ k p v
 101 |      -> Nat
 102 | size (MkHashPSQ npsq) =
 103 |   NatPSQ.fold (\_, _, (MkBucket _ _ opsq), acc => 1 + OrdPSQ.size opsq + acc)
 104 |               0
 105 |               npsq
 106 | 
 107 | ||| The priority and value of a given key, or Nothing if the
 108 | ||| key is not bound. O(min(n ,W))
 109 | export
 110 | lookup :  Hashable k
 111 |        => Ord k
 112 |        => Ord p
 113 |        => k
 114 |        -> HashPSQ k p v
 115 |        -> Maybe (p, v)
 116 | lookup k (MkHashPSQ npsq) =
 117 |   let nsq' = NatPSQ.lookup (cast $ hash k) npsq
 118 |     in case nsq' of
 119 |          Nothing                      =>
 120 |            Nothing
 121 |          Just (p0, MkBucket k0 v0 os) =>
 122 |            case k0 == k of
 123 |              True  =>
 124 |                Just (p0, v0)
 125 |              False =>
 126 |                OrdPSQ.lookup k os
 127 | 
 128 | ||| Check if a key is present in the queue. O(min(n, W))
 129 | export
 130 | member :  Hashable k
 131 |        => Ord k
 132 |        => Ord p
 133 |        => k
 134 |        -> HashPSQ k p v
 135 |        -> Bool
 136 | member k =
 137 |   isJust . lookup k
 138 | 
 139 | ||| Is the key not a member of the queue? See also member. O(log n)
 140 | export
 141 | notMember :  Hashable k
 142 |           => Ord k
 143 |           => Ord p
 144 |           => k
 145 |           -> HashPSQ k p v
 146 |           -> Bool
 147 | notMember k m =
 148 |   not $
 149 |     member k m
 150 | 
 151 | ||| The element with the lowest priority. O(1)
 152 | export
 153 | findMin :  Hashable k
 154 |         => Ord k
 155 |         => Ord p
 156 |         => HashPSQ k p v
 157 |         -> Maybe (k, p, v)
 158 | findMin (MkHashPSQ npsq) =
 159 |   case NatPSQ.findMin npsq of
 160 |     Nothing                     =>
 161 |       Nothing
 162 |     Just (_, p, MkBucket k x _) =>
 163 |       Just (k, p, x)
 164 | 
 165 | --------------------------------------------------------------------------------
 166 | --          Views
 167 | --------------------------------------------------------------------------------
 168 | 
 169 | private
 170 | deleteV :  Eq k
 171 |         => Ord k
 172 |         => Ord p
 173 |         => k
 174 |         -> Maybe (p, Bucket k p v)
 175 |         -> (Maybe (p, v), Maybe (p, Bucket k p v))
 176 | deleteV _ Nothing                         =
 177 |   (Nothing, Nothing)
 178 | deleteV k (Just (p, MkBucket bk bx opsq)) =
 179 |   case k == bk of
 180 |     True  =>
 181 |       case OrdPSQ.minView opsq of
 182 |         Nothing                  =>
 183 |           (Just (p, bx), Nothing)
 184 |         Just (k', p', x', opsq') =>
 185 |           (Just (p, bx), Just (p', MkBucket k' x' opsq'))
 186 |     False =>
 187 |       case OrdPSQ.deleteView k opsq of
 188 |         Nothing              =>
 189 |           (Nothing,       Nothing)
 190 |         Just (p', x', opsq') =>
 191 |           (Just (p', x'), Just (p, MkBucket bk bx opsq'))
 192 | 
 193 | ||| Delete a key and its priority and value from the queue. If
 194 | ||| the key was present, the associated priority and value are returned in
 195 | ||| addition to the updated queue. O(min(n, W))
 196 | export
 197 | deleteView :  Hashable k
 198 |            => Ord k
 199 |            => Ord p
 200 |            => k
 201 |            -> HashPSQ k p v
 202 |            -> Maybe (p, v, HashPSQ k p v)
 203 | deleteView k (MkHashPSQ npsq) =
 204 |   case NatPSQ.alter (\x => deleteV k x) (cast $ hash k) npsq of
 205 |     (Nothing,     _    ) =>
 206 |       Nothing
 207 |     (Just (p, x), npsq') =>
 208 |       Just (p, x, MkHashPSQ npsq')
 209 | 
 210 | ||| Insert a new key, priority and value into the queue. If the key
 211 | ||| is already present in the queue, then the evicted priority and value can be
 212 | ||| found the first element of the returned tuple. O(min(n, W))
 213 | export
 214 | insertView :  Hashable k
 215 |            => Ord k
 216 |            => Ord p
 217 |            => k
 218 |            -> p
 219 |            -> v
 220 |            -> HashPSQ k p v
 221 |            -> (Maybe (p, v), HashPSQ k p v)
 222 | insertView k p x t =
 223 |   case deleteView k t of
 224 |     Nothing          =>
 225 |       (Nothing,       insert k p x t)
 226 |     Just (p', x', _) =>
 227 |       (Just (p', x'), insert k p x t)
 228 | 
 229 | private
 230 | minV :  Ord k
 231 |      => Ord p
 232 |      => Maybe (a, b, Bucket k p v)
 233 |      -> (Maybe (k, b, v), Maybe (a, p, Bucket k p v))
 234 | minV Nothing                        =
 235 |   (Nothing, Nothing)
 236 | minV (Just (h, p, MkBucket k x os)) =
 237 |   case OrdPSQ.minView os of
 238 |     Nothing                =>
 239 |       (Just (k, p, x), Nothing)
 240 |     Just (k', p', x', os') =>
 241 |       (Just (k, p, x), Just (h, p', MkBucket k' x' os'))
 242 | 
 243 | ||| Retrieve the binding with the least priority, and the
 244 | ||| rest of the queue stripped of that binding. O(min(n, W))
 245 | export
 246 | minView :  Hashable k
 247 |         => Ord k
 248 |         => Ord p
 249 |         => HashPSQ k p v
 250 |         -> Maybe (k, p, v, HashPSQ k p v)
 251 | minView (MkHashPSQ npsq) =
 252 |   case NatPSQ.alterMin minV npsq of
 253 |     (Nothing       , _    ) =>
 254 |       Nothing
 255 |     (Just (k, p, x), npsq') =>
 256 |       Just (k, p, x, MkHashPSQ npsq')
 257 | 
 258 | ||| Return a list of elements ordered by key whose priorities are at most pt,
 259 | ||| and the rest of the queue stripped of these elements.
 260 | ||| The returned list of elements can be in any order: no guarantees there.
 261 | export
 262 | atMostView :  Hashable k
 263 |            => Ord k
 264 |            => Ord p
 265 |            => p
 266 |            -> HashPSQ k p v
 267 |            -> (List (k, p, v), HashPSQ k p v)
 268 | atMostView pt (MkHashPSQ t0) =
 269 |   let -- First we use NatPSQ.atMostView to get a collection of buckets that have
 270 |       -- AT LEAST one element with a low priority.  Buckets will usually only
 271 |       -- contain a single element.
 272 |       (buckets, t1) = NatPSQ.atMostView pt t0
 273 |       -- We now need to run through the buckets.  This will give us a list of
 274 |       -- elements to return and a bunch of buckets to re-insert.
 275 |       (returns, reinserts) = go Nil Nil buckets
 276 |       -- Now we can do the re-insertion pass.
 277 |       t2 = foldl
 278 |              (\t, (p, b@(MkBucket k _ _)) => NatPSQ.unsafeInsertNew (cast $ hash k) p b t)
 279 |              t1
 280 |              reinserts
 281 | 
 282 |     in (returns, MkHashPSQ t2)
 283 |     where
 284 |       -- We use two accumulators, for returns and re-inserts.
 285 |       go :  List (k, p, v)
 286 |          -> List (p, Bucket k p v)
 287 |          -> List (a, p, Bucket k p v)
 288 |          -> (List (k, p, v), List (p, Bucket k p v))
 289 |       go rets reins []                                =
 290 |         (rets, reins)
 291 |       go rets reins ((_, p, MkBucket k v opsq) :: bs) =
 292 |           -- Note that 'elems' should be very small, ideally a null list.
 293 |           let (elems, opsq') = OrdPSQ.atMostView pt opsq
 294 |               rets'          = (k, p, v) :: elems ++ rets
 295 |               reins'         = case toBucket opsq' of
 296 |                                  Nothing      =>
 297 |                                    reins
 298 |                                  Just (p', b) =>
 299 |                                    ((p', b) :: reins)
 300 |             in go rets' reins' bs
 301 | 
 302 | --------------------------------------------------------------------------------
 303 | --          Delete
 304 | --------------------------------------------------------------------------------
 305 | 
 306 | ||| Delete a key and its priority and value from the queue.
 307 | ||| When the key is not a member of the queue, the original queue is returned. O(min(n, W))
 308 | export
 309 | delete :  Hashable k
 310 |        => Ord k
 311 |        => Ord p
 312 |        => k
 313 |        -> HashPSQ k p v
 314 |        -> HashPSQ k p v
 315 | delete k t =
 316 |   case deleteView k t of
 317 |     Nothing         =>
 318 |       t
 319 |     Just (_, _, t') =>
 320 |       t'
 321 | 
 322 | ||| Delete the binding with the least priority, and return the
 323 | ||| rest of the queue stripped of that binding.
 324 | ||| In case the queue is empty, the empty queue is returned again. O(min(n, W))
 325 | export
 326 | deleteMin :  Hashable k
 327 |           => Ord k
 328 |           => Ord p
 329 |           => HashPSQ k p v
 330 |           -> HashPSQ k p v
 331 | deleteMin t =
 332 |   case minView t of
 333 |     Nothing            =>
 334 |       t
 335 |     Just (_, _, _, t') =>
 336 |       t'
 337 | 
 338 | --------------------------------------------------------------------------------
 339 | --          Alter
 340 | --------------------------------------------------------------------------------
 341 | 
 342 | ||| The expression, alter f k queue, alters the value x at k, or absence thereof.
 343 | ||| alter can be used to insert, delete, or update a value in a queue.
 344 | ||| It also allows you to calculate an additional value b. O(min(n, w))
 345 | export
 346 | alter :  Hashable k
 347 |       => Ord k
 348 |       => Ord p
 349 |       => (Maybe (p, v) -> (b, Maybe (p, v)))
 350 |       -> k
 351 |       -> HashPSQ k p v
 352 |       -> (b, HashPSQ k p v)
 353 | alter f k (MkHashPSQ npsq) =
 354 |   let h = cast $ hash k
 355 |     in case NatPSQ.deleteView h npsq of
 356 |          Nothing =>
 357 |            case f Nothing of
 358 |              (b, Nothing)     =>
 359 |                (b, MkHashPSQ npsq)
 360 |              (b, Just (p, x)) =>
 361 |                (b, MkHashPSQ $ NatPSQ.unsafeInsertNew h p (MkBucket k x OrdPSQ.empty) npsq)
 362 |          Just (bp, MkBucket bk bx opsq, npsq') =>
 363 |            case k == bk of
 364 |              True  =>
 365 |                case f (Just (bp, bx)) of
 366 |                  (b, Nothing) =>
 367 |                    case toBucket opsq of
 368 |                      Nothing             =>
 369 |                        (b, MkHashPSQ npsq')
 370 |                      Just (bp', bucket') =>
 371 |                        (b, MkHashPSQ $ NatPSQ.unsafeInsertNew h bp' bucket' npsq')
 372 |                  (b, Just (p, x)) =>
 373 |                    let (bp', bucket') = mkBucket k p x opsq
 374 |                      in (b, MkHashPSQ $ NatPSQ.unsafeInsertNew h bp' bucket' npsq')
 375 |              False =>
 376 |                let (b, opsq')     = OrdPSQ.alter f k opsq
 377 |                    (bp', bucket') = mkBucket bk bp bx opsq'
 378 |                  in (b, MkHashPSQ $ NatPSQ.unsafeInsertNew h bp' bucket' npsq')
 379 | 
 380 | ||| A variant of alter which works on the element with the
 381 | ||| minimum priority.
 382 | ||| Unlike alter, this variant also allows you to change the key of the element. O(min(n, W))
 383 | export
 384 | alterMin :  Hashable k
 385 |          => Ord k
 386 |          => Ord p
 387 |          => (Maybe (k, p, v) -> (b, Maybe (k, p, v)))
 388 |          -> HashPSQ k p v
 389 |          -> (b, HashPSQ k p v)
 390 | alterMin f t0 =
 391 |   let (t, mbx)  = case minView t0 of
 392 |                     Nothing             =>
 393 |                       (t0, Nothing)
 394 |                     Just (k, p, x, t0') =>
 395 |                       (t0', Just (k, p, x))
 396 |       (b, mbx') = f mbx
 397 |     in case mbx' of
 398 |          Nothing        =>
 399 |            (b, t)
 400 |          Just (k, p, x) =>
 401 |            (b, insert k p x t)
 402 | 
 403 | --------------------------------------------------------------------------------
 404 | --          Traversal
 405 | --------------------------------------------------------------------------------
 406 | 
 407 | ||| Modify every value in the queue. O(n)
 408 | export
 409 | map :  (v -> w)
 410 |     -> HashPSQ k p v
 411 |     -> HashPSQ k p w
 412 | map f (MkHashPSQ npsq) =
 413 |   MkHashPSQ (NatPSQ.map mapBucket npsq)
 414 |   where
 415 |     mapBucket :  Bucket k p v
 416 |               -> Bucket k p w
 417 |     mapBucket (MkBucket k v opsq) =
 418 |       MkBucket k (f v) (OrdPSQ.map (\k', p', v' => f v') opsq)
 419 | 
 420 | --------------------------------------------------------------------------------
 421 | --          Conversion
 422 | --------------------------------------------------------------------------------
 423 | 
 424 | ||| Build a queue from a list of (key, priority, value) tuples.
 425 | ||| If the list contains more than one priority and value for the same key, the
 426 | ||| last priority and value for the key is retained. O(min(n, W))
 427 | export
 428 | fromList :  Hashable k
 429 |          => Ord k
 430 |          => Ord p
 431 |          => List (k, p, v)
 432 |          -> HashPSQ k p v
 433 | fromList =
 434 |   foldl (\psq, (k, p, x) => insert k p x psq) empty
 435 | 
 436 | ||| Convert a queue to a list of (key, priority, value) tuples. The
 437 | ||| order of the list is not specified. O(n)
 438 | export
 439 | toList :  Hashable k
 440 |        => Ord k
 441 |        => Ord p
 442 |        => HashPSQ k p v
 443 |        -> List (k, p, v)
 444 | toList (MkHashPSQ npsq) =
 445 |     [ (k', p', x')
 446 |     | (_, p, (MkBucket k x opsq)) <- NatPSQ.toList npsq
 447 |     , (k', p', x')                <- (k, p, x) :: OrdPSQ.toList opsq
 448 |     ]
 449 | 
 450 | ||| Obtain the list of present keys in the queue. O(n)
 451 | export
 452 | keys :  Hashable k
 453 |      => Ord k
 454 |      => Ord p
 455 |      => HashPSQ k p v
 456 |      -> List k
 457 | keys t =
 458 |   [k | (k, _, _) <- toList t]
 459 | 
 460 | --------------------------------------------------------------------------------
 461 | --          Interfaces
 462 | --------------------------------------------------------------------------------
 463 | 
 464 | export
 465 | Functor (HashPSQ k p) where
 466 |   map = Data.HashPSQ.map
 467 |