Code shared from the Rust Playground

playground.rs

use std::borrow::Borrow;
use std::collections::HashMap;
use std::hash::Hash;
use std::ops::Deref;
use std::pin::Pin;
use std::ptr::NonNull;

/// Wrapper around `Pin<Box<T>>`.
///
/// This is necessary in order to implement `Borrow<T>` for `Pin<Box<T>>`.
#[derive(Eq, Hash, PartialEq)]
struct PinBox<T> {
    inner: Pin<Box<T>>,
}

impl<T> PinBox<T> {
    /// Creates a `PinBox` that pins the given value.
    fn new(value: T) -> Self {
        Self {
            inner: Box::pin(value),
        }
    }
}

impl<T> Borrow<T> for PinBox<T> {
    fn borrow(&self) -> &T {
        self.inner.borrow()
    }
}

/// A bijective hashmap.
///
/// Maintains the invariant that each left value is associated with exactly one right value and
/// vice-versa.
pub struct BiHashMap<L, R> {
    // each map owns its keys and has values that are pointers to the corresponding key in the
    // other map. the keys are pinned so that the pointers are guaranteed to be valid. the
    // invariant implies that each `PinBox` has exactly one pointer pointing to it.
    left_map: HashMap<PinBox<L>, NonNull<R>>,
    right_map: HashMap<PinBox<R>, NonNull<L>>,
}

impl<L, R> BiHashMap<L, R>
where
    L: Eq + Hash,
    R: Eq + Hash,
{
    /// Creates an empty `BiHashMap`.
    pub fn new() -> Self {
        Self {
            left_map: HashMap::new(),
            right_map: HashMap::new(),
        }
    }

    /// Attempts to insert the given left-right pair into the `BiHashMap`.
    ///
    /// If either the left value or the right value already exists in the bimap, the insertion
    /// fails and the values are returned. This is necessary in order to maintain the bijection
    /// invariant.
    pub fn insert(&mut self, left: L, right: R) -> Result<(), (L, R)> {
        // check if either value already exists in the bimap, and return with an error if so.
        if self.left_map.contains_key(&left) || self.right_map.contains_key(&right) {
            Err((left, right))
        } else {
            // pin the left value
            let left_pin = PinBox::new(left);
            // get a pointer to the pinned left value
            let left_ptr = NonNull::from(left_pin.inner.deref());

            // same for the right value
            let right_pin = PinBox::new(right);
            let right_ptr = NonNull::from(right_pin.inner.deref());

            // perform the insertions
            self.left_map.insert(left_pin, right_ptr);
            self.right_map.insert(right_pin, left_ptr);

            Ok(())
        }
    }

    /// Returns a reference to the right value corresponding to the given left value.
    pub fn get_by_left(&self, left: &L) -> Option<&R> {
        self.left_map
            .get(left)
            // safe to dereference `right_ptr` because the pinned right value that it points to is
            // guaranteed to exist as a key in `self.right_map`.
            .map(|right_ptr| unsafe { right_ptr.as_ref() })
    }

    /// Removes the left-right pair corresponding to the given left value.
    pub fn remove_by_left(&mut self, left: &L) -> Option<(L, R)> {
        self.left_map
            .remove_entry(left)
            .map(|(left_pin, right_ptr)| {
                // safe to dereference `right_ptr` because the pinned right value that it points to
                // is guaranteed to exist as a key in `self.right_map`.
                let right = unsafe { right_ptr.as_ref() };

                // safe to unwrap because the key is guaranteed to exist as part of the bijection
                // invariant.
                let (right_pin, _left_ptr) = self.right_map.remove_entry(right).unwrap();

                // at this point, `_left_ptr` and `right_ptr` are no longer used and dropped at the
                // end of this method, so it's safe to unpin `left_pin` and `right_pin` since
                // nothing points to them anymore.
                unsafe {
                    (
                        *Pin::into_inner_unchecked(left_pin.inner),
                        *Pin::into_inner_unchecked(right_pin.inner),
                    )
                }
            })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test() {
        let mut bimap = BiHashMap::new();

        assert!(bimap.insert('a', 1).is_ok());
        assert!(bimap.insert('b', 2).is_ok());
        assert!(bimap.insert('c', 3).is_ok());
        assert!(bimap.insert('a', 4).is_err());

        assert_eq!(bimap.get_by_left(&'c'), Some(&3));
        assert_eq!(bimap.get_by_left(&'d'), None);

        assert_eq!(bimap.remove_by_left(&'a'), Some(('a', 1)));
        assert_eq!(bimap.remove_by_left(&'b'), Some(('b', 2)));
        assert_eq!(bimap.remove_by_left(&'b'), None);
    }
}