multi_bimap 0.5.0

// SPDX-FileCopyrightText: 2025 multi_bimap contributors
//
// SPDX-License-Identifier: MIT OR Apache-2.0

#![doc(html_root_url = "https://docs.rs/multi_bimap")]
#![doc = include_str!("../README.md")]
#![cfg_attr(docsrs, doc = "\n## Feature flags\n")]
#![cfg_attr(docsrs, doc = document_features::document_features!())]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![deny(missing_docs)]
#![forbid(unsafe_code)]
#![no_std]

// There's a lot of `<... as Container>` down here because it fixes a compiler
// error that seems to happen due to some type circularity.

#[cfg(feature = "std")]
extern crate std;

extern crate alloc;

use core::borrow::Borrow;

use maplike::containers::Container;
use maplike::one::One;
use maplike::ops::{Clear, Get, Insert, Modify, Put, Remove, WithOne};

#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
/// One-to-one bimap made of two antiparallel hash maps.
pub type HashBimap<L, R> =
    MultiBimap<std::collections::HashMap<L, One<R>>, std::collections::HashMap<R, One<L>>>;
/// One-to-one bimap made of two antiparallel B-tree maps.
pub type BTreeBimap<L, R> =
    MultiBimap<alloc::collections::BTreeMap<L, One<R>>, alloc::collections::BTreeMap<R, One<L>>>;

#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
/// Many-to-many bimap made of two antiparallel hash-set-valued hash maps.
pub type HashMultiBimap<L, R> = MultiBimap<
    std::collections::HashMap<L, std::collections::HashSet<R>>,
    std::collections::HashMap<R, std::collections::HashSet<L>>,
>;
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
/// Many-to-many bimap made of two antiparallel hash-set-valued hash maps.
pub type HashHashMultiBimap<L, R> = MultiBimap<
    std::collections::HashMap<L, std::collections::HashSet<R>>,
    std::collections::HashMap<R, std::collections::HashSet<L>>,
>;
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
/// Many-to-many bimap made of two antiparallel B-tree-set-valued hash maps.
pub type HashBTreeMultiBimap<L, R> = MultiBimap<
    std::collections::HashMap<L, alloc::collections::BTreeSet<R>>,
    std::collections::HashMap<R, alloc::collections::BTreeSet<L>>,
>;
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
/// Many-to-many bimap made of two antiparallel vec-valued hash maps.
pub type HashVecMultiBimap<L, R> = MultiBimap<
    std::collections::HashMap<L, alloc::collections::BTreeSet<R>>,
    std::collections::HashMap<R, alloc::collections::BTreeSet<L>>,
>;

/// Many-to-many bimap made of two antiparallel B-tree-set-valued B-tree maps.
pub type BTreeMultiBimap<L, R> = MultiBimap<
    alloc::collections::BTreeMap<L, alloc::collections::BTreeSet<R>>,
    alloc::collections::BTreeMap<R, alloc::collections::BTreeSet<L>>,
>;
/// Many-to-many bimap made of two antiparallel B-tree-set-valued B-tree maps.
pub type BTreeBTreeMultiBimap<L, R> = MultiBimap<
    alloc::collections::BTreeMap<L, alloc::collections::BTreeSet<R>>,
    alloc::collections::BTreeMap<R, alloc::collections::BTreeSet<L>>,
>;
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
/// Many-to-many bimap made of two antiparallel B-tree-set-valued B-tree maps.
pub type BTreeHashMultiBimap<L, R> = MultiBimap<
    alloc::collections::BTreeMap<L, std::collections::HashSet<R>>,
    alloc::collections::BTreeMap<R, std::collections::HashSet<L>>,
>;
/// Many-to-many bimap made of two antiparallel vec-valued B-tree maps.
pub type BTreeVecMultiBimap<L, R> = MultiBimap<
    alloc::collections::BTreeMap<L, alloc::vec::Vec<R>>,
    alloc::collections::BTreeMap<R, alloc::vec::Vec<L>>,
>;

/// Many-to-many bidirectional map made of two antiparallel maps.
#[derive(Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "undoredo", derive(undoredo::Delta))]
pub struct MultiBimap<L2R, R2L> {
    left_to_right: L2R,
    right_to_left: R2L,
}

impl<L2R, R2L> MultiBimap<L2R, R2L> {
    /// Returns a reference to the left-to-right map.
    pub fn left_to_right(&self) -> &L2R {
        &self.left_to_right
    }

    /// Returns a reference to the right-to-left map.
    pub fn right_to_left(&self) -> &R2L {
        &self.right_to_left
    }
}

impl<L2R: Default, R2L: Default> MultiBimap<L2R, R2L> {
    /// Creates a new, empty `MultiBimap`.
    pub fn new() -> Self {
        MultiBimap {
            left_to_right: Default::default(),
            right_to_left: Default::default(),
        }
    }
}

impl<L2R, R2L> MultiBimap<L2R, R2L>
where
    L2R: Container,
    R2L: Container,
{
    /// Returns the container holding right-side values associated with given
    /// left-side key.
    ///
    /// # Examples
    ///
    /// ```
    /// use multi_bimap::MultiBimap;
    /// use std::collections::{HashMap, HashSet};
    ///
    /// let mut m: MultiBimap<HashMap<&str, HashSet<i32>>, HashMap<i32, HashSet<&str>>> = MultiBimap::new();
    ///
    /// m.insert("a", 1);
    /// m.insert("a", 2);
    ///
    /// assert_eq!(m.get_by_left("a"), Some(&HashSet::from([1, 2])));
    /// assert_eq!(m.get_by_left("missing"), None);
    /// ```
    pub fn get_by_left<Q: ?Sized>(&self, left: &Q) -> Option<&<L2R as Container>::Value>
    where
        L2R: Get<<L2R as Container>::Key, Q>,
        <L2R as Container>::Key: Borrow<Q>,
    {
        self.left_to_right.get(left)
    }

    /// Returns the container holding right-side values associated with given
    /// left-side key.
    ///
    /// # Examples
    ///
    /// ```
    /// use multi_bimap::MultiBimap;
    /// use std::collections::{HashMap, HashSet};
    ///
    /// let mut m: MultiBimap<HashMap<&str, HashSet<i32>>, HashMap<i32, HashSet<&str>>> = MultiBimap::new();
    ///
    /// m.insert("a", 1);
    /// m.insert("b", 1);
    ///
    /// assert_eq!(m.get_by_right(&1), Some(&HashSet::from(["a", "b"])));
    /// assert_eq!(m.get_by_right(&2), None);
    /// ```
    pub fn get_by_right<Q: ?Sized>(&self, right: &Q) -> Option<&<R2L as Container>::Value>
    where
        R2L: Get<<R2L as Container>::Key, Q>,
        <R2L as Container>::Key: Borrow<Q>,
    {
        self.right_to_left.get(right)
    }
}

impl<L2R, R2L> Container for MultiBimap<L2R, R2L>
where
    L2R: Container,
    R2L: Container,
{
    type Key = <L2R as Container>::Key;
    type Value = <R2L as Container>::Key;
}

impl<L2R, R2L> Insert<<L2R as Container>::Key> for MultiBimap<L2R, R2L>
where
    L2R: Container,
    R2L: Container,
    L2R: Get<<L2R as Container>::Key>
        + Insert<<L2R as Container>::Key>
        + Modify<<L2R as Container>::Key>
        + Remove<<L2R as Container>::Key>,
    R2L: Get<<R2L as Container>::Key>
        + Insert<<R2L as Container>::Key>
        + Modify<<R2L as Container>::Key>
        + Remove<<R2L as Container>::Key>,
    <L2R as Container>::Value: WithOne<<R2L as Container>::Key> + Put<<R2L as Container>::Key>,
    <R2L as Container>::Value: WithOne<<L2R as Container>::Key> + Put<<L2R as Container>::Key>,
    <L2R as Container>::Key: Clone + PartialEq,
    <R2L as Container>::Key: Clone + PartialEq,
{
    type Output = (
        Option<<R2L as Container>::Key>,
        Option<<L2R as Container>::Key>,
    );

    fn insert(
        &mut self,
        key: <L2R as Container>::Key,
        value: <R2L as Container>::Key,
    ) -> Self::Output {
        MultiBimap::insert(self, key, value)
    }
}

impl<L2R, R2L> MultiBimap<L2R, R2L>
where
    L2R: Container,
    R2L: Container,
    L2R: Get<<L2R as Container>::Key>
        + Insert<<L2R as Container>::Key>
        + Modify<<L2R as Container>::Key>
        + Remove<<L2R as Container>::Key>,
    R2L: Get<<R2L as Container>::Key>
        + Insert<<R2L as Container>::Key>
        + Modify<<R2L as Container>::Key>
        + Remove<<R2L as Container>::Key>,
    <L2R as Container>::Value: WithOne<<R2L as Container>::Key> + Put<<R2L as Container>::Key>,
    <R2L as Container>::Value: WithOne<<L2R as Container>::Key> + Put<<L2R as Container>::Key>,
{
    /// Insert a left-right association into the bimap.
    ///
    /// Both sides may map to multiple values.
    ///
    /// Returns any values that have been displaced on each side by this
    /// insertion. This can happen if a value container with finite maximum
    /// number of elements is used. For example, [`maplike::one::One`] and
    /// [`Box`] can hold only one element and will always have it displaced and
    /// returned upon insertion.
    ///
    /// # Examples
    ///
    /// ```
    /// use multi_bimap::MultiBimap;
    /// use std::collections::{HashMap, HashSet};
    ///
    /// let mut m: MultiBimap<HashMap<&str, HashSet<i32>>, HashMap<i32, HashSet<&str>>> = MultiBimap::new();
    ///
    /// assert_eq!(m.insert("a", 1), (None, None));
    /// m.insert("a", 2);
    /// m.insert("b", 1);
    ///
    /// assert_eq!(m.get_by_left("a"), Some(&HashSet::from([1, 2])));
    /// ```
    pub fn insert(
        &mut self,
        left: <L2R as Container>::Key,
        right: <R2L as Container>::Key,
    ) -> (
        Option<<R2L as Container>::Key>,
        Option<<L2R as Container>::Key>,
    )
    where
        <L2R as Container>::Key: Clone + PartialEq,
        <R2L as Container>::Key: Clone + PartialEq,
    {
        // PERF: Using Entry API may be faster here, but not all collections
        // support it.

        let left_out = if self.left_to_right.get(&left).is_some() {
            let mut out = None;

            self.left_to_right.modify(&left, |rights| {
                out = rights.put(right.clone());
            });

            out
        } else {
            self.left_to_right
                .insert(left.clone(), WithOne::with_one(right.clone()));
            None
        };

        // A displaced right-side value is no longer associated with `left`, so
        // drop its reverse entry. Skip when it is the same as `right`; the
        // right-side update further below will refresh that entry.
        if let Some(ref old_right) = left_out
            && old_right != &right
        {
            self.right_to_left.remove(old_right);
        }

        let right_out = if self.right_to_left.get(&right).is_some() {
            let mut out = None;
            self.right_to_left.modify(&right, |lefts| {
                out = lefts.put(left.clone());
            });
            out
        } else {
            self.right_to_left
                .insert(right.clone(), WithOne::with_one(left.clone()));
            None
        };

        // Same as above for a displaced left-side key.
        if let Some(ref old_left) = right_out
            && old_left != &left
        {
            self.left_to_right.remove(old_left);
        }

        (left_out, right_out)
    }
}

impl<L2R, R2L> Remove<(<L2R as Container>::Key, <R2L as Container>::Key)> for MultiBimap<L2R, R2L>
where
    L2R: Container,
    R2L: Container,
    L2R: Get<<L2R as Container>::Key>
        + Modify<<L2R as Container>::Key>
        + Remove<<L2R as Container>::Key>,
    R2L: Get<<R2L as Container>::Key>
        + Modify<<R2L as Container>::Key>
        + Remove<<R2L as Container>::Key>,
    <L2R as Container>::Value:
        Remove<<R2L as Container>::Key, Output = Option<()>> + Default + PartialEq,
    <R2L as Container>::Value:
        Remove<<L2R as Container>::Key, Output = Option<()>> + Default + PartialEq,
    <L2R as Container>::Key: Clone,
    <R2L as Container>::Key: Clone,
{
    type Output = Option<(<L2R as Container>::Key, <R2L as Container>::Key)>;

    fn remove(
        &mut self,
        key: &(<L2R as Container>::Key, <R2L as Container>::Key),
    ) -> Option<(<L2R as Container>::Key, <R2L as Container>::Key)> {
        MultiBimap::remove(self, &key.0, &key.1)
    }
}

impl<L2R, R2L> MultiBimap<L2R, R2L>
where
    L2R: Container,
    R2L: Container,
    L2R: Get<<L2R as Container>::Key>
        + Modify<<L2R as Container>::Key>
        + Remove<<L2R as Container>::Key>,
    R2L: Get<<R2L as Container>::Key>
        + Modify<<R2L as Container>::Key>
        + Remove<<R2L as Container>::Key>,
    <L2R as Container>::Value:
        Remove<<R2L as Container>::Key, Output = Option<()>> + Default + PartialEq,
    <R2L as Container>::Value:
        Remove<<L2R as Container>::Key, Output = Option<()>> + Default + PartialEq,
    <L2R as Container>::Key: Clone,
    <R2L as Container>::Key: Clone,
{
    /// Remove a left-right association from the bimap.
    ///
    /// Empty keys are dropped from both sides.
    ///
    /// Returns the removed pair if it was present.
    ///
    /// # Examples
    ///
    /// ```
    /// use multi_bimap::MultiBimap;
    /// use std::collections::{HashMap, HashSet};
    ///
    /// let mut m: MultiBimap<HashMap<&str, HashSet<i32>>, HashMap<i32, HashSet<&str>>> = MultiBimap::new();
    ///
    /// m.insert("a", 1);
    /// m.insert("a", 2);
    /// m.insert("b", 1);
    ///
    /// assert_eq!(m.remove(&"a", &1), Some(("a", 1)));
    /// assert_eq!(m.remove(&"a", &1), None);
    ///
    /// assert_eq!(m.get_by_left("a"), Some(&HashSet::from([2])));
    /// ```
    pub fn remove(
        &mut self,
        left: &<L2R as Container>::Key,
        right: &<R2L as Container>::Key,
    ) -> Option<(<L2R as Container>::Key, <R2L as Container>::Key)> {
        let mut present = false;

        if self.left_to_right.get(left).is_some() {
            self.left_to_right.modify(left, |rights| {
                present = rights.remove(right).is_some();
            });

            if present && self.left_to_right.get(left) == Some(&Default::default()) {
                self.left_to_right.remove(left);
            }
        }

        if self.right_to_left.get(right).is_some() {
            let mut present_right = false;

            self.right_to_left.modify(right, |lefts| {
                present_right = lefts.remove(left).is_some();
            });

            if present_right && self.right_to_left.get(right) == Some(&Default::default()) {
                self.right_to_left.remove(right);
            }
        }

        present.then(|| (left.clone(), right.clone()))
    }
}

impl<L2R, R2L> Clear for MultiBimap<L2R, R2L>
where
    L2R: Clear,
    R2L: Clear,
{
    fn clear(&mut self) {
        MultiBimap::clear(self)
    }
}

impl<L2R, R2L> MultiBimap<L2R, R2L>
where
    L2R: Clear,
    R2L: Clear,
{
    /// Remove all associations from the bimap.
    ///
    /// Clears both the left-to-right and right-to-left sides.
    ///
    /// # Examples
    ///
    /// ```
    /// use multi_bimap::MultiBimap;
    /// use std::collections::{HashMap, HashSet};
    ///
    /// let mut m: MultiBimap<HashMap<&str, HashSet<i32>>, HashMap<i32, HashSet<&str>>> = MultiBimap::new();
    ///
    /// m.insert("a", 1);
    /// m.insert("b", 2);
    /// m.clear();
    ///
    /// assert!(m.left_to_right().is_empty());
    /// assert!(m.right_to_left().is_empty());
    /// ```
    pub fn clear(&mut self) {
        self.left_to_right.clear();
        self.right_to_left.clear();
    }
}

#[cfg(all(test, feature = "std"))]
mod tests {
    use super::*;
    use std::collections::{HashMap, HashSet};

    #[test]
    fn insert_allows_many_on_both_sides() {
        let mut m: MultiBimap<HashMap<&str, HashSet<i32>>, HashMap<i32, HashSet<&str>>> =
            MultiBimap::new();

        m.insert("a", 1);
        m.insert("a", 2);
        m.insert("b", 1);

        assert_eq!(m.get_by_left("a"), Some(&HashSet::from([1, 2])));
        assert_eq!(m.get_by_left("b"), Some(&HashSet::from([1])));
        assert_eq!(m.get_by_right(&1), Some(&HashSet::from(["a", "b"])));
        assert_eq!(m.get_by_right(&2), Some(&HashSet::from(["a"])));
    }

    #[test]
    fn remove_pair_and_drop_empty_keys() {
        let mut m: MultiBimap<HashMap<&str, HashSet<i32>>, HashMap<i32, HashSet<&str>>> =
            MultiBimap::new();

        m.insert("a", 1);
        m.insert("a", 2);
        m.insert("b", 1);

        assert_eq!(m.remove(&"a", &1), Some(("a", 1)));
        assert_eq!(m.remove(&"a", &1), None);

        assert_eq!(m.get_by_left("a"), Some(&HashSet::from([2])));
        assert_eq!(m.get_by_right(&1), Some(&HashSet::from(["b"])));

        assert_eq!(m.remove(&"a", &2), Some(("a", 2)));
        assert!(m.get_by_left("a").is_none());
        assert!(m.get_by_right(&2).is_none());
    }

    #[test]
    fn clear_empties_both_sides() {
        let mut m: MultiBimap<HashMap<&str, HashSet<i32>>, HashMap<i32, HashSet<&str>>> =
            MultiBimap::new();

        m.insert("a", 1);
        m.insert("b", 2);
        m.clear();

        assert!(m.left_to_right().is_empty());
        assert!(m.right_to_left().is_empty());
    }

    #[test]
    fn one_to_one_insert_drops_displaced_reverse() {
        let mut m: HashBimap<&str, &str> = HashBimap::new();

        m.insert("Lithuania", "Vilnius");
        assert_eq!(m.insert("Lithuania", "Kaunas"), (Some("Vilnius"), None));
        assert_eq!(m.get_by_left("Lithuania"), Some(&One::new("Kaunas")));
        assert!(m.get_by_right("Vilnius").is_none());
        assert_eq!(m.get_by_right("Kaunas"), Some(&One::new("Lithuania")));

        assert_eq!(m.insert("Lithuania", "Vilnius"), (Some("Kaunas"), None));
        assert_eq!(m.get_by_left("Lithuania"), Some(&One::new("Vilnius")));
        assert!(m.get_by_right("Kaunas").is_none());
        assert_eq!(m.get_by_right("Vilnius"), Some(&One::new("Lithuania")));
    }

    #[test]
    fn one_to_one_insert_steals_existing_right() {
        let mut m: HashBimap<&str, &str> = HashBimap::new();

        m.insert("Poland", "Warsaw");
        assert_eq!(m.insert("Lithuania", "Warsaw"), (None, Some("Poland")));
        assert_eq!(m.get_by_left("Lithuania"), Some(&One::new("Warsaw")));
        assert!(m.get_by_left("Poland").is_none());
        assert_eq!(m.get_by_right("Warsaw"), Some(&One::new("Lithuania")));
    }
}