{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}

module Language.REST.Internal.MultiSet
  ( MultiSet
  , delete
  , deleteMany
  , distinctElems
  , empty
  , filter
  , insert
  , member
  , null
  , toList
  , toOccurList
  , singleton
  , fromList
  , toSet
  ) where

import Prelude hiding (null, filter)

import GHC.Generics
import Data.Hashable
import qualified Data.List as L
import qualified Data.HashMap.Strict as M
import qualified Data.HashSet as S

data MultiSet a = MultiSet (M.HashMap a Int) deriving (MultiSet a -> MultiSet a -> Bool
(MultiSet a -> MultiSet a -> Bool)
-> (MultiSet a -> MultiSet a -> Bool) -> Eq (MultiSet a)
forall a. Eq a => MultiSet a -> MultiSet a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: MultiSet a -> MultiSet a -> Bool
$c/= :: forall a. Eq a => MultiSet a -> MultiSet a -> Bool
== :: MultiSet a -> MultiSet a -> Bool
$c== :: forall a. Eq a => MultiSet a -> MultiSet a -> Bool
Eq, (forall x. MultiSet a -> Rep (MultiSet a) x)
-> (forall x. Rep (MultiSet a) x -> MultiSet a)
-> Generic (MultiSet a)
forall x. Rep (MultiSet a) x -> MultiSet a
forall x. MultiSet a -> Rep (MultiSet a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (MultiSet a) x -> MultiSet a
forall a x. MultiSet a -> Rep (MultiSet a) x
$cto :: forall a x. Rep (MultiSet a) x -> MultiSet a
$cfrom :: forall a x. MultiSet a -> Rep (MultiSet a) x
Generic, Int -> MultiSet a -> Int
MultiSet a -> Int
(Int -> MultiSet a -> Int)
-> (MultiSet a -> Int) -> Hashable (MultiSet a)
forall a. Hashable a => Int -> MultiSet a -> Int
forall a. Hashable a => MultiSet a -> Int
forall a. (Int -> a -> Int) -> (a -> Int) -> Hashable a
hash :: MultiSet a -> Int
$chash :: forall a. Hashable a => MultiSet a -> Int
hashWithSalt :: Int -> MultiSet a -> Int
$chashWithSalt :: forall a. Hashable a => Int -> MultiSet a -> Int
Hashable, Eq (MultiSet a)
Eq (MultiSet a)
-> (MultiSet a -> MultiSet a -> Ordering)
-> (MultiSet a -> MultiSet a -> Bool)
-> (MultiSet a -> MultiSet a -> Bool)
-> (MultiSet a -> MultiSet a -> Bool)
-> (MultiSet a -> MultiSet a -> Bool)
-> (MultiSet a -> MultiSet a -> MultiSet a)
-> (MultiSet a -> MultiSet a -> MultiSet a)
-> Ord (MultiSet a)
MultiSet a -> MultiSet a -> Bool
MultiSet a -> MultiSet a -> Ordering
MultiSet a -> MultiSet a -> MultiSet a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. Ord a => Eq (MultiSet a)
forall a. Ord a => MultiSet a -> MultiSet a -> Bool
forall a. Ord a => MultiSet a -> MultiSet a -> Ordering
forall a. Ord a => MultiSet a -> MultiSet a -> MultiSet a
min :: MultiSet a -> MultiSet a -> MultiSet a
$cmin :: forall a. Ord a => MultiSet a -> MultiSet a -> MultiSet a
max :: MultiSet a -> MultiSet a -> MultiSet a
$cmax :: forall a. Ord a => MultiSet a -> MultiSet a -> MultiSet a
>= :: MultiSet a -> MultiSet a -> Bool
$c>= :: forall a. Ord a => MultiSet a -> MultiSet a -> Bool
> :: MultiSet a -> MultiSet a -> Bool
$c> :: forall a. Ord a => MultiSet a -> MultiSet a -> Bool
<= :: MultiSet a -> MultiSet a -> Bool
$c<= :: forall a. Ord a => MultiSet a -> MultiSet a -> Bool
< :: MultiSet a -> MultiSet a -> Bool
$c< :: forall a. Ord a => MultiSet a -> MultiSet a -> Bool
compare :: MultiSet a -> MultiSet a -> Ordering
$ccompare :: forall a. Ord a => MultiSet a -> MultiSet a -> Ordering
$cp1Ord :: forall a. Ord a => Eq (MultiSet a)
Ord)

instance Show a => Show (MultiSet a) where
  show :: MultiSet a -> String
show MultiSet a
ms = String
"{" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
L.intercalate String
", " ((a -> String) -> [a] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map a -> String
forall a. Show a => a -> String
show ([a] -> [String]) -> [a] -> [String]
forall a b. (a -> b) -> a -> b
$ MultiSet a -> [a]
forall a. MultiSet a -> [a]
toList MultiSet a
ms) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"}"

-- | @delete k m@ removes a single instance of @k@ from the multiset @m.
--   If @k is not in the multiset, the original multiset is returned
delete :: (Hashable a, Eq a) => a -> MultiSet a -> MultiSet a
delete :: a -> MultiSet a -> MultiSet a
delete a
k = a -> Int -> MultiSet a -> MultiSet a
forall a.
(Hashable a, Eq a) =>
a -> Int -> MultiSet a -> MultiSet a
deleteMany a
k Int
1

-- | @delete k n m@ removes @n@ instances of @k@ from the multiset @m@.
--   If there are less than @n@ instances of @k@ in the multiset, all
--   instances are removed.
deleteMany :: (Hashable a, Eq a) => a -> Int -> MultiSet a -> MultiSet a
deleteMany :: a -> Int -> MultiSet a -> MultiSet a
deleteMany a
k Int
v (MultiSet HashMap a Int
ms) | Just Int
c <- a -> HashMap a Int -> Maybe Int
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup a
k HashMap a Int
ms
                             , Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
v = HashMap a Int -> MultiSet a
forall a. HashMap a Int -> MultiSet a
MultiSet (HashMap a Int -> MultiSet a) -> HashMap a Int -> MultiSet a
forall a b. (a -> b) -> a -> b
$ a -> Int -> HashMap a Int -> HashMap a Int
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert a
k (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
v) HashMap a Int
ms
deleteMany a
k Int
_ (MultiSet HashMap a Int
ms) | Bool
otherwise = HashMap a Int -> MultiSet a
forall a. HashMap a Int -> MultiSet a
MultiSet (HashMap a Int -> MultiSet a) -> HashMap a Int -> MultiSet a
forall a b. (a -> b) -> a -> b
$ a -> HashMap a Int -> HashMap a Int
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
M.delete a
k HashMap a Int
ms

distinctElems :: MultiSet a -> [a]
distinctElems :: MultiSet a -> [a]
distinctElems (MultiSet HashMap a Int
ms) = HashMap a Int -> [a]
forall k v. HashMap k v -> [k]
M.keys HashMap a Int
ms

empty :: MultiSet a
empty :: MultiSet a
empty = HashMap a Int -> MultiSet a
forall a. HashMap a Int -> MultiSet a
MultiSet HashMap a Int
forall k v. HashMap k v
M.empty

toOccurList :: MultiSet a -> [(a, Int)]
toOccurList :: MultiSet a -> [(a, Int)]
toOccurList (MultiSet HashMap a Int
ms) = HashMap a Int -> [(a, Int)]
forall k v. HashMap k v -> [(k, v)]
M.toList HashMap a Int
ms

filter :: (a -> Bool) -> MultiSet a -> MultiSet a
filter :: (a -> Bool) -> MultiSet a -> MultiSet a
filter a -> Bool
f (MultiSet HashMap a Int
ms) = HashMap a Int -> MultiSet a
forall a. HashMap a Int -> MultiSet a
MultiSet (HashMap a Int -> MultiSet a) -> HashMap a Int -> MultiSet a
forall a b. (a -> b) -> a -> b
$ (a -> Int -> Bool) -> HashMap a Int -> HashMap a Int
forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v
M.filterWithKey a -> Int -> Bool
forall p. a -> p -> Bool
f' HashMap a Int
ms
  where
    f' :: a -> p -> Bool
f' a
k p
_ = a -> Bool
f a
k

null :: MultiSet a -> Bool
null :: MultiSet a -> Bool
null (MultiSet HashMap a Int
ms) = HashMap a Int -> Bool
forall k v. HashMap k v -> Bool
M.null HashMap a Int
ms

-- | @member k m@ returns @true@ iff there is at least one instance of @k@
--   in @m@
member :: (Eq a, Hashable a) => a -> MultiSet a -> Bool
member :: a -> MultiSet a -> Bool
member a
k (MultiSet HashMap a Int
ms) = a -> HashMap a Int -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
M.member a
k HashMap a Int
ms

toList :: MultiSet a -> [a]
toList :: MultiSet a -> [a]
toList MultiSet a
ms = ((a, Int) -> [a]) -> [(a, Int)] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (a, Int) -> [a]
forall a. (a, Int) -> [a]
go (MultiSet a -> [(a, Int)]
forall a. MultiSet a -> [(a, Int)]
toOccurList MultiSet a
ms)
  where
    go :: (a, Int) -> [a]
go (a
k, Int
num) = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
num ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ a -> [a]
forall a. a -> [a]
repeat a
k

insert :: (Eq a, Hashable a) => a -> MultiSet a -> MultiSet a
insert :: a -> MultiSet a -> MultiSet a
insert a
k (MultiSet HashMap a Int
ms) | Just Int
c <- a -> HashMap a Int -> Maybe Int
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup a
k HashMap a Int
ms
                       = HashMap a Int -> MultiSet a
forall a. HashMap a Int -> MultiSet a
MultiSet (HashMap a Int -> MultiSet a) -> HashMap a Int -> MultiSet a
forall a b. (a -> b) -> a -> b
$ a -> Int -> HashMap a Int -> HashMap a Int
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert a
k (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) HashMap a Int
ms
insert a
k (MultiSet HashMap a Int
ms) | Bool
otherwise
                       = HashMap a Int -> MultiSet a
forall a. HashMap a Int -> MultiSet a
MultiSet (HashMap a Int -> MultiSet a) -> HashMap a Int -> MultiSet a
forall a b. (a -> b) -> a -> b
$ a -> Int -> HashMap a Int -> HashMap a Int
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert a
k Int
1 HashMap a Int
ms

singleton :: (Eq a, Hashable a) => a -> MultiSet a
singleton :: a -> MultiSet a
singleton a
k = HashMap a Int -> MultiSet a
forall a. HashMap a Int -> MultiSet a
MultiSet (a -> Int -> HashMap a Int
forall k v. Hashable k => k -> v -> HashMap k v
M.singleton a
k Int
1)

fromList  :: (Eq a, Hashable a) => [a] -> MultiSet a
fromList :: [a] -> MultiSet a
fromList = (MultiSet a -> a -> MultiSet a) -> MultiSet a -> [a] -> MultiSet a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((a -> MultiSet a -> MultiSet a) -> MultiSet a -> a -> MultiSet a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> MultiSet a -> MultiSet a
forall a. (Eq a, Hashable a) => a -> MultiSet a -> MultiSet a
insert) MultiSet a
forall a. MultiSet a
empty

toSet :: MultiSet a -> S.HashSet a
toSet :: MultiSet a -> HashSet a
toSet (MultiSet HashMap a Int
ms) = HashMap a Int -> HashSet a
forall k a. HashMap k a -> HashSet k
M.keysSet HashMap a Int
ms