-- | A left-unique relation.
module Unison.Util.BiMultimap
  ( BiMultimap,
    Unison.Util.BiMultimap.empty,

    -- ** Basic queries
    isEmpty,

    -- ** Lookup
    memberDom,
    lookupDom,
    lookupRan,
    unsafeLookupRan,
    lookupPreimage,

    -- ** Mapping / traversing
    unsafeTraverseDom,

    -- ** Filtering
    filter,
    filterDom,
    filterDomain,
    restrictDom,
    restrictRan,
    withoutDom,
    withoutRan,

    -- ** Maps
    domain,
    range,
    unsafeFromDomain,
    fromRange,

    -- ** Sets
    dom,
    ran,

    -- ** Relations
    toRelation,

    -- ** Insert
    insert,
    unsafeInsert,

    -- ** Union
    unsafeUnion,
  )
where

import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import Data.Set.NonEmpty (NESet)
import Data.Set.NonEmpty qualified as Set.NonEmpty
import Unison.Prelude
import Unison.Util.Map qualified as Map
import Unison.Util.Relation (Relation)
import Unison.Util.Relation qualified as Relation
import Prelude hiding (filter)

-- | A left-unique relation.
--
-- "Left-unique" means that for all @(x, y)@ in the relation, @y@ is related only to @x@.
data BiMultimap a b = BiMultimap
  { forall a b. BiMultimap a b -> Map a (NESet b)
toMultimap :: !(Map a (NESet b)),
    forall a b. BiMultimap a b -> Map b a
toMapR :: !(Map b a)
  }
  deriving (BiMultimap a b -> BiMultimap a b -> Bool
(BiMultimap a b -> BiMultimap a b -> Bool)
-> (BiMultimap a b -> BiMultimap a b -> Bool)
-> Eq (BiMultimap a b)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall a b.
(Eq a, Eq b) =>
BiMultimap a b -> BiMultimap a b -> Bool
$c== :: forall a b.
(Eq a, Eq b) =>
BiMultimap a b -> BiMultimap a b -> Bool
== :: BiMultimap a b -> BiMultimap a b -> Bool
$c/= :: forall a b.
(Eq a, Eq b) =>
BiMultimap a b -> BiMultimap a b -> Bool
/= :: BiMultimap a b -> BiMultimap a b -> Bool
Eq, Eq (BiMultimap a b)
Eq (BiMultimap a b) =>
(BiMultimap a b -> BiMultimap a b -> Ordering)
-> (BiMultimap a b -> BiMultimap a b -> Bool)
-> (BiMultimap a b -> BiMultimap a b -> Bool)
-> (BiMultimap a b -> BiMultimap a b -> Bool)
-> (BiMultimap a b -> BiMultimap a b -> Bool)
-> (BiMultimap a b -> BiMultimap a b -> BiMultimap a b)
-> (BiMultimap a b -> BiMultimap a b -> BiMultimap a b)
-> Ord (BiMultimap a b)
BiMultimap a b -> BiMultimap a b -> Bool
BiMultimap a b -> BiMultimap a b -> Ordering
BiMultimap a b -> BiMultimap a b -> BiMultimap a b
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 b. (Ord a, Ord b) => Eq (BiMultimap a b)
forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> Bool
forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> Ordering
forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> BiMultimap a b
$ccompare :: forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> Ordering
compare :: BiMultimap a b -> BiMultimap a b -> Ordering
$c< :: forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> Bool
< :: BiMultimap a b -> BiMultimap a b -> Bool
$c<= :: forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> Bool
<= :: BiMultimap a b -> BiMultimap a b -> Bool
$c> :: forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> Bool
> :: BiMultimap a b -> BiMultimap a b -> Bool
$c>= :: forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> Bool
>= :: BiMultimap a b -> BiMultimap a b -> Bool
$cmax :: forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> BiMultimap a b
max :: BiMultimap a b -> BiMultimap a b -> BiMultimap a b
$cmin :: forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> BiMultimap a b
min :: BiMultimap a b -> BiMultimap a b -> BiMultimap a b
Ord, Int -> BiMultimap a b -> ShowS
[BiMultimap a b] -> ShowS
BiMultimap a b -> String
(Int -> BiMultimap a b -> ShowS)
-> (BiMultimap a b -> String)
-> ([BiMultimap a b] -> ShowS)
-> Show (BiMultimap a b)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. (Show a, Show b) => Int -> BiMultimap a b -> ShowS
forall a b. (Show a, Show b) => [BiMultimap a b] -> ShowS
forall a b. (Show a, Show b) => BiMultimap a b -> String
$cshowsPrec :: forall a b. (Show a, Show b) => Int -> BiMultimap a b -> ShowS
showsPrec :: Int -> BiMultimap a b -> ShowS
$cshow :: forall a b. (Show a, Show b) => BiMultimap a b -> String
show :: BiMultimap a b -> String
$cshowList :: forall a b. (Show a, Show b) => [BiMultimap a b] -> ShowS
showList :: [BiMultimap a b] -> ShowS
Show)

-- | An empty left-unique relation.
empty :: (Ord a, Ord b) => BiMultimap a b
empty :: forall a b. (Ord a, Ord b) => BiMultimap a b
empty = Map a (NESet b) -> Map b a -> BiMultimap a b
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap Map a (NESet b)
forall a. Monoid a => a
mempty Map b a
forall a. Monoid a => a
mempty

-- | Is a left-unique relation empty?
isEmpty :: BiMultimap a b -> Bool
isEmpty :: forall a b. BiMultimap a b -> Bool
isEmpty =
  Map a (NESet b) -> Bool
forall k a. Map k a -> Bool
Map.null (Map a (NESet b) -> Bool)
-> (BiMultimap a b -> Map a (NESet b)) -> BiMultimap a b -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
domain

memberDom :: (Ord a) => a -> BiMultimap a b -> Bool
memberDom :: forall a b. Ord a => a -> BiMultimap a b -> Bool
memberDom a
x =
  a -> Map a (NESet b) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member a
x (Map a (NESet b) -> Bool)
-> (BiMultimap a b -> Map a (NESet b)) -> BiMultimap a b -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
domain

-- | Look up the set of @b@ related to an @a@.
--
-- /O(log a)/.
lookupDom :: (Ord a) => a -> BiMultimap a b -> Set b
lookupDom :: forall a b. Ord a => a -> BiMultimap a b -> Set b
lookupDom a
a =
  a -> Map a (NESet b) -> Set b
forall a b. Ord a => a -> Map a (NESet b) -> Set b
lookupDom_ a
a (Map a (NESet b) -> Set b)
-> (BiMultimap a b -> Map a (NESet b)) -> BiMultimap a b -> Set b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
domain

lookupDom_ :: (Ord a) => a -> Map a (NESet b) -> Set b
lookupDom_ :: forall a b. Ord a => a -> Map a (NESet b) -> Set b
lookupDom_ a
x Map a (NESet b)
xs =
  Set b -> (NESet b -> Set b) -> Maybe (NESet b) -> Set b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set b
forall a. Set a
Set.empty NESet b -> Set b
forall a. NESet a -> Set a
Set.NonEmpty.toSet (a -> Map a (NESet b) -> Maybe (NESet b)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
x Map a (NESet b)
xs)

-- | Look up the @a@ related to a @b@.
--
-- /O(log b)/.
lookupRan :: (Ord b) => b -> BiMultimap a b -> Maybe a
lookupRan :: forall b a. Ord b => b -> BiMultimap a b -> Maybe a
lookupRan b
b (BiMultimap Map a (NESet b)
_ Map b a
r) =
  b -> Map b a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup b
b Map b a
r

-- | Look up the @a@ related to a @b@.
--
-- /O(log b)/.
unsafeLookupRan :: (Ord b) => b -> BiMultimap a b -> a
unsafeLookupRan :: forall b a. Ord b => b -> BiMultimap a b -> a
unsafeLookupRan b
b (BiMultimap Map a (NESet b)
_ Map b a
r) =
  Map b a
r Map b a -> b -> a
forall k a. Ord k => Map k a -> k -> a
Map.! b
b

-- | Look up the preimage of a @b@, that is, the set of @b@ that are related to the same @a@ as the input @b@.
--
-- /O(log a + log b)
lookupPreimage :: (Ord a, Ord b) => b -> BiMultimap a b -> Set b
lookupPreimage :: forall a b. (Ord a, Ord b) => b -> BiMultimap a b -> Set b
lookupPreimage b
y (BiMultimap Map a (NESet b)
domain Map b a
range) =
  Set b -> (a -> Set b) -> Maybe a -> Set b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set b
forall a. Set a
Set.empty (\a
x -> a -> Map a (NESet b) -> Set b
forall a b. Ord a => a -> Map a (NESet b) -> Set b
lookupDom_ a
x Map a (NESet b)
domain) (b -> Map b a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup b
y Map b a
range)

-- | Traverse over the domain a left-unique relation.
--
-- The caller is responsible for maintaining left-uniqueness.
unsafeTraverseDom :: forall a b m x. (Monad m, Ord b, Ord x) => (a -> m b) -> BiMultimap a x -> m (BiMultimap b x)
unsafeTraverseDom :: forall a b (m :: * -> *) x.
(Monad m, Ord b, Ord x) =>
(a -> m b) -> BiMultimap a x -> m (BiMultimap b x)
unsafeTraverseDom a -> m b
f BiMultimap a x
m =
  ((a, NESet x)
 -> (BiMultimap b x -> m (BiMultimap b x))
 -> BiMultimap b x
 -> m (BiMultimap b x))
-> (BiMultimap b x -> m (BiMultimap b x))
-> [(a, NESet x)]
-> BiMultimap b x
-> m (BiMultimap b x)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a, NESet x)
-> (BiMultimap b x -> m (BiMultimap b x))
-> BiMultimap b x
-> m (BiMultimap b x)
g BiMultimap b x -> m (BiMultimap b x)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map a (NESet x) -> [(a, NESet x)]
forall k a. Map k a -> [(k, a)]
Map.toList (BiMultimap a x -> Map a (NESet x)
forall a b. BiMultimap a b -> Map a (NESet b)
domain BiMultimap a x
m)) BiMultimap b x
forall a b. (Ord a, Ord b) => BiMultimap a b
Unison.Util.BiMultimap.empty
  where
    g :: (a, NESet x) -> (BiMultimap b x -> m (BiMultimap b x)) -> (BiMultimap b x -> m (BiMultimap b x))
    g :: (a, NESet x)
-> (BiMultimap b x -> m (BiMultimap b x))
-> BiMultimap b x
-> m (BiMultimap b x)
g (a
a, NESet x
xs) BiMultimap b x -> m (BiMultimap b x)
acc (BiMultimap Map b (NESet x)
domain0 Map x b
range0) = do
      !b
b <- a -> m b
f a
a
      BiMultimap b x -> m (BiMultimap b x)
acc (BiMultimap b x -> m (BiMultimap b x))
-> BiMultimap b x -> m (BiMultimap b x)
forall a b. (a -> b) -> a -> b
$! Map b (NESet x) -> Map x b -> BiMultimap b x
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap (b -> NESet x -> Map b (NESet x) -> Map b (NESet x)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert b
b NESet x
xs Map b (NESet x)
domain0) (b -> NESet x -> Map x b -> Map x b
forall b a. Ord b => a -> NESet b -> Map b a -> Map b a
deriveRangeFromDomain b
b NESet x
xs Map x b
range0)

-- | Filter a left-unique relation, keeping only members @(a, b)@ that satisfy a predicate.
filter :: (Ord a, Ord b) => (a -> b -> Bool) -> BiMultimap a b -> BiMultimap a b
filter :: forall a b.
(Ord a, Ord b) =>
(a -> b -> Bool) -> BiMultimap a b -> BiMultimap a b
filter a -> b -> Bool
p (BiMultimap Map a (NESet b)
domain Map b a
range) =
  Map a (NESet b) -> Map b a -> BiMultimap a b
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap
    ( (a -> NESet b -> Maybe (NESet b))
-> Map a (NESet b) -> Map a (NESet b)
forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybeWithKey
        ( \a
x NESet b
ys ->
            NESet b
ys
              NESet b -> (NESet b -> Set b) -> Set b
forall a b. a -> (a -> b) -> b
& (b -> Bool) -> NESet b -> Set b
forall a. (a -> Bool) -> NESet a -> Set a
Set.NonEmpty.filter (a -> b -> Bool
p a
x)
              Set b -> (Set b -> Maybe (NESet b)) -> Maybe (NESet b)
forall a b. a -> (a -> b) -> b
& Set b -> Maybe (NESet b)
forall a. Set a -> Maybe (NESet a)
Set.NonEmpty.nonEmptySet
        )
        Map a (NESet b)
domain
    )
    ((b -> a -> Bool) -> Map b a -> Map b a
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey ((a -> b -> Bool) -> b -> a -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> Bool
p) Map b a
range)

-- | Filter a left-unique relation, keeping only members @(a, b)@ whose @a@ satisfies a predicate.
filterDom :: (Ord a, Ord b) => (a -> Bool) -> BiMultimap a b -> BiMultimap a b
filterDom :: forall a b.
(Ord a, Ord b) =>
(a -> Bool) -> BiMultimap a b -> BiMultimap a b
filterDom a -> Bool
f BiMultimap a b
m =
  Map a (NESet b) -> BiMultimap a b
forall b a. Ord b => Map a (NESet b) -> BiMultimap a b
unsafeFromDomain ((a -> NESet b -> Bool) -> Map a (NESet b) -> Map a (NESet b)
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (\a
x NESet b
_ -> a -> Bool
f a
x) (BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
domain BiMultimap a b
m))

-- | Filter a left-unique relation, keeping only members @(a, b)@ whose @a@ and set of @b@ satisfies a predicate.
filterDomain :: (Ord a, Ord b) => (a -> NESet b -> Bool) -> BiMultimap a b -> BiMultimap a b
filterDomain :: forall a b.
(Ord a, Ord b) =>
(a -> NESet b -> Bool) -> BiMultimap a b -> BiMultimap a b
filterDomain a -> NESet b -> Bool
f BiMultimap a b
m =
  Map a (NESet b) -> BiMultimap a b
forall b a. Ord b => Map a (NESet b) -> BiMultimap a b
unsafeFromDomain ((a -> NESet b -> Bool) -> Map a (NESet b) -> Map a (NESet b)
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey a -> NESet b -> Bool
f (BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
domain BiMultimap a b
m))

-- | Restrict a left-unique relation to only those @(a, b)@ members whose @a@ is in the given set.
restrictDom :: (Ord a, Ord b) => Set a -> BiMultimap a b -> BiMultimap a b
restrictDom :: forall a b.
(Ord a, Ord b) =>
Set a -> BiMultimap a b -> BiMultimap a b
restrictDom Set a
xs BiMultimap a b
m =
  Map a (NESet b) -> BiMultimap a b
forall b a. Ord b => Map a (NESet b) -> BiMultimap a b
unsafeFromDomain (Map a (NESet b) -> Set a -> Map a (NESet b)
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.restrictKeys (BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
domain BiMultimap a b
m) Set a
xs)

-- | Restrict a left-unique relation to only those @(a, b)@ members whose @b@ is in the given set.
restrictRan :: (Ord a, Ord b) => Set b -> BiMultimap a b -> BiMultimap a b
restrictRan :: forall a b.
(Ord a, Ord b) =>
Set b -> BiMultimap a b -> BiMultimap a b
restrictRan Set b
ys BiMultimap a b
m =
  Map b a -> BiMultimap a b
forall a b. (Ord a, Ord b) => Map b a -> BiMultimap a b
fromRange (Map b a -> Set b -> Map b a
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.restrictKeys (BiMultimap a b -> Map b a
forall a b. BiMultimap a b -> Map b a
range BiMultimap a b
m) Set b
ys)

-- | Restrict a left-unique relation to only those @(a, b)@ members whose @a@ is not in the given set.
withoutDom :: (Ord a, Ord b) => Set a -> BiMultimap a b -> BiMultimap a b
withoutDom :: forall a b.
(Ord a, Ord b) =>
Set a -> BiMultimap a b -> BiMultimap a b
withoutDom Set a
xs BiMultimap a b
m =
  Map a (NESet b) -> BiMultimap a b
forall b a. Ord b => Map a (NESet b) -> BiMultimap a b
unsafeFromDomain (Map a (NESet b) -> Set a -> Map a (NESet b)
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys (BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
domain BiMultimap a b
m) Set a
xs)

-- | Restrict a left-unique relation to only those @(a, b)@ members whose @b@ is not in the given set.
withoutRan :: (Ord a, Ord b) => Set b -> BiMultimap a b -> BiMultimap a b
withoutRan :: forall a b.
(Ord a, Ord b) =>
Set b -> BiMultimap a b -> BiMultimap a b
withoutRan Set b
ys BiMultimap a b
m =
  Map b a -> BiMultimap a b
forall a b. (Ord a, Ord b) => Map b a -> BiMultimap a b
fromRange (Map b a -> Set b -> Map b a
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys (BiMultimap a b -> Map b a
forall a b. BiMultimap a b -> Map b a
range BiMultimap a b
m) Set b
ys)

domain :: BiMultimap a b -> Map a (NESet b)
domain :: forall a b. BiMultimap a b -> Map a (NESet b)
domain = BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
toMultimap

-- | /O(1)/.
range :: BiMultimap a b -> Map b a
range :: forall a b. BiMultimap a b -> Map b a
range = BiMultimap a b -> Map b a
forall a b. BiMultimap a b -> Map b a
toMapR

-- | Construct a left-unique relation from a mapping from its left-elements to set-of-right-elements. The caller is
-- responsible for ensuring that no right-element is mapped to by two different left-elements.
unsafeFromDomain :: (Ord b) => Map a (NESet b) -> BiMultimap a b
unsafeFromDomain :: forall b a. Ord b => Map a (NESet b) -> BiMultimap a b
unsafeFromDomain Map a (NESet b)
domain =
  Map a (NESet b) -> Map b a -> BiMultimap a b
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap Map a (NESet b)
domain (Map a (NESet b) -> Map b a
forall a b. Ord b => Map a (NESet b) -> Map b a
invertDomain Map a (NESet b)
domain)

invertDomain :: forall a b. (Ord b) => Map a (NESet b) -> Map b a
invertDomain :: forall a b. Ord b => Map a (NESet b) -> Map b a
invertDomain =
  (Map b a -> a -> NESet b -> Map b a)
-> Map b a -> Map a (NESet b) -> Map b a
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
Map.foldlWithKey' Map b a -> a -> NESet b -> Map b a
f Map b a
forall k a. Map k a
Map.empty
  where
    f :: Map b a -> a -> NESet b -> Map b a
    f :: Map b a -> a -> NESet b -> Map b a
f Map b a
acc a
x NESet b
ys =
      (Map b a -> b -> Map b a) -> Map b a -> NESet b -> Map b a
forall a b. (a -> b -> a) -> a -> NESet b -> a
Set.NonEmpty.foldl' (a -> Map b a -> b -> Map b a
g a
x) Map b a
acc NESet b
ys

    g :: a -> Map b a -> b -> Map b a
    g :: a -> Map b a -> b -> Map b a
g a
x Map b a
acc b
y =
      b -> a -> Map b a -> Map b a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert b
y a
x Map b a
acc

-- | Construct a left-unique relation from a mapping from its right-elements to its left-elements.
fromRange :: (Ord a, Ord b) => Map b a -> BiMultimap a b
fromRange :: forall a b. (Ord a, Ord b) => Map b a -> BiMultimap a b
fromRange Map b a
m =
  Map a (NESet b) -> Map b a -> BiMultimap a b
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap ((Map a (NESet b) -> b -> a -> Map a (NESet b))
-> Map a (NESet b) -> Map b a -> Map a (NESet b)
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
Map.foldlWithKey' Map a (NESet b) -> b -> a -> Map a (NESet b)
forall {k} {a}.
(Ord k, Ord a) =>
Map k (NESet a) -> a -> k -> Map k (NESet a)
f Map a (NESet b)
forall k a. Map k a
Map.empty Map b a
m) Map b a
m
  where
    f :: Map k (NESet a) -> a -> k -> Map k (NESet a)
f Map k (NESet a)
acc a
k k
v =
      (NESet a -> NESet a -> NESet a)
-> k -> NESet a -> Map k (NESet a) -> Map k (NESet a)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith NESet a -> NESet a -> NESet a
forall a. Ord a => NESet a -> NESet a -> NESet a
Set.NonEmpty.union k
v (a -> NESet a
forall a. a -> NESet a
Set.NonEmpty.singleton a
k) Map k (NESet a)
acc

-- | Returns the domain of the relation, as a Set, in its entirety.
--
-- /O(a)/.
dom :: BiMultimap a b -> Set a
dom :: forall a b. BiMultimap a b -> Set a
dom =
  Map a (NESet b) -> Set a
forall k a. Map k a -> Set k
Map.keysSet (Map a (NESet b) -> Set a)
-> (BiMultimap a b -> Map a (NESet b)) -> BiMultimap a b -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
toMultimap

-- | Returns the range of the relation, as a Set, in its entirety.
--
-- /O(a)/.
ran :: BiMultimap a b -> Set b
ran :: forall a b. BiMultimap a b -> Set b
ran =
  Map b a -> Set b
forall k a. Map k a -> Set k
Map.keysSet (Map b a -> Set b)
-> (BiMultimap a b -> Map b a) -> BiMultimap a b -> Set b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMultimap a b -> Map b a
forall a b. BiMultimap a b -> Map b a
toMapR

-- | Convert a left-unique relation to a relation (forgetting its left-uniqueness).
toRelation :: (Ord a, Ord b) => BiMultimap a b -> Relation a b
toRelation :: forall a b. (Ord a, Ord b) => BiMultimap a b -> Relation a b
toRelation =
  Map a (Set b) -> Relation a b
forall a b. (Ord a, Ord b) => Map a (Set b) -> Relation a b
Relation.fromMultimap (Map a (Set b) -> Relation a b)
-> (BiMultimap a b -> Map a (Set b))
-> BiMultimap a b
-> Relation a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NESet b -> Set b) -> Map a (NESet b) -> Map a (Set b)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map NESet b -> Set b
forall a. NESet a -> Set a
Set.NonEmpty.toSet (Map a (NESet b) -> Map a (Set b))
-> (BiMultimap a b -> Map a (NESet b))
-> BiMultimap a b
-> Map a (Set b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
domain

-- | Insert a pair into a left-unique relation, maintaining left-uniqueness, preferring the latest inserted element.
--
-- That is, if a left-unique relation already contains the pair @(x, y)@, then inserting the pair @(z, y)@ will cause
-- the @(x, y)@ pair to be deleted.
insert :: (Ord a, Ord b) => a -> b -> BiMultimap a b -> BiMultimap a b
insert :: forall a b.
(Ord a, Ord b) =>
a -> b -> BiMultimap a b -> BiMultimap a b
insert a
a b
b m :: BiMultimap a b
m@(BiMultimap Map a (NESet b)
l Map b a
r) =
  case (Maybe a -> (UpsertResult a, Maybe a))
-> b -> Map b a -> (UpsertResult a, Map b a)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
Map.alterF (a -> Maybe a -> (UpsertResult a, Maybe a)
forall a. Eq a => a -> Maybe a -> (UpsertResult a, Maybe a)
upsertFunc a
a) b
b Map b a
r of
    (UpsertResult a
Ignored, Map b a
_) -> BiMultimap a b
m
    (UpsertResult a
Inserted, Map b a
r') -> Map a (NESet b) -> Map b a -> BiMultimap a b
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap Map a (NESet b)
l' Map b a
r'
    (Replaced a
old, Map b a
r') ->
      let l'' :: Map a (NESet b)
l'' = (NESet b -> Maybe (NESet b))
-> a -> Map a (NESet b) -> Map a (NESet b)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (Set b -> Maybe (NESet b)
forall a. Set a -> Maybe (NESet a)
Set.NonEmpty.nonEmptySet (Set b -> Maybe (NESet b))
-> (NESet b -> Set b) -> NESet b -> Maybe (NESet b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> NESet b -> Set b
forall a. Ord a => a -> NESet a -> Set a
Set.NonEmpty.delete b
b) a
old Map a (NESet b)
l'
       in Map a (NESet b) -> Map b a -> BiMultimap a b
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap Map a (NESet b)
l'' Map b a
r'
  where
    l' :: Map a (NESet b)
l' = (Maybe (NESet b) -> NESet b)
-> a -> Map a (NESet b) -> Map a (NESet b)
forall k v. Ord k => (Maybe v -> v) -> k -> Map k v -> Map k v
Map.upsert (NESet b -> (NESet b -> NESet b) -> Maybe (NESet b) -> NESet b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (b -> NESet b
forall a. a -> NESet a
Set.NonEmpty.singleton b
b) (b -> NESet b -> NESet b
forall a. Ord a => a -> NESet a -> NESet a
Set.NonEmpty.insert b
b)) a
a Map a (NESet b)
l

-- @upsertFunc x@ returns a function that upserts @x@, suitable for passing to @Map.alterF@.
upsertFunc :: (Eq a) => a -> Maybe a -> (UpsertResult a, Maybe a)
upsertFunc :: forall a. Eq a => a -> Maybe a -> (UpsertResult a, Maybe a)
upsertFunc a
new Maybe a
existing =
  case Maybe a
existing of
    Maybe a
Nothing -> (UpsertResult a
forall old. UpsertResult old
Inserted, a -> Maybe a
forall a. a -> Maybe a
Just a
new)
    Just a
old
      | a
old a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
new -> (UpsertResult a
forall old. UpsertResult old
Ignored, Maybe a
existing)
      | Bool
otherwise -> (a -> UpsertResult a
forall old. old -> UpsertResult old
Replaced a
old, a -> Maybe a
forall a. a -> Maybe a
Just a
new)

data UpsertResult old
  = Ignored -- Ignored because an equivalent thing was already there
  | Inserted -- Inserted something new
  | Replaced old -- Replaced what was there, here's the old thing

-- | Like @insert x y@, but the caller is responsible maintaining left-uniqueness.
unsafeInsert :: (Ord a, Ord b) => a -> b -> BiMultimap a b -> BiMultimap a b
unsafeInsert :: forall a b.
(Ord a, Ord b) =>
a -> b -> BiMultimap a b -> BiMultimap a b
unsafeInsert a
x b
y (BiMultimap Map a (NESet b)
xs Map b a
ys) =
  Map a (NESet b) -> Map b a -> BiMultimap a b
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap
    ((Maybe (NESet b) -> NESet b)
-> a -> Map a (NESet b) -> Map a (NESet b)
forall k v. Ord k => (Maybe v -> v) -> k -> Map k v -> Map k v
Map.upsert (NESet b -> (NESet b -> NESet b) -> Maybe (NESet b) -> NESet b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (b -> NESet b
forall a. a -> NESet a
Set.NonEmpty.singleton b
y) (b -> NESet b -> NESet b
forall a. Ord a => a -> NESet a -> NESet a
Set.NonEmpty.insert b
y)) a
x Map a (NESet b)
xs)
    (b -> a -> Map b a -> Map b a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert b
y a
x Map b a
ys)

-- | Union two left-unique relations together.
--
-- The caller is responsible for maintaining left-uniqueness.
unsafeUnion :: (Ord a, Ord b) => BiMultimap a b -> BiMultimap a b -> BiMultimap a b
unsafeUnion :: forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> BiMultimap a b
unsafeUnion BiMultimap a b
xs BiMultimap a b
ys =
  Map a (NESet b) -> Map b a -> BiMultimap a b
forall a b. Map a (NESet b) -> Map b a -> BiMultimap a b
BiMultimap
    ((NESet b -> NESet b -> NESet b)
-> Map a (NESet b) -> Map a (NESet b) -> Map a (NESet b)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith NESet b -> NESet b -> NESet b
forall a. Ord a => NESet a -> NESet a -> NESet a
Set.NonEmpty.union (BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
toMultimap BiMultimap a b
xs) (BiMultimap a b -> Map a (NESet b)
forall a b. BiMultimap a b -> Map a (NESet b)
toMultimap BiMultimap a b
ys))
    (Map b a -> Map b a -> Map b a
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union (BiMultimap a b -> Map b a
forall a b. BiMultimap a b -> Map b a
toMapR BiMultimap a b
xs) (BiMultimap a b -> Map b a
forall a b. BiMultimap a b -> Map b a
toMapR BiMultimap a b
ys))

------------------------------------------------------------------------------------------------------------------------

-- @deriveRangeFromDomain x ys range@ is a helper that inserts @(x, y1)@, @(x, y2)@, ... into range @r@.
deriveRangeFromDomain :: (Ord b) => a -> NESet b -> Map b a -> Map b a
deriveRangeFromDomain :: forall b a. Ord b => a -> NESet b -> Map b a -> Map b a
deriveRangeFromDomain a
x NESet b
ys Map b a
acc =
  (b -> Map b a -> Map b a) -> Map b a -> NESet b -> Map b a
forall a b. (a -> b -> b) -> b -> NESet a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((b -> a -> Map b a -> Map b a) -> a -> b -> Map b a -> Map b a
forall a b c. (a -> b -> c) -> b -> a -> c
flip b -> a -> Map b a -> Map b a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
x) Map b a
acc NESet b
ys
{-# INLINE deriveRangeFromDomain #-}