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

   1 {-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
   2 module Distribution.Solver.Modular.PSQ
   3     ( PSQ(..)  -- Unit test needs constructor access
   4     , casePSQ
   5     , cons
   6     , length
   7     , lookup
   8     , filter
   9     , filterIfAny
  10     , filterIfAnyByKeys
  11     , filterKeys
  12     , firstOnly
  13     , fromList
  14     , isZeroOrOne
  15     , keys
  16     , map
  17     , mapKeys
  18     , mapWithKey
  19     , maximumBy
  20     , minimumBy
  21     , null
  22     , prefer
  23     , preferByKeys
  24     , snoc
  25     , sortBy
  26     , sortByKeys
  27     , toList
  28     , union
  29     ) where
  30
  31 -- Priority search queues.
  32 --
  33 -- I am not yet sure what exactly is needed. But we need a data structure with
  34 -- key-based lookup that can be sorted. We're using a sequence right now with
  35 -- (inefficiently implemented) lookup, because I think that queue-based
  36 -- operations and sorting turn out to be more efficiency-critical in practice.
  37
  38 import Control.Arrow (first, second)
  39
  40 import qualified Data.Foldable as F
  41 import Data.Function
  42 import qualified Data.List as S
  43 import Data.Ord (comparing)
  44 import Data.Traversable
  45 import Prelude hiding (foldr, length, lookup, filter, null, map)
  46
  47 newtype PSQ k v = PSQ [(k, v)]
  48   deriving (Eq, Show, Functor, F.Foldable, Traversable) -- Qualified Foldable to avoid issues with FTP
  49
  50 keys :: PSQ k v -> [k]
  51 keys (PSQ xs) = fmap fst xs
  52
  53 lookup :: Eq k => k -> PSQ k v -> Maybe v
  54 lookup k (PSQ xs) = S.lookup k xs
  55
  56 map :: (v1 -> v2) -> PSQ k v1 -> PSQ k v2
  57 map f (PSQ xs) = PSQ (fmap (second f) xs)
  58
  59 mapKeys :: (k1 -> k2) -> PSQ k1 v -> PSQ k2 v
  60 mapKeys f (PSQ xs) = PSQ (fmap (first f) xs)
  61
  62 mapWithKey :: (k -> a -> b) -> PSQ k a -> PSQ k b
  63 mapWithKey f (PSQ xs) = PSQ (fmap (\ (k, v) -> (k, f k v)) xs)
  64
  65 fromList :: [(k, a)] -> PSQ k a
  66 fromList = PSQ
  67
  68 cons :: k -> a -> PSQ k a -> PSQ k a
  69 cons k x (PSQ xs) = PSQ ((k, x) : xs)
  70
  71 snoc :: PSQ k a -> k -> a -> PSQ k a
  72 snoc (PSQ xs) k x = PSQ (xs ++ [(k, x)])
  73
  74 casePSQ :: PSQ k a -> r -> (k -> a -> PSQ k a -> r) -> r
  75 casePSQ (PSQ xs) n c =
  76   case xs of
  77     []          -> n
  78     (k, v) : ys -> c k v (PSQ ys)
  79
  80 sortBy :: (a -> a -> Ordering) -> PSQ k a -> PSQ k a
  81 sortBy cmp (PSQ xs) = PSQ (S.sortBy (cmp `on` snd) xs)
  82
  83 sortByKeys :: (k -> k -> Ordering) -> PSQ k a -> PSQ k a
  84 sortByKeys cmp (PSQ xs) = PSQ (S.sortBy (cmp `on` fst) xs)
  85
  86 maximumBy :: (k -> Int) -> PSQ k a -> (k, a)
  87 maximumBy sel (PSQ xs) =
  88   S.minimumBy (flip (comparing (sel . fst))) xs
  89
  90 minimumBy :: (a -> Int) -> PSQ k a -> PSQ k a
  91 minimumBy sel (PSQ xs) =
  92   PSQ [snd (S.minimumBy (comparing fst) (S.map (\ x -> (sel (snd x), x)) xs))]
  93
  94 -- | Sort the list so that values satisfying the predicate are first.
  95 prefer :: (a -> Bool) -> PSQ k a -> PSQ k a
  96 prefer p = sortBy $ flip (comparing p)
  97
  98 -- | Sort the list so that keys satisfying the predicate are first.
  99 preferByKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
 100 preferByKeys p = sortByKeys $ flip (comparing p)
 101
 102 -- | Will partition the list according to the predicate. If
 103 -- there is any element that satisfies the predicate, then only
 104 -- the elements satisfying the predicate are returned.
 105 -- Otherwise, the rest is returned.
 106 --
 107 filterIfAny :: (a -> Bool) -> PSQ k a -> PSQ k a
 108 filterIfAny p (PSQ xs) =
 109   let
 110     (pro, con) = S.partition (p . snd) xs
 111   in
 112     if S.null pro then PSQ con else PSQ pro
 113
 114 -- | Variant of 'filterIfAny' that takes a predicate on the keys
 115 -- rather than on the values.
 116 --
 117 filterIfAnyByKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
 118 filterIfAnyByKeys p (PSQ xs) =
 119   let
 120     (pro, con) = S.partition (p . fst) xs
 121   in
 122     if S.null pro then PSQ con else PSQ pro
 123
 124 filterKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
 125 filterKeys p (PSQ xs) = PSQ (S.filter (p . fst) xs)
 126
 127 filter :: (a -> Bool) -> PSQ k a -> PSQ k a
 128 filter p (PSQ xs) = PSQ (S.filter (p . snd) xs)
 129
 130 length :: PSQ k a -> Int
 131 length (PSQ xs) = S.length xs
 132
 133 null :: PSQ k a -> Bool
 134 null (PSQ xs) = S.null xs
 135
 136 isZeroOrOne :: PSQ k a -> Bool
 137 isZeroOrOne (PSQ [])  = True
 138 isZeroOrOne (PSQ [_]) = True
 139 isZeroOrOne _         = False
 140
 141 firstOnly :: PSQ k a -> PSQ k a
 142 firstOnly (PSQ [])      = PSQ []
 143 firstOnly (PSQ (x : _)) = PSQ [x]
 144
 145 toList :: PSQ k a -> [(k, a)]
 146 toList (PSQ xs) = xs
 147
 148 union :: PSQ k a -> PSQ k a -> PSQ k a
 149 union (PSQ xs) (PSQ ys) = PSQ (xs ++ ys)