1{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
2module 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
38import Control.Arrow (first, second)
39
40import qualified Data.Foldable as F
41import Data.Function
42import qualified Data.List as S
43import Data.Ord (comparing)
44import Data.Traversable
45import Prelude hiding (foldr, length, lookup, filter, null, map)
46
47newtype PSQ k v = PSQ [(k, v)]
48  deriving (Eq, Show, Functor, F.Foldable, Traversable) -- Qualified Foldable to avoid issues with FTP
49
50keys :: PSQ k v -> [k]
51keys (PSQ xs) = fmap fst xs
52
53lookup :: Eq k => k -> PSQ k v -> Maybe v
54lookup k (PSQ xs) = S.lookup k xs
55
56map :: (v1 -> v2) -> PSQ k v1 -> PSQ k v2
57map f (PSQ xs) = PSQ (fmap (second f) xs)
58
59mapKeys :: (k1 -> k2) -> PSQ k1 v -> PSQ k2 v
60mapKeys f (PSQ xs) = PSQ (fmap (first f) xs)
61
62mapWithKey :: (k -> a -> b) -> PSQ k a -> PSQ k b
63mapWithKey f (PSQ xs) = PSQ (fmap (\ (k, v) -> (k, f k v)) xs)
64
65fromList :: [(k, a)] -> PSQ k a
66fromList = PSQ
67
68cons :: k -> a -> PSQ k a -> PSQ k a
69cons k x (PSQ xs) = PSQ ((k, x) : xs)
70
71snoc :: PSQ k a -> k -> a -> PSQ k a
72snoc (PSQ xs) k x = PSQ (xs ++ [(k, x)])
73
74casePSQ :: PSQ k a -> r -> (k -> a -> PSQ k a -> r) -> r
75casePSQ (PSQ xs) n c =
76  case xs of
77    []          -> n
78    (k, v) : ys -> c k v (PSQ ys)
79
80sortBy :: (a -> a -> Ordering) -> PSQ k a -> PSQ k a
81sortBy cmp (PSQ xs) = PSQ (S.sortBy (cmp `on` snd) xs)
82
83sortByKeys :: (k -> k -> Ordering) -> PSQ k a -> PSQ k a
84sortByKeys cmp (PSQ xs) = PSQ (S.sortBy (cmp `on` fst) xs)
85
86maximumBy :: (k -> Int) -> PSQ k a -> (k, a)
87maximumBy sel (PSQ xs) =
88  S.minimumBy (flip (comparing (sel . fst))) xs
89
90minimumBy :: (a -> Int) -> PSQ k a -> PSQ k a
91minimumBy 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.
95prefer :: (a -> Bool) -> PSQ k a -> PSQ k a
96prefer p = sortBy $ flip (comparing p)
97
98-- | Sort the list so that keys satisfying the predicate are first.
99preferByKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
100preferByKeys p = sortByKeys $ flip (comparing p)
101
102-- | Will partition the list according to the predicate. If
103-- there is any element that satisfies the precidate, then only
104-- the elements satisfying the predicate are returned.
105-- Otherwise, the rest is returned.
106--
107filterIfAny :: (a -> Bool) -> PSQ k a -> PSQ k a
108filterIfAny 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--
117filterIfAnyByKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
118filterIfAnyByKeys 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
124filterKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
125filterKeys p (PSQ xs) = PSQ (S.filter (p . fst) xs)
126
127filter :: (a -> Bool) -> PSQ k a -> PSQ k a
128filter p (PSQ xs) = PSQ (S.filter (p . snd) xs)
129
130length :: PSQ k a -> Int
131length (PSQ xs) = S.length xs
132
133null :: PSQ k a -> Bool
134null (PSQ xs) = S.null xs
135
136isZeroOrOne :: PSQ k a -> Bool
137isZeroOrOne (PSQ [])  = True
138isZeroOrOne (PSQ [_]) = True
139isZeroOrOne _         = False
140
141firstOnly :: PSQ k a -> PSQ k a
142firstOnly (PSQ [])      = PSQ []
143firstOnly (PSQ (x : _)) = PSQ [x]
144
145toList :: PSQ k a -> [(k, a)]
146toList (PSQ xs) = xs
147
148union :: PSQ k a -> PSQ k a -> PSQ k a
149union (PSQ xs) (PSQ ys) = PSQ (xs ++ ys)
150