content: {-# 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
      , splits
      , 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 (Eq, Show, Functor, F.Foldable, Traversable) -- Qualified Foldable to avoid issues with FTP
  
  keys :: PSQ k v -> [k]
  keys (PSQ xs) = fmap fst xs
  
  lookup :: Eq k => k -> PSQ k v -> Maybe v
  lookup k (PSQ xs) = S.lookup k xs
  
  map :: (v1 -> v2) -> PSQ k v1 -> PSQ k v2
  map f (PSQ xs) = PSQ (fmap (second f) xs)
  
  mapKeys :: (k1 -> k2) -> PSQ k1 v -> PSQ k2 v
  mapKeys f (PSQ xs) = PSQ (fmap (first f) xs)
  
  mapWithKey :: (k -> a -> b) -> PSQ k a -> PSQ k b
  mapWithKey f (PSQ xs) = PSQ (fmap (\ (k, v) -> (k, f k v)) xs)
  
  fromList :: [(k, a)] -> PSQ k a
  fromList = PSQ
  
  cons :: k -> a -> PSQ k a -> PSQ k a
  cons k x (PSQ xs) = PSQ ((k, x) : xs)
  
  snoc :: PSQ k a -> k -> a -> PSQ k a
  snoc (PSQ xs) k x = PSQ (xs ++ [(k, x)])
  
  casePSQ :: PSQ k a -> r -> (k -> a -> PSQ k a -> r) -> r
  casePSQ (PSQ xs) n c =
    case xs of
      []          -> n
      (k, v) : ys -> c k v (PSQ ys)
  
  splits :: PSQ k a -> PSQ k (a, PSQ k a)
  splits = go id
    where
      go f xs = casePSQ xs
          (PSQ [])
          (\ k v ys -> cons k (v, f ys) (go (f . cons k v) ys))
  
  sortBy :: (a -> a -> Ordering) -> PSQ k a -> PSQ k a
  sortBy cmp (PSQ xs) = PSQ (S.sortBy (cmp `on` snd) xs)
  
  sortByKeys :: (k -> k -> Ordering) -> PSQ k a -> PSQ k a
  sortByKeys cmp (PSQ xs) = PSQ (S.sortBy (cmp `on` fst) xs)
  
  maximumBy :: (k -> Int) -> PSQ k a -> (k, a)
  maximumBy sel (PSQ xs) =
    S.minimumBy (flip (comparing (sel . fst))) xs
  
  minimumBy :: (a -> Int) -> PSQ k a -> PSQ k a
  minimumBy sel (PSQ xs) =
    PSQ [snd (S.minimumBy (comparing fst) (S.map (\ x -> (sel (snd x), x)) xs))]
  
  -- | Sort the list so that values satisfying the predicate are first.
  prefer :: (a -> Bool) -> PSQ k a -> PSQ k a
  prefer p = sortBy $ flip (comparing p)
  
  -- | Sort the list so that keys satisfying the predicate are first.
  preferByKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
  preferByKeys p = sortByKeys $ flip (comparing p)
  
  -- | Will partition the list according to the predicate. If
  -- there is any element that satisfies the precidate, then only
  -- the elements satisfying the predicate are returned.
  -- Otherwise, the rest is returned.
  --
  filterIfAny :: (a -> Bool) -> PSQ k a -> PSQ k a
  filterIfAny p (PSQ xs) =
    let
      (pro, con) = S.partition (p . snd) xs
    in
      if S.null pro then PSQ con else PSQ 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 p (PSQ xs) =
    let
      (pro, con) = S.partition (p . fst) xs
    in
      if S.null pro then PSQ con else PSQ pro
  
  filterKeys :: (k -> Bool) -> PSQ k a -> PSQ k a
  filterKeys p (PSQ xs) = PSQ (S.filter (p . fst) xs)
  
  filter :: (a -> Bool) -> PSQ k a -> PSQ k a
  filter p (PSQ xs) = PSQ (S.filter (p . snd) xs)
  
  length :: PSQ k a -> Int
  length (PSQ xs) = S.length xs
  
  null :: PSQ k a -> Bool
  null (PSQ xs) = S.null xs
  
  isZeroOrOne :: PSQ k a -> Bool
  isZeroOrOne (PSQ [])  = True
  isZeroOrOne (PSQ [_]) = True
  isZeroOrOne _         = False
  
  firstOnly :: PSQ k a -> PSQ k a
  firstOnly (PSQ [])      = PSQ []
  firstOnly (PSQ (x : _)) = PSQ [x]
  
  toList :: PSQ k a -> [(k, a)]
  toList (PSQ xs) = xs
  
  union :: PSQ k a -> PSQ k a -> PSQ k a
  union (PSQ xs) (PSQ ys) = PSQ (xs ++ ys)
  