cabal-install-solver/src/Distribution/Solver/Modular/WeightedPSQ.hs

   1 {-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
   2 {-# LANGUAGE ScopedTypeVariables #-}
   3 {-# LANGUAGE TupleSections #-}
   4 module Distribution.Solver.Modular.WeightedPSQ (
   5     WeightedPSQ
   6   , fromList
   7   , toList
   8   , keys
   9   , weights
  10   , isZeroOrOne
  11   , filter
  12   , lookup
  13   , mapWithKey
  14   , mapWeightsWithKey
  15   , traverseWithKey
  16   , union
  17   , takeUntil
  18   ) where
  19
  20 import qualified Data.Foldable as F
  21 import qualified Data.List as L
  22 import Data.Ord (comparing)
  23 import qualified Data.Traversable as T
  24 import Prelude hiding (filter, lookup)
  25
  26 -- | An association list that is sorted by weight.
  27 --
  28 -- Each element has a key ('k'), value ('v'), and weight ('w'). All operations
  29 -- that add elements or modify weights stably sort the elements by weight.
  30 newtype WeightedPSQ w k v = WeightedPSQ [(w, k, v)]
  31   deriving (Eq, Show, Functor, F.Foldable, T.Traversable)
  32
  33 -- | /O(N)/.
  34 filter :: (v -> Bool) -> WeightedPSQ k w v -> WeightedPSQ k w v
  35 filter p (WeightedPSQ xs) = WeightedPSQ (L.filter (p . triple_3) xs)
  36
  37 -- | /O(1)/. Return @True@ if the @WeightedPSQ@ contains zero or one elements.
  38 isZeroOrOne :: WeightedPSQ w k v -> Bool
  39 isZeroOrOne (WeightedPSQ [])  = True
  40 isZeroOrOne (WeightedPSQ [_]) = True
  41 isZeroOrOne _                 = False
  42
  43 -- | /O(1)/. Return the elements in order.
  44 toList :: WeightedPSQ w k v -> [(w, k, v)]
  45 toList (WeightedPSQ xs) = xs
  46
  47 -- | /O(N log N)/.
  48 fromList :: Ord w => [(w, k, v)] -> WeightedPSQ w k v
  49 fromList = WeightedPSQ . L.sortBy (comparing triple_1)
  50
  51 -- | /O(N)/. Return the weights in order.
  52 weights :: WeightedPSQ w k v -> [w]
  53 weights (WeightedPSQ xs) = L.map triple_1 xs
  54
  55 -- | /O(N)/. Return the keys in order.
  56 keys :: WeightedPSQ w k v -> [k]
  57 keys (WeightedPSQ xs) = L.map triple_2 xs
  58
  59 -- | /O(N)/. Return the value associated with the first occurrence of the give
  60 -- key, if it exists.
  61 lookup :: Eq k => k -> WeightedPSQ w k v -> Maybe v
  62 lookup k (WeightedPSQ xs) = triple_3 `fmap` L.find ((k ==) . triple_2) xs
  63
  64 -- | /O(N log N)/. Update the weights.
  65 mapWeightsWithKey :: Ord w2
  66                   => (k -> w1 -> w2)
  67                   -> WeightedPSQ w1 k v
  68                   -> WeightedPSQ w2 k v
  69 mapWeightsWithKey f (WeightedPSQ xs) = fromList $
  70                                        L.map (\ (w, k, v) -> (f k w, k, v)) xs
  71
  72 -- | /O(N)/. Update the values.
  73 mapWithKey :: (k -> v1 -> v2) -> WeightedPSQ w k v1 -> WeightedPSQ w k v2
  74 mapWithKey f (WeightedPSQ xs) = WeightedPSQ $
  75                                 L.map (\ (w, k, v) -> (w, k, f k v)) xs
  76
  77 -- | /O(N)/. Traverse and update values in some applicative functor.
  78 traverseWithKey
  79   :: Applicative f
  80   => (k -> v -> f v')
  81   -> WeightedPSQ w k v
  82   -> f (WeightedPSQ w k v')
  83 traverseWithKey f (WeightedPSQ q) = WeightedPSQ <$>
  84   traverse (\(w,k,v) -> (w,k,) <$> f k v) q
  85
  86 -- | /O((N + M) log (N + M))/. Combine two @WeightedPSQ@s, preserving all
  87 -- elements. Elements from the first @WeightedPSQ@ come before elements in the
  88 -- second when they have the same weight.
  89 union :: Ord w => WeightedPSQ w k v -> WeightedPSQ w k v -> WeightedPSQ w k v
  90 union (WeightedPSQ xs) (WeightedPSQ ys) = fromList (xs ++ ys)
  91
  92 -- | /O(N)/. Return the prefix of values ending with the first element that
  93 -- satisfies p, or all elements if none satisfy p.
  94 takeUntil :: forall w k v. (v -> Bool) -> WeightedPSQ w k v -> WeightedPSQ w k v
  95 takeUntil p (WeightedPSQ xs) = WeightedPSQ (go xs)
  96   where
  97     go :: [(w, k, v)] -> [(w, k, v)]
  98     go []       = []
  99     go (y : ys) = y : if p (triple_3 y) then [] else go ys
 100
 101 triple_1 :: (x, y, z) -> x
 102 triple_1 (x, _, _) = x
 103
 104 triple_2 :: (x, y, z) -> y
 105 triple_2 (_, y, _) = y
 106
 107 triple_3 :: (x, y, z) -> z
 108 triple_3 (_, _, z) = z