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