-- | A simple overlay over Data.Map to manage unordered sets with duplicates.

module Agda.Utils.Bag where

import Prelude hiding (null, map)

import Text.Show.Functions () -- instance only

import qualified Data.List as List
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Semigroup

import Agda.Utils.Functor

import Agda.Utils.Impossible

-- | A set with duplicates.
--   Faithfully stores elements which are equal with regard to (==).
newtype Bag a = Bag
  { forall a. Bag a -> Map a [a]
bag :: Map a [a]
      -- ^ The list contains all occurrences of @a@ (not just the duplicates!).
      --   Hence, the invariant: the list is never empty.
  }
  deriving (Bag a -> Bag a -> Bool
forall a. Eq a => Bag a -> Bag a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Bag a -> Bag a -> Bool
$c/= :: forall a. Eq a => Bag a -> Bag a -> Bool
== :: Bag a -> Bag a -> Bool
$c== :: forall a. Eq a => Bag a -> Bag a -> Bool
Eq, Bag a -> Bag a -> Bool
Bag a -> Bag a -> Ordering
Bag a -> Bag a -> Bag 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 (Bag a)
forall a. Ord a => Bag a -> Bag a -> Bool
forall a. Ord a => Bag a -> Bag a -> Ordering
forall a. Ord a => Bag a -> Bag a -> Bag a
min :: Bag a -> Bag a -> Bag a
$cmin :: forall a. Ord a => Bag a -> Bag a -> Bag a
max :: Bag a -> Bag a -> Bag a
$cmax :: forall a. Ord a => Bag a -> Bag a -> Bag a
>= :: Bag a -> Bag a -> Bool
$c>= :: forall a. Ord a => Bag a -> Bag a -> Bool
> :: Bag a -> Bag a -> Bool
$c> :: forall a. Ord a => Bag a -> Bag a -> Bool
<= :: Bag a -> Bag a -> Bool
$c<= :: forall a. Ord a => Bag a -> Bag a -> Bool
< :: Bag a -> Bag a -> Bool
$c< :: forall a. Ord a => Bag a -> Bag a -> Bool
compare :: Bag a -> Bag a -> Ordering
$ccompare :: forall a. Ord a => Bag a -> Bag a -> Ordering
Ord)
  -- The list contains all occurrences of @a@ (not just the duplicates!).
  -- Hence the invariant: the list is never empty.
  --
  -- This is slightly wasteful, but much easier to implement
  -- in terms of @Map@ as the alternative, which is to store
  -- only the duplicates in the list.
  -- See, e.g., implementation of 'union' which would be impossible
  -- to do in the other representation.  We would need a
  -- 'Map.unionWithKey' that passes us *both* keys.
  -- But Map works under the assumption that Eq for keys is identity,
  -- it does not honor information in keys that goes beyond Ord.

------------------------------------------------------------------------
-- * Query
------------------------------------------------------------------------

-- | Is the bag empty?
null :: Bag a -> Bool
null :: forall a. Bag a -> Bool
null = forall k a. Map k a -> Bool
Map.null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> Map a [a]
bag

-- | Number of elements in the bag.  Duplicates count. O(n).
size :: Bag a -> Int
size :: forall a. Bag a -> Int
size = forall a. Sum a -> a
getSum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a. a -> Sum a
Sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Int
length) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> Map a [a]
bag

-- | @(bag ! a)@ finds all elements equal to @a@.  O(log n).
--   Total function, returns @[]@ if none are.
(!) :: Ord a => Bag a -> a -> [a]
! :: forall a. Ord a => Bag a -> a -> [a]
(!) (Bag Map a [a]
b) a
a = forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault [] a
a Map a [a]
b
  -- Note: not defined infix because of BangPatterns.

-- | O(log n).
member :: Ord a => a -> Bag a -> Bool
member :: forall a. Ord a => a -> Bag a -> Bool
member a
a = Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Ord a => a -> Bag a -> Bool
notMember a
a

-- | O(log n).
notMember :: Ord a => a -> Bag a -> Bool
notMember :: forall a. Ord a => a -> Bag a -> Bool
notMember a
a Bag a
b = forall (t :: * -> *) a. Foldable t => t a -> Bool
List.null (Bag a
b forall a. Ord a => Bag a -> a -> [a]
! a
a)

-- | Return the multiplicity of the given element. O(log n + count _ _).
count :: Ord a => a -> Bag a -> Int
count :: forall a. Ord a => a -> Bag a -> Int
count a
a Bag a
b = forall (t :: * -> *) a. Foldable t => t a -> Int
length (Bag a
b forall a. Ord a => Bag a -> a -> [a]
! a
a)

------------------------------------------------------------------------
-- * Construction
------------------------------------------------------------------------

-- | O(1)
empty :: Bag a
empty :: forall a. Bag a
empty = forall a. Map a [a] -> Bag a
Bag forall a b. (a -> b) -> a -> b
$ forall k a. Map k a
Map.empty

-- | O(1)
singleton :: a -> Bag a
singleton :: forall a. a -> Bag a
singleton a
a = forall a. Map a [a] -> Bag a
Bag forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
Map.singleton a
a [a
a]

union :: Ord a => Bag a -> Bag a -> Bag a
union :: forall a. Ord a => Bag a -> Bag a -> Bag a
union (Bag Map a [a]
b) (Bag Map a [a]
c) = forall a. Map a [a] -> Bag a
Bag forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith forall a. [a] -> [a] -> [a]
(++) Map a [a]
b Map a [a]
c

unions :: Ord a => [Bag a] -> Bag a
unions :: forall a. Ord a => [Bag a] -> Bag a
unions = forall a. Map a [a] -> Bag a
Bag forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith forall a. [a] -> [a] -> [a]
(++)  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
List.map forall a. Bag a -> Map a [a]
bag

-- | @insert a b = union b (singleton a)@
insert :: Ord a => a -> Bag a -> Bag a
insert :: forall a. Ord a => a -> Bag a -> Bag a
insert a
a = forall a. Map a [a] -> Bag a
Bag forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith forall a. [a] -> [a] -> [a]
(++) a
a [a
a] forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> Map a [a]
bag

-- | @fromList = unions . map singleton@
fromList :: Ord a => [a] -> Bag a
fromList :: forall a. Ord a => [a] -> Bag a
fromList = forall a. Map a [a] -> Bag a
Bag forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith forall a. [a] -> [a] -> [a]
(++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
List.map (\ a
a -> (a
a,[a
a]))

------------------------------------------------------------------------
-- * Destruction
------------------------------------------------------------------------

-- | Returns the elements of the bag, grouped by equality (==).
groups :: Bag a -> [[a]]
groups :: forall a. Bag a -> [[a]]
groups = forall k a. Map k a -> [a]
Map.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> Map a [a]
bag

-- | Returns the bag, with duplicates.
toList :: Bag a -> [a]
toList :: forall a. Bag a -> [a]
toList = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> [[a]]
groups

-- | Returns the bag without duplicates.
keys :: Bag a -> [a]
keys :: forall a. Bag a -> [a]
keys = forall k a. Map k a -> [k]
Map.keys forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> Map a [a]
bag
-- Works because of the invariant!
-- keys = catMaybes . map listToMaybe . Map.elems . bag
--   -- Map.keys does not work, as zero copies @(a,[])@
--   -- should count as not present in the bag.

-- | Returns the bag, with duplicates.
elems :: Bag a -> [a]
elems :: forall a. Bag a -> [a]
elems = forall a. Bag a -> [a]
toList

toAscList :: Bag a -> [a]
toAscList :: forall a. Bag a -> [a]
toAscList = forall a. Bag a -> [a]
toList

------------------------------------------------------------------------
-- * Traversal
------------------------------------------------------------------------

map :: Ord b => (a -> b) -> Bag a -> Bag b
map :: forall b a. Ord b => (a -> b) -> Bag a -> Bag b
map a -> b
f = forall a. Map a [a] -> Bag a
Bag forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith forall a. [a] -> [a] -> [a]
(++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
List.map [a] -> (b, [b])
ff forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [a]
Map.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> Map a [a]
bag
  where
    ff :: [a] -> (b, [b])
ff (a
a : [a]
as) = (b
b, b
b forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
List.map a -> b
f [a]
as) where b :: b
b = a -> b
f a
a
    ff []       = forall a. HasCallStack => a
__IMPOSSIBLE__

traverse' :: forall a b m . (Applicative m, Ord b) =>
             (a -> m b) -> Bag a -> m (Bag b)
traverse' :: forall a b (m :: * -> *).
(Applicative m, Ord b) =>
(a -> m b) -> Bag a -> m (Bag b)
traverse' a -> m b
f = (forall a. Map a [a] -> Bag a
Bag forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith forall a. [a] -> [a] -> [a]
(++)) forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> (a -> m b) -> a -> m c
<.> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse [a] -> m (b, [b])
trav forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [a]
Map.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> Map a [a]
bag
  where
    trav :: [a] -> m (b, [b])
    trav :: [a] -> m (b, [b])
trav (a
a : [a]
as) = (\ b
b [b]
bs -> (b
b, b
bforall a. a -> [a] -> [a]
:[b]
bs)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m b
f a
a forall (f :: * -> *) a b. Applicative f => 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 a -> m b
f [a]
as
    trav []       = forall a. HasCallStack => a
__IMPOSSIBLE__

------------------------------------------------------------------------
-- Instances
------------------------------------------------------------------------

instance Show a => Show (Bag a) where
  showsPrec :: Int -> Bag a -> ShowS
showsPrec Int
_ (Bag Map a [a]
b) = (String
"Agda.Utils.Bag.Bag (" forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> ShowS
shows Map a [a]
b forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char
')'forall a. a -> [a] -> [a]
:)

instance Ord a => Semigroup (Bag a) where
  <> :: Bag a -> Bag a -> Bag a
(<>) = forall a. Ord a => Bag a -> Bag a -> Bag a
union

instance Ord a => Monoid (Bag a) where
  mempty :: Bag a
mempty  = forall a. Bag a
empty
  mappend :: Bag a -> Bag a -> Bag a
mappend = forall a. Semigroup a => a -> a -> a
(<>)
  mconcat :: [Bag a] -> Bag a
mconcat = forall a. Ord a => [Bag a] -> Bag a
unions

instance Foldable Bag where
  foldMap :: forall m a. Monoid m => (a -> m) -> Bag a -> m
foldMap a -> m
f = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Bag a -> [a]
toList

-- not a Functor (only works for 'Ord'ered types)
-- not Traversable (only works for 'Ord'ered types)