{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
module Distribution.Solver.Modular.PSQ
    ( PSQ(..)  -- Unit test needs constructor access
    , casePSQ
    , cons
    , length
    , lookup
    , filter
    , filterIfAny
    , filterIfAnyByKeys
    , filterKeys
    , firstOnly
    , fromList
    , isZeroOrOne
    , keys
    , map
    , mapKeys
    , mapWithKey
    , maximumBy
    , minimumBy
    , null
    , prefer
    , preferByKeys
    , snoc
    , sortBy
    , sortByKeys
    , toList
    , union
    ) where

-- Priority search queues.
--
-- I am not yet sure what exactly is needed. But we need a data structure with
-- key-based lookup that can be sorted. We're using a sequence right now with
-- (inefficiently implemented) lookup, because I think that queue-based
-- operations and sorting turn out to be more efficiency-critical in practice.

import Control.Arrow (first, second)

import qualified Data.Foldable as F
import Data.Function
import qualified Data.List as S
import Data.Ord (comparing)
import Data.Traversable
import Prelude hiding (foldr, length, lookup, filter, null, map)

newtype PSQ k v = PSQ [(k, v)]
  deriving (PSQ k v -> PSQ k v -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => PSQ k v -> PSQ k v -> Bool
/= :: PSQ k v -> PSQ k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => PSQ k v -> PSQ k v -> Bool
== :: PSQ k v -> PSQ k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => PSQ k v -> PSQ k v -> Bool
Eq, Int -> PSQ k v -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> PSQ k v -> ShowS
forall k v. (Show k, Show v) => [PSQ k v] -> ShowS
forall k v. (Show k, Show v) => PSQ k v -> String
showList :: [PSQ k v] -> ShowS
$cshowList :: forall k v. (Show k, Show v) => [PSQ k v] -> ShowS
show :: PSQ k v -> String
$cshow :: forall k v. (Show k, Show v) => PSQ k v -> String
showsPrec :: Int -> PSQ k v -> ShowS
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> PSQ k v -> ShowS
Show, forall a b. a -> PSQ k b -> PSQ k a
forall a b. (a -> b) -> PSQ k a -> PSQ k b
forall k a b. a -> PSQ k b -> PSQ k a
forall k a b. (a -> b) -> PSQ k a -> PSQ k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> PSQ k b -> PSQ k a
$c<$ :: forall k a b. a -> PSQ k b -> PSQ k a
fmap :: forall a b. (a -> b) -> PSQ k a -> PSQ k b
$cfmap :: forall k a b. (a -> b) -> PSQ k a -> PSQ k b
Functor, forall a. PSQ k a -> Bool
forall k a. Eq a => a -> PSQ k a -> Bool
forall k a. Num a => PSQ k a -> a
forall k a. Ord a => PSQ k a -> a
forall m a. Monoid m => (a -> m) -> PSQ k a -> m
forall k m. Monoid m => PSQ k m -> m
forall k a. PSQ k a -> Bool
forall k a. PSQ k a -> Int
forall k a. PSQ k a -> [a]
forall a b. (a -> b -> b) -> b -> PSQ k a -> b
forall k a. (a -> a -> a) -> PSQ k a -> a
forall k m a. Monoid m => (a -> m) -> PSQ k a -> m
forall k b a. (b -> a -> b) -> b -> PSQ k a -> b
forall k a b. (a -> b -> b) -> b -> PSQ k a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => PSQ k a -> a
$cproduct :: forall k a. Num a => PSQ k a -> a
sum :: forall a. Num a => PSQ k a -> a
$csum :: forall k a. Num a => PSQ k a -> a
minimum :: forall a. Ord a => PSQ k a -> a
$cminimum :: forall k a. Ord a => PSQ k a -> a
maximum :: forall a. Ord a => PSQ k a -> a
$cmaximum :: forall k a. Ord a => PSQ k a -> a
elem :: forall a. Eq a => a -> PSQ k a -> Bool
$celem :: forall k a. Eq a => a -> PSQ k a -> Bool
length :: forall a. PSQ k a -> Int
$clength :: forall k a. PSQ k a -> Int
null :: forall a. PSQ k a -> Bool
$cnull :: forall k a. PSQ k a -> Bool
toList :: forall a. PSQ k a -> [a]
$ctoList :: forall k a. PSQ k a -> [a]
foldl1 :: forall a. (a -> a -> a) -> PSQ k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> PSQ k a -> a
foldr1 :: forall a. (a -> a -> a) -> PSQ k a -> a
$cfoldr1 :: forall k a. (a -> a -> a) -> PSQ k a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> PSQ k a -> b
$cfoldl' :: forall k b a. (b -> a -> b) -> b -> PSQ k a -> b
foldl :: forall b a. (b -> a -> b) -> b -> PSQ k a -> b
$cfoldl :: forall k b a. (b -> a -> b) -> b -> PSQ k a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> PSQ k a -> b
$cfoldr' :: forall k a b. (a -> b -> b) -> b -> PSQ k a -> b
foldr :: forall a b. (a -> b -> b) -> b -> PSQ k a -> b
$cfoldr :: forall k a b. (a -> b -> b) -> b -> PSQ k a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> PSQ k a -> m
$cfoldMap' :: forall k m a. Monoid m => (a -> m) -> PSQ k a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> PSQ k a -> m
$cfoldMap :: forall k m a. Monoid m => (a -> m) -> PSQ k a -> m
fold :: forall m. Monoid m => PSQ k m -> m
$cfold :: forall k m. Monoid m => PSQ k m -> m
F.Foldable, forall k. Functor (PSQ k)
forall k. Foldable (PSQ k)
forall k (m :: * -> *) a. Monad m => PSQ k (m a) -> m (PSQ k a)
forall k (f :: * -> *) a.
Applicative f =>
PSQ k (f a) -> f (PSQ k a)
forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> PSQ k a -> m (PSQ k b)
forall k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PSQ k a -> f (PSQ k b)
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PSQ k a -> f (PSQ k b)
sequence :: forall (m :: * -> *) a. Monad m => PSQ k (m a) -> m (PSQ k a)
$csequence :: forall k (m :: * -> *) a. Monad m => PSQ k (m a) -> m (PSQ k a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> PSQ k a -> m (PSQ k b)
$cmapM :: forall k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> PSQ k a -> m (PSQ k b)
sequenceA :: forall (f :: * -> *) a. Applicative f => PSQ k (f a) -> f (PSQ k a)
$csequenceA :: forall k (f :: * -> *) a.
Applicative f =>
PSQ k (f a) -> f (PSQ k a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PSQ k a -> f (PSQ k b)
$ctraverse :: forall k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> PSQ k a -> f (PSQ k b)
Traversable) -- Qualified Foldable to avoid issues with FTP

keys :: PSQ k v -> [k]
keys :: forall k v. PSQ k v -> [k]
keys (PSQ [(k, v)]
xs) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst [(k, v)]
xs

lookup :: Eq k => k -> PSQ k v -> Maybe v
lookup :: forall k v. Eq k => k -> PSQ k v -> Maybe v
lookup k
k (PSQ [(k, v)]
xs) = forall a b. Eq a => a -> [(a, b)] -> Maybe b
S.lookup k
k [(k, v)]
xs

map :: (v1 -> v2) -> PSQ k v1 -> PSQ k v2
map :: forall v1 v2 k. (v1 -> v2) -> PSQ k v1 -> PSQ k v2
map v1 -> v2
f (PSQ [(k, v1)]
xs) = forall k v. [(k, v)] -> PSQ k v
PSQ (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second v1 -> v2
f) [(k, v1)]
xs)

mapKeys :: (k1 -> k2) -> PSQ k1 v -> PSQ k2 v
mapKeys :: forall k1 k2 v. (k1 -> k2) -> PSQ k1 v -> PSQ k2 v
mapKeys k1 -> k2
f (PSQ [(k1, v)]
xs) = forall k v. [(k, v)] -> PSQ k v
PSQ (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first k1 -> k2
f) [(k1, v)]
xs)

mapWithKey :: (k -> a -> b) -> PSQ k a -> PSQ k b
mapWithKey :: forall k a b. (k -> a -> b) -> PSQ k a -> PSQ k b
mapWithKey k -> a -> b
f (PSQ [(k, a)]
xs) = forall k v. [(k, v)] -> PSQ k v
PSQ (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ (k
k, a
v) -> (k
k, k -> a -> b
f k
k a
v)) [(k, a)]
xs)

fromList :: [(k, a)] -> PSQ k a
fromList :: forall k v. [(k, v)] -> PSQ k v
fromList = forall k v. [(k, v)] -> PSQ k v
PSQ

cons :: k -> a -> PSQ k a -> PSQ k a
cons :: forall k a. k -> a -> PSQ k a -> PSQ k a
cons k
k a
x (PSQ [(k, a)]
xs) = forall k v. [(k, v)] -> PSQ k v
PSQ ((k
k, a
x) forall a. a -> [a] -> [a]
: [(k, a)]
xs)

snoc :: PSQ k a -> k -> a -> PSQ k a
snoc :: forall k a. PSQ k a -> k -> a -> PSQ k a
snoc (PSQ [(k, a)]
xs) k
k a
x = forall k v. [(k, v)] -> PSQ k v
PSQ ([(k, a)]
xs forall a. [a] -> [a] -> [a]
++ [(k
k, a
x)])

casePSQ :: PSQ k a -> r -> (k -> a -> PSQ k a -> r) -> r
casePSQ :: forall k a r. PSQ k a -> r -> (k -> a -> PSQ k a -> r) -> r
casePSQ (PSQ [(k, a)]
xs) r
n k -> a -> PSQ k a -> r
c =
  case [(k, a)]
xs of
    []          -> r
n
    (k
k, a
v) : [(k, a)]
ys -> k -> a -> PSQ k a -> r
c k
k a
v (forall k v. [(k, v)] -> PSQ k v
PSQ [(k, a)]
ys)

sortBy :: (a -> a -> Ordering) -> PSQ k a -> PSQ k a
sortBy :: forall a k. (a -> a -> Ordering) -> PSQ k a -> PSQ k a
sortBy a -> a -> Ordering
cmp (PSQ [(k, a)]
xs) = forall k v. [(k, v)] -> PSQ k v
PSQ (forall a. (a -> a -> Ordering) -> [a] -> [a]
S.sortBy (a -> a -> Ordering
cmp forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> b
snd) [(k, a)]
xs)

sortByKeys :: (k -> k -> Ordering) -> PSQ k a -> PSQ k a
sortByKeys :: forall k a. (k -> k -> Ordering) -> PSQ k a -> PSQ k a
sortByKeys k -> k -> Ordering
cmp (PSQ [(k, a)]
xs) = forall k v. [(k, v)] -> PSQ k v
PSQ (forall a. (a -> a -> Ordering) -> [a] -> [a]
S.sortBy (k -> k -> Ordering
cmp forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> a
fst) [(k, a)]
xs)

maximumBy :: (k -> Int) -> PSQ k a -> (k, a)
maximumBy :: forall k a. (k -> Int) -> PSQ k a -> (k, a)
maximumBy k -> Int
sel (PSQ [(k, a)]
xs) =
  forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
S.minimumBy (forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (k -> Int
sel forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst))) [(k, a)]
xs

minimumBy :: (a -> Int) -> PSQ k a -> PSQ k a
minimumBy :: forall a k. (a -> Int) -> PSQ k a -> PSQ k a
minimumBy a -> Int
sel (PSQ [(k, a)]
xs) =
  forall k v. [(k, v)] -> PSQ k v
PSQ [forall a b. (a, b) -> b
snd (forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
S.minimumBy (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing forall a b. (a, b) -> a
fst) (forall a b. (a -> b) -> [a] -> [b]
S.map (\ (k, a)
x -> (a -> Int
sel (forall a b. (a, b) -> b
snd (k, a)
x), (k, a)
x)) [(k, a)]
xs))]

-- | Sort the list so that values satisfying the predicate are first.
prefer :: (a -> Bool) -> PSQ k a -> PSQ k a
prefer :: forall a k. (a -> Bool) -> PSQ k a -> PSQ k a
prefer a -> Bool
p = forall a k. (a -> a -> Ordering) -> PSQ k a -> PSQ k a
sortBy forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing a -> Bool
p)

-- | Sort the list so that keys satisfying the predicate are first.
preferByKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
preferByKeys :: forall k a. (k -> Bool) -> PSQ k a -> PSQ k a
preferByKeys k -> Bool
p = forall k a. (k -> k -> Ordering) -> PSQ k a -> PSQ k a
sortByKeys forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing k -> Bool
p)

-- | Will partition the list according to the predicate. If
-- there is any element that satisfies the predicate, then only
-- the elements satisfying the predicate are returned.
-- Otherwise, the rest is returned.
--
filterIfAny :: (a -> Bool) -> PSQ k a -> PSQ k a
filterIfAny :: forall a k. (a -> Bool) -> PSQ k a -> PSQ k a
filterIfAny a -> Bool
p (PSQ [(k, a)]
xs) =
  let
    ([(k, a)]
pro, [(k, a)]
con) = forall a. (a -> Bool) -> [a] -> ([a], [a])
S.partition (a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(k, a)]
xs
  in
    if forall (t :: * -> *) a. Foldable t => t a -> Bool
S.null [(k, a)]
pro then forall k v. [(k, v)] -> PSQ k v
PSQ [(k, a)]
con else forall k v. [(k, v)] -> PSQ k v
PSQ [(k, a)]
pro

-- | Variant of 'filterIfAny' that takes a predicate on the keys
-- rather than on the values.
--
filterIfAnyByKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
filterIfAnyByKeys :: forall k a. (k -> Bool) -> PSQ k a -> PSQ k a
filterIfAnyByKeys k -> Bool
p (PSQ [(k, a)]
xs) =
  let
    ([(k, a)]
pro, [(k, a)]
con) = forall a. (a -> Bool) -> [a] -> ([a], [a])
S.partition (k -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(k, a)]
xs
  in
    if forall (t :: * -> *) a. Foldable t => t a -> Bool
S.null [(k, a)]
pro then forall k v. [(k, v)] -> PSQ k v
PSQ [(k, a)]
con else forall k v. [(k, v)] -> PSQ k v
PSQ [(k, a)]
pro

filterKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
filterKeys :: forall k a. (k -> Bool) -> PSQ k a -> PSQ k a
filterKeys k -> Bool
p (PSQ [(k, a)]
xs) = forall k v. [(k, v)] -> PSQ k v
PSQ (forall a. (a -> Bool) -> [a] -> [a]
S.filter (k -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(k, a)]
xs)

filter :: (a -> Bool) -> PSQ k a -> PSQ k a
filter :: forall a k. (a -> Bool) -> PSQ k a -> PSQ k a
filter a -> Bool
p (PSQ [(k, a)]
xs) = forall k v. [(k, v)] -> PSQ k v
PSQ (forall a. (a -> Bool) -> [a] -> [a]
S.filter (a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(k, a)]
xs)

length :: PSQ k a -> Int
length :: forall k a. PSQ k a -> Int
length (PSQ [(k, a)]
xs) = forall (t :: * -> *) a. Foldable t => t a -> Int
S.length [(k, a)]
xs

null :: PSQ k a -> Bool
null :: forall k a. PSQ k a -> Bool
null (PSQ [(k, a)]
xs) = forall (t :: * -> *) a. Foldable t => t a -> Bool
S.null [(k, a)]
xs

isZeroOrOne :: PSQ k a -> Bool
isZeroOrOne :: forall k a. PSQ k a -> Bool
isZeroOrOne (PSQ [])  = Bool
True
isZeroOrOne (PSQ [(k, a)
_]) = Bool
True
isZeroOrOne PSQ k a
_         = Bool
False

firstOnly :: PSQ k a -> PSQ k a
firstOnly :: forall k a. PSQ k a -> PSQ k a
firstOnly (PSQ [])      = forall k v. [(k, v)] -> PSQ k v
PSQ []
firstOnly (PSQ ((k, a)
x : [(k, a)]
_)) = forall k v. [(k, v)] -> PSQ k v
PSQ [(k, a)
x]

toList :: PSQ k a -> [(k, a)]
toList :: forall k a. PSQ k a -> [(k, a)]
toList (PSQ [(k, a)]
xs) = [(k, a)]
xs

union :: PSQ k a -> PSQ k a -> PSQ k a
union :: forall k a. PSQ k a -> PSQ k a -> PSQ k a
union (PSQ [(k, a)]
xs) (PSQ [(k, a)]
ys) = forall k v. [(k, v)] -> PSQ k v
PSQ ([(k, a)]
xs forall a. [a] -> [a] -> [a]
++ [(k, a)]
ys)