{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
module Distribution.Solver.Modular.WeightedPSQ (
    WeightedPSQ
  , fromList
  , toList
  , keys
  , weights
  , isZeroOrOne
  , filter
  , lookup
  , mapWithKey
  , mapWeightsWithKey
  , traverseWithKey
  , union
  , takeUntil
  ) where

import qualified Data.Foldable as F
import qualified Data.List as L
import Data.Ord (comparing)
import qualified Data.Traversable as T
import Prelude hiding (filter, lookup)

-- | An association list that is sorted by weight.
--
-- Each element has a key ('k'), value ('v'), and weight ('w'). All operations
-- that add elements or modify weights stably sort the elements by weight.
newtype WeightedPSQ w k v = WeightedPSQ [(w, k, v)]
  deriving (WeightedPSQ w k v -> WeightedPSQ w k v -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall w k v.
(Eq w, Eq k, Eq v) =>
WeightedPSQ w k v -> WeightedPSQ w k v -> Bool
/= :: WeightedPSQ w k v -> WeightedPSQ w k v -> Bool
$c/= :: forall w k v.
(Eq w, Eq k, Eq v) =>
WeightedPSQ w k v -> WeightedPSQ w k v -> Bool
== :: WeightedPSQ w k v -> WeightedPSQ w k v -> Bool
$c== :: forall w k v.
(Eq w, Eq k, Eq v) =>
WeightedPSQ w k v -> WeightedPSQ w k v -> Bool
Eq, Int -> WeightedPSQ w k v -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall w k v.
(Show w, Show k, Show v) =>
Int -> WeightedPSQ w k v -> ShowS
forall w k v.
(Show w, Show k, Show v) =>
[WeightedPSQ w k v] -> ShowS
forall w k v.
(Show w, Show k, Show v) =>
WeightedPSQ w k v -> String
showList :: [WeightedPSQ w k v] -> ShowS
$cshowList :: forall w k v.
(Show w, Show k, Show v) =>
[WeightedPSQ w k v] -> ShowS
show :: WeightedPSQ w k v -> String
$cshow :: forall w k v.
(Show w, Show k, Show v) =>
WeightedPSQ w k v -> String
showsPrec :: Int -> WeightedPSQ w k v -> ShowS
$cshowsPrec :: forall w k v.
(Show w, Show k, Show v) =>
Int -> WeightedPSQ w k v -> ShowS
Show, forall a b. a -> WeightedPSQ w k b -> WeightedPSQ w k a
forall a b. (a -> b) -> WeightedPSQ w k a -> WeightedPSQ w k b
forall w k a b. a -> WeightedPSQ w k b -> WeightedPSQ w k a
forall w k a b. (a -> b) -> WeightedPSQ w k a -> WeightedPSQ w 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 -> WeightedPSQ w k b -> WeightedPSQ w k a
$c<$ :: forall w k a b. a -> WeightedPSQ w k b -> WeightedPSQ w k a
fmap :: forall a b. (a -> b) -> WeightedPSQ w k a -> WeightedPSQ w k b
$cfmap :: forall w k a b. (a -> b) -> WeightedPSQ w k a -> WeightedPSQ w k b
Functor, forall a. WeightedPSQ w k a -> Bool
forall m a. Monoid m => (a -> m) -> WeightedPSQ w k a -> m
forall a b. (a -> b -> b) -> b -> WeightedPSQ w k a -> b
forall w k a. Eq a => a -> WeightedPSQ w k a -> Bool
forall w k a. Num a => WeightedPSQ w k a -> a
forall w k a. Ord a => WeightedPSQ w k a -> a
forall w k m. Monoid m => WeightedPSQ w k m -> m
forall w k a. WeightedPSQ w k a -> Bool
forall w k a. WeightedPSQ w k a -> Int
forall w k a. WeightedPSQ w k a -> [a]
forall w k a. (a -> a -> a) -> WeightedPSQ w k a -> a
forall w k m a. Monoid m => (a -> m) -> WeightedPSQ w k a -> m
forall w k b a. (b -> a -> b) -> b -> WeightedPSQ w k a -> b
forall w k a b. (a -> b -> b) -> b -> WeightedPSQ w 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 => WeightedPSQ w k a -> a
$cproduct :: forall w k a. Num a => WeightedPSQ w k a -> a
sum :: forall a. Num a => WeightedPSQ w k a -> a
$csum :: forall w k a. Num a => WeightedPSQ w k a -> a
minimum :: forall a. Ord a => WeightedPSQ w k a -> a
$cminimum :: forall w k a. Ord a => WeightedPSQ w k a -> a
maximum :: forall a. Ord a => WeightedPSQ w k a -> a
$cmaximum :: forall w k a. Ord a => WeightedPSQ w k a -> a
elem :: forall a. Eq a => a -> WeightedPSQ w k a -> Bool
$celem :: forall w k a. Eq a => a -> WeightedPSQ w k a -> Bool
length :: forall a. WeightedPSQ w k a -> Int
$clength :: forall w k a. WeightedPSQ w k a -> Int
null :: forall a. WeightedPSQ w k a -> Bool
$cnull :: forall w k a. WeightedPSQ w k a -> Bool
toList :: forall a. WeightedPSQ w k a -> [a]
$ctoList :: forall w k a. WeightedPSQ w k a -> [a]
foldl1 :: forall a. (a -> a -> a) -> WeightedPSQ w k a -> a
$cfoldl1 :: forall w k a. (a -> a -> a) -> WeightedPSQ w k a -> a
foldr1 :: forall a. (a -> a -> a) -> WeightedPSQ w k a -> a
$cfoldr1 :: forall w k a. (a -> a -> a) -> WeightedPSQ w k a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> WeightedPSQ w k a -> b
$cfoldl' :: forall w k b a. (b -> a -> b) -> b -> WeightedPSQ w k a -> b
foldl :: forall b a. (b -> a -> b) -> b -> WeightedPSQ w k a -> b
$cfoldl :: forall w k b a. (b -> a -> b) -> b -> WeightedPSQ w k a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> WeightedPSQ w k a -> b
$cfoldr' :: forall w k a b. (a -> b -> b) -> b -> WeightedPSQ w k a -> b
foldr :: forall a b. (a -> b -> b) -> b -> WeightedPSQ w k a -> b
$cfoldr :: forall w k a b. (a -> b -> b) -> b -> WeightedPSQ w k a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> WeightedPSQ w k a -> m
$cfoldMap' :: forall w k m a. Monoid m => (a -> m) -> WeightedPSQ w k a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> WeightedPSQ w k a -> m
$cfoldMap :: forall w k m a. Monoid m => (a -> m) -> WeightedPSQ w k a -> m
fold :: forall m. Monoid m => WeightedPSQ w k m -> m
$cfold :: forall w k m. Monoid m => WeightedPSQ w k m -> m
F.Foldable, forall w k. Functor (WeightedPSQ w k)
forall w k. Foldable (WeightedPSQ w k)
forall w k (m :: * -> *) a.
Monad m =>
WeightedPSQ w k (m a) -> m (WeightedPSQ w k a)
forall w k (f :: * -> *) a.
Applicative f =>
WeightedPSQ w k (f a) -> f (WeightedPSQ w k a)
forall w k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> WeightedPSQ w k a -> m (WeightedPSQ w k b)
forall w k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> WeightedPSQ w k a -> f (WeightedPSQ w 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) -> WeightedPSQ w k a -> f (WeightedPSQ w k b)
sequence :: forall (m :: * -> *) a.
Monad m =>
WeightedPSQ w k (m a) -> m (WeightedPSQ w k a)
$csequence :: forall w k (m :: * -> *) a.
Monad m =>
WeightedPSQ w k (m a) -> m (WeightedPSQ w k a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> WeightedPSQ w k a -> m (WeightedPSQ w k b)
$cmapM :: forall w k (m :: * -> *) a b.
Monad m =>
(a -> m b) -> WeightedPSQ w k a -> m (WeightedPSQ w k b)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
WeightedPSQ w k (f a) -> f (WeightedPSQ w k a)
$csequenceA :: forall w k (f :: * -> *) a.
Applicative f =>
WeightedPSQ w k (f a) -> f (WeightedPSQ w k a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> WeightedPSQ w k a -> f (WeightedPSQ w k b)
$ctraverse :: forall w k (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> WeightedPSQ w k a -> f (WeightedPSQ w k b)
T.Traversable)

-- | /O(N)/.
filter :: (v -> Bool) -> WeightedPSQ k w v -> WeightedPSQ k w v
filter :: forall v k w. (v -> Bool) -> WeightedPSQ k w v -> WeightedPSQ k w v
filter v -> Bool
p (WeightedPSQ [(k, w, v)]
xs) = forall w k v. [(w, k, v)] -> WeightedPSQ w k v
WeightedPSQ (forall a. (a -> Bool) -> [a] -> [a]
L.filter (v -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x y z. (x, y, z) -> z
triple_3) [(k, w, v)]
xs)

-- | /O(1)/. Return @True@ if the @WeightedPSQ@ contains zero or one elements.
isZeroOrOne :: WeightedPSQ w k v -> Bool
isZeroOrOne :: forall w k a. WeightedPSQ w k a -> Bool
isZeroOrOne (WeightedPSQ [])  = Bool
True
isZeroOrOne (WeightedPSQ [(w, k, v)
_]) = Bool
True
isZeroOrOne WeightedPSQ w k v
_                 = Bool
False

-- | /O(1)/. Return the elements in order.
toList :: WeightedPSQ w k v -> [(w, k, v)]
toList :: forall w k v. WeightedPSQ w k v -> [(w, k, v)]
toList (WeightedPSQ [(w, k, v)]
xs) = [(w, k, v)]
xs

-- | /O(N log N)/.
fromList :: Ord w => [(w, k, v)] -> WeightedPSQ w k v
fromList :: forall w k v. Ord w => [(w, k, v)] -> WeightedPSQ w k v
fromList = forall w k v. [(w, k, v)] -> WeightedPSQ w k v
WeightedPSQ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing forall x y z. (x, y, z) -> x
triple_1)

-- | /O(N)/. Return the weights in order.
weights :: WeightedPSQ w k v -> [w]
weights :: forall w k v. WeightedPSQ w k v -> [w]
weights (WeightedPSQ [(w, k, v)]
xs) = forall a b. (a -> b) -> [a] -> [b]
L.map forall x y z. (x, y, z) -> x
triple_1 [(w, k, v)]
xs

-- | /O(N)/. Return the keys in order.
keys :: WeightedPSQ w k v -> [k]
keys :: forall w k v. WeightedPSQ w k v -> [k]
keys (WeightedPSQ [(w, k, v)]
xs) = forall a b. (a -> b) -> [a] -> [b]
L.map forall x y z. (x, y, z) -> y
triple_2 [(w, k, v)]
xs

-- | /O(N)/. Return the value associated with the first occurrence of the give
-- key, if it exists.
lookup :: Eq k => k -> WeightedPSQ w k v -> Maybe v
lookup :: forall k w v. Eq k => k -> WeightedPSQ w k v -> Maybe v
lookup k
k (WeightedPSQ [(w, k, v)]
xs) = forall x y z. (x, y, z) -> z
triple_3 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
L.find ((k
k forall a. Eq a => a -> a -> Bool
==) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x y z. (x, y, z) -> y
triple_2) [(w, k, v)]
xs

-- | /O(N log N)/. Update the weights.
mapWeightsWithKey :: Ord w2
                  => (k -> w1 -> w2)
                  -> WeightedPSQ w1 k v
                  -> WeightedPSQ w2 k v
mapWeightsWithKey :: forall w2 k w1 v.
Ord w2 =>
(k -> w1 -> w2) -> WeightedPSQ w1 k v -> WeightedPSQ w2 k v
mapWeightsWithKey k -> w1 -> w2
f (WeightedPSQ [(w1, k, v)]
xs) = forall w k v. Ord w => [(w, k, v)] -> WeightedPSQ w k v
fromList forall a b. (a -> b) -> a -> b
$
                                       forall a b. (a -> b) -> [a] -> [b]
L.map (\ (w1
w, k
k, v
v) -> (k -> w1 -> w2
f k
k w1
w, k
k, v
v)) [(w1, k, v)]
xs

-- | /O(N)/. Update the values.
mapWithKey :: (k -> v1 -> v2) -> WeightedPSQ w k v1 -> WeightedPSQ w k v2
mapWithKey :: forall k v1 v2 w.
(k -> v1 -> v2) -> WeightedPSQ w k v1 -> WeightedPSQ w k v2
mapWithKey k -> v1 -> v2
f (WeightedPSQ [(w, k, v1)]
xs) = forall w k v. [(w, k, v)] -> WeightedPSQ w k v
WeightedPSQ forall a b. (a -> b) -> a -> b
$
                                forall a b. (a -> b) -> [a] -> [b]
L.map (\ (w
w, k
k, v1
v) -> (w
w, k
k, k -> v1 -> v2
f k
k v1
v)) [(w, k, v1)]
xs

-- | /O(N)/. Traverse and update values in some applicative functor.
traverseWithKey
  :: Applicative f
  => (k -> v -> f v')
  -> WeightedPSQ w k v
  -> f (WeightedPSQ w k v')
traverseWithKey :: forall (f :: * -> *) k v v' w.
Applicative f =>
(k -> v -> f v') -> WeightedPSQ w k v -> f (WeightedPSQ w k v')
traverseWithKey k -> v -> f v'
f (WeightedPSQ [(w, k, v)]
q) = forall w k v. [(w, k, v)] -> WeightedPSQ w k v
WeightedPSQ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
  forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\(w
w,k
k,v
v) -> (w
w,k
k,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> v -> f v'
f k
k v
v) [(w, k, v)]
q

-- | /O((N + M) log (N + M))/. Combine two @WeightedPSQ@s, preserving all
-- elements. Elements from the first @WeightedPSQ@ come before elements in the
-- second when they have the same weight.
union :: Ord w => WeightedPSQ w k v -> WeightedPSQ w k v -> WeightedPSQ w k v
union :: forall w k v.
Ord w =>
WeightedPSQ w k v -> WeightedPSQ w k v -> WeightedPSQ w k v
union (WeightedPSQ [(w, k, v)]
xs) (WeightedPSQ [(w, k, v)]
ys) = forall w k v. Ord w => [(w, k, v)] -> WeightedPSQ w k v
fromList ([(w, k, v)]
xs forall a. [a] -> [a] -> [a]
++ [(w, k, v)]
ys)

-- | /O(N)/. Return the prefix of values ending with the first element that
-- satisfies p, or all elements if none satisfy p.
takeUntil :: forall w k v. (v -> Bool) -> WeightedPSQ w k v -> WeightedPSQ w k v
takeUntil :: forall w k v. (v -> Bool) -> WeightedPSQ w k v -> WeightedPSQ w k v
takeUntil v -> Bool
p (WeightedPSQ [(w, k, v)]
xs) = forall w k v. [(w, k, v)] -> WeightedPSQ w k v
WeightedPSQ ([(w, k, v)] -> [(w, k, v)]
go [(w, k, v)]
xs)
  where
    go :: [(w, k, v)] -> [(w, k, v)]
    go :: [(w, k, v)] -> [(w, k, v)]
go []       = []
    go ((w, k, v)
y : [(w, k, v)]
ys) = (w, k, v)
y forall a. a -> [a] -> [a]
: if v -> Bool
p (forall x y z. (x, y, z) -> z
triple_3 (w, k, v)
y) then [] else [(w, k, v)] -> [(w, k, v)]
go [(w, k, v)]
ys

triple_1 :: (x, y, z) -> x
triple_1 :: forall x y z. (x, y, z) -> x
triple_1 (x
x, y
_, z
_) = x
x

triple_2 :: (x, y, z) -> y
triple_2 :: forall x y z. (x, y, z) -> y
triple_2 (x
_, y
y, z
_) = y
y

triple_3 :: (x, y, z) -> z
triple_3 :: forall x y z. (x, y, z) -> z
triple_3 (x
_, y
_, z
z) = z
z