{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TupleSections #-}

{- | A simple implementation of multisets as hash maps to counts.
 Supports just what is needed for protovoice slices.
-}
module Internal.MultiSet where

import Control.DeepSeq (NFData)
import Control.Monad as M
import Data.Foldable
  ( all
  , foldl'
  )
import qualified Data.HashMap.Strict as HM
import Data.HashSet (HashSet)
import Data.Hashable (Hashable)
import GHC.Generics (Generic)
import Prelude hiding (lookup)

{- | A hash multiset.
 A unordered collection of items that can occur several times.
-}
newtype MultiSet a = MS {forall a. MultiSet a -> HashMap a Int
unMS :: HM.HashMap a Int}
  deriving (MultiSet a -> MultiSet a -> Bool
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, 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
Ord, Int -> MultiSet a -> ShowS
forall a. Show a => Int -> MultiSet a -> ShowS
forall a. Show a => [MultiSet a] -> ShowS
forall a. Show a => MultiSet a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MultiSet a] -> ShowS
$cshowList :: forall a. Show a => [MultiSet a] -> ShowS
show :: MultiSet a -> String
$cshow :: forall a. Show a => MultiSet a -> String
showsPrec :: Int -> MultiSet a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> MultiSet a -> ShowS
Show, 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
forall a. Eq a -> (Int -> a -> Int) -> (a -> Int) -> Hashable a
forall {a}. Hashable a => Eq (MultiSet a)
forall a. Hashable a => Int -> MultiSet a -> Int
forall a. Hashable a => MultiSet a -> Int
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, MultiSet a -> ()
forall a. NFData a => MultiSet a -> ()
forall a. (a -> ()) -> NFData a
rnf :: MultiSet a -> ()
$crnf :: forall a. NFData a => MultiSet a -> ()
NFData)

instance (Eq a, Hashable a) => Semigroup (MultiSet a) where
  (MS HashMap a Int
a) <> :: MultiSet a -> MultiSet a -> MultiSet a
<> (MS HashMap a Int
b) = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
HM.unionWith forall a. Num a => a -> a -> a
(+) HashMap a Int
a HashMap a Int
b

-- instance Foldable MultiSet where
toList :: MultiSet a -> [a]
toList :: forall a. MultiSet a -> [a]
toList (MS HashMap a Int
a) = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall k v1 v2. (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
HM.mapWithKey (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Int -> a -> [a]
replicate) HashMap a Int
a

all :: (a -> Bool) -> MultiSet a -> Bool
all :: forall a. (a -> Bool) -> MultiSet a -> Bool
all a -> Bool
p (MS HashMap a Int
a) = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
Data.Foldable.all a -> Bool
p forall a b. (a -> b) -> a -> b
$ forall k a. HashMap k a -> HashSet k
HM.keysSet HashMap a Int
a

toOccurList :: MultiSet k -> [(k, Int)]
toOccurList :: forall k. MultiSet k -> [(k, Int)]
toOccurList (MS HashMap k Int
m) = forall k v. HashMap k v -> [(k, v)]
HM.toList HashMap k Int
m

distinctElems :: MultiSet k -> [k]
distinctElems :: forall a. MultiSet a -> [a]
distinctElems (MS HashMap k Int
m) = forall k v. HashMap k v -> [k]
HM.keys HashMap k Int
m

union :: (Eq a, Hashable a) => MultiSet a -> MultiSet a -> MultiSet a
union :: forall a.
(Eq a, Hashable a) =>
MultiSet a -> MultiSet a -> MultiSet a
union (MS HashMap a Int
a) (MS HashMap a Int
b) = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
HM.unionWith forall a. Num a => a -> a -> a
(+) HashMap a Int
a HashMap a Int
b

unions :: (Foldable t, Eq a0, Hashable a0) => t (MultiSet a0) -> MultiSet a0
unions :: forall (t :: * -> *) a0.
(Foldable t, Eq a0, Hashable a0) =>
t (MultiSet a0) -> MultiSet a0
unions = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a.
(Eq a, Hashable a) =>
MultiSet a -> MultiSet a -> MultiSet a
union forall a. MultiSet a
empty

maxUnion :: (Eq a, Hashable a) => MultiSet a -> MultiSet a -> MultiSet a
maxUnion :: forall a.
(Eq a, Hashable a) =>
MultiSet a -> MultiSet a -> MultiSet a
maxUnion (MS HashMap a Int
a) (MS HashMap a Int
b) = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
HM.unionWith forall a. Ord a => a -> a -> a
max HashMap a Int
a HashMap a Int
b

deleteN :: Int -> Int -> Maybe Int
deleteN :: Int -> Int -> Maybe Int
deleteN Int
n Int
m
  | Int
m forall a. Ord a => a -> a -> Bool
<= Int
n = forall a. Maybe a
Nothing
  | Bool
otherwise = forall a. a -> Maybe a
Just (Int
m forall a. Num a => a -> a -> a
- Int
n)

(\\) :: (Eq a, Hashable a) => MultiSet a -> MultiSet a -> MultiSet a
(MS HashMap a Int
a) \\ :: forall a.
(Eq a, Hashable a) =>
MultiSet a -> MultiSet a -> MultiSet a
\\ (MS HashMap a Int
b) = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall k v w.
(Eq k, Hashable k) =>
(v -> w -> Maybe v) -> HashMap k v -> HashMap k w -> HashMap k v
HM.differenceWith (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> Maybe Int
deleteN) HashMap a Int
a HashMap a Int
b

toSet :: MultiSet k -> HashSet k
toSet :: forall k. MultiSet k -> HashSet k
toSet (MS HashMap k Int
a) = forall k a. HashMap k a -> HashSet k
HM.keysSet HashMap k Int
a

size :: MultiSet a -> Int
size :: forall a. MultiSet a -> Int
size (MS HashMap a Int
a) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(+) Int
0 HashMap a Int
a

map :: (Eq b, Hashable b) => (a -> b) -> MultiSet a -> MultiSet b
map :: forall b a.
(Eq b, Hashable b) =>
(a -> b) -> MultiSet a -> MultiSet b
map a -> b
f (MS HashMap a Int
a) =
  forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
HM.foldlWithKey' (\HashMap b Int
m a
k Int
o -> forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HM.insertWith forall a. Num a => a -> a -> a
(+) (a -> b
f a
k) Int
o HashMap b Int
m) forall k v. HashMap k v
HM.empty HashMap a Int
a

traverse
  :: (Eq b, Hashable b, Applicative f)
  => (a -> f b)
  -> MultiSet a
  -> f (MultiSet b)
traverse :: forall b (f :: * -> *) a.
(Eq b, Hashable b, Applicative f) =>
(a -> f b) -> MultiSet a -> f (MultiSet b)
traverse a -> f b
f (MS HashMap a Int
m) =
  forall a. HashMap a Int -> MultiSet a
MS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> [(k, v)] -> HashMap k v
HM.fromListWith forall a. Num a => a -> a -> a
(+)
    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)
Prelude.traverse
      (\(a
k, Int
v) -> (,Int
v) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
k)
      (forall k v. HashMap k v -> [(k, v)]
HM.toList HashMap a Int
m)

filter :: (a -> Bool) -> MultiSet a -> MultiSet a
filter :: forall a. (a -> Bool) -> MultiSet a -> MultiSet a
filter a -> Bool
f (MS HashMap a Int
a) = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v
HM.filterWithKey (\a
k Int
_ -> a -> Bool
f a
k) HashMap a Int
a

fromList :: (Foldable t, Eq a, Hashable a) => t a -> MultiSet a
fromList :: forall (t :: * -> *) a.
(Foldable t, Eq a, Hashable a) =>
t a -> MultiSet a
fromList t a
xs = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\HashMap a Int
m a
x -> forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HM.insertWith forall a. Num a => a -> a -> a
(+) a
x Int
1 HashMap a Int
m) forall k v. HashMap k v
HM.empty t a
xs

fromSet :: (Foldable t, Eq a, Hashable a) => t a -> MultiSet a
fromSet :: forall (t :: * -> *) a.
(Foldable t, Eq a, Hashable a) =>
t a -> MultiSet a
fromSet t a
xs = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\HashMap a Int
m a
x -> forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert a
x Int
1 HashMap a Int
m) forall k v. HashMap k v
HM.empty t a
xs

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

delete :: (Eq a, Hashable a) => a -> MultiSet a -> MultiSet a
delete :: forall a. (Eq a, Hashable a) => a -> MultiSet a -> MultiSet a
delete a
x = forall a. HashMap a Int -> MultiSet a
MS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a.
(Eq k, Hashable k) =>
(a -> Maybe a) -> k -> HashMap k a -> HashMap k a
HM.update (Int -> Int -> Maybe Int
deleteN Int
1) a
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. MultiSet a -> HashMap a Int
unMS

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

foldM :: Monad m => (b -> a -> m b) -> b -> MultiSet a -> m b
foldM :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> MultiSet a -> m b
foldM b -> a -> m b
f b
b (MS HashMap a Int
as) = forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
HM.foldlWithKey' m b -> a -> Int -> m b
doFold (forall (f :: * -> *) a. Applicative f => a -> f a
pure b
b) HashMap a Int
as
 where
  doFold :: m b -> a -> Int -> m b
doFold m b
mb a
a Int
i = do
    b
b' <- m b
mb
    forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
M.foldM b -> a -> m b
f b
b' (forall a. Int -> a -> [a]
replicate Int
i a
a)

insertMany :: (Eq a, Hashable a) => a -> Int -> MultiSet a -> MultiSet a
insertMany :: forall a.
(Eq a, Hashable a) =>
a -> Int -> MultiSet a -> MultiSet a
insertMany a
a Int
n (MS HashMap a Int
as)
  | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = forall a. HashMap a Int -> MultiSet a
MS HashMap a Int
as
  | Bool
otherwise = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HM.insertWith forall a. Num a => a -> a -> a
(+) a
a Int
n HashMap a Int
as

insert :: (Eq a, Hashable a) => a -> MultiSet a -> MultiSet a
insert :: forall a. (Eq a, Hashable a) => a -> MultiSet a -> MultiSet a
insert a
a = forall a.
(Eq a, Hashable a) =>
a -> Int -> MultiSet a -> MultiSet a
insertMany a
a Int
1

singleton :: Hashable a => a -> MultiSet a
singleton :: forall a. Hashable a => a -> MultiSet a
singleton a
a = forall a. HashMap a Int -> MultiSet a
MS forall a b. (a -> b) -> a -> b
$ forall k v. Hashable k => k -> v -> HashMap k v
HM.singleton a
a Int
1

member :: (Eq k, Hashable k) => k -> MultiSet k -> Bool
member :: forall k. (Eq k, Hashable k) => k -> MultiSet k -> Bool
member k
a MultiSet k
as = forall k. (Eq k, Hashable k) => k -> MultiSet k -> Int
lookup k
a MultiSet k
as forall a. Ord a => a -> a -> Bool
> Int
0

lookup :: (Eq k, Hashable k) => k -> MultiSet k -> Int
lookup :: forall k. (Eq k, Hashable k) => k -> MultiSet k -> Int
lookup k
a (MS HashMap k Int
as) = forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
HM.lookupDefault Int
0 k
a HashMap k Int
as

(!) :: (Eq k, Hashable k) => MultiSet k -> k -> Int
MultiSet k
as ! :: forall k. (Eq k, Hashable k) => MultiSet k -> k -> Int
! k
a = forall k. (Eq k, Hashable k) => k -> MultiSet k -> Int
lookup k
a MultiSet k
as