Parameterize the index type of maps and sets

examples/custom_index.rs

@@ -0,0 +1,66 @@
+use fxhash::FxBuildHasher;
+use indexmap::{IndexMap, IndexSet, Indexable};
+use std::ops::{Index, IndexMut};
+
+fn main() {
+    let map: FancyMap<i32, i32> = (-10..=10).map(|i| (i, i * i)).collect();
+    let set: FancySet<i32> = map.values().cloned().collect();
+
+    println!("map to squares: {:?}", map);
+    assert_eq!(map[&-10], map[&10]); // index by key
+    assert_eq!(map[FancyIndex(0)], map[FancyIndex(20)]); // index by position
+
+    println!("unique squares: {:?}", set);
+    assert_eq!(set[FancyIndex(0)], 100); // index by position
+    assert_eq!(set[FancyIndex(10)], 0); // index by position
+}
+
+/// A custom index newtype ensures it can't be confused with indexes for
+/// unrelated containers. This one is also smaller to reduce map memory,
+/// which also reduces the maximum capacity.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub struct FancyIndex(u16);
+
+pub type FancyMap<K, V> = IndexMap<K, V, FxBuildHasher, FancyIndex>;
+pub type FancySet<T> = IndexSet<T, FxBuildHasher, FancyIndex>;
+
+impl Indexable for FancyIndex {}
+
+impl TryFrom<usize> for FancyIndex {
+    type Error = <u16 as TryFrom<usize>>::Error;
+
+    fn try_from(value: usize) -> Result<Self, Self::Error> {
+        Ok(FancyIndex(u16::try_from(value)?))
+    }
+}
+
+impl From<FancyIndex> for usize {
+    fn from(i: FancyIndex) -> usize {
+        i.0.into()
+    }
+}
+
+// Unfortunately, `Index` and `IndexMut for IndexMap` are not automatically implemented for all
+// `Idx: Indexable`, because the compiler considers `Index<Idx>` could potentially overlap with
+// `Index<&Q>` for keys, since references are fundamental. Therefore, we have to implement indexing
+// with `FancyIndex` ourselves. `IndexSet` does already implement all `Index<Idx>` though.
+
+impl<K, V> Index<FancyIndex> for FancyMap<K, V> {
+    type Output = V;
+
+    fn index(&self, index: FancyIndex) -> &V {
+        let (_key, value) = self
+            .get_index(index)
+            .expect("IndexMap: index out of bounds");
+        value
+    }
+}
+
+impl<K, V> IndexMut<FancyIndex> for FancyMap<K, V> {
+    fn index_mut(&mut self, index: FancyIndex) -> &mut V {
+        let (_key, value) = self
+            .get_index_mut(index)
+            .expect("IndexMap: index out of bounds");
+        value
+    }
+}

src/indexable.rs

@@ -0,0 +1,53 @@
+use crate::IndexMap;
+use core::ops::{Index, IndexMut};
+
+/// Indexable types can convert to and from `usize`.
+pub trait Indexable: Sized + Copy + Ord + Send + Sync + TryFrom<usize> + TryInto<usize> {
+    /// Creates an index from a `usize`; panics on failure.
+    #[inline]
+    #[track_caller]
+    fn from_usize(i: usize) -> Self {
+        match Self::try_from(i) {
+            Ok(i) => i,
+            Err(_) => panic!("invalid index!"),
+        }
+    }
+
+    /// Converts the index to a `usize`; panics on failure.
+    #[inline]
+    #[track_caller]
+    fn into_usize(self) -> usize {
+        match self.try_into() {
+            Ok(i) => i,
+            Err(_) => panic!("invalid index!"),
+        }
+    }
+}
+
+impl Indexable for usize {}
+
+macro_rules! impl_indexable {
+    ($($ty:ident)*) => {$(
+        impl Indexable for $ty {}
+
+        impl<K, V, S> Index<$ty> for IndexMap<K, V, S, $ty> {
+            type Output = V;
+
+            fn index(&self, index: $ty) -> &V {
+                self.get_index(index)
+                    .expect("IndexMap: index out of bounds")
+                    .1
+            }
+        }
+
+        impl<K, V, S> IndexMut<$ty> for IndexMap<K, V, S, $ty> {
+            fn index_mut(&mut self, index: $ty) -> &mut V {
+                self.get_index_mut(index)
+                    .expect("IndexMap: index out of bounds")
+                    .1
+            }
+        }
+    )*}
+}
+
+impl_indexable!(u8 u16 u32 u64 u128);

src/lib.rs

@@ -87,6 +87,7 @@ use alloc::vec::{self, Vec};
 #[macro_use]
 mod macros;
 mod equivalent;
+mod indexable;
 mod mutable_keys;
 #[cfg(feature = "serde")]
 mod serde;
@@ -106,6 +107,7 @@ mod rayon;
 mod rustc;
 
 pub use crate::equivalent::Equivalent;
+pub use crate::indexable::Indexable;
 pub use crate::map::IndexMap;
 pub use crate::set::IndexSet;
 
@@ -181,11 +183,16 @@ impl<K, V> Bucket<K, V> {
     }
 }
 
+/// Get basic unbounded access to entries
 trait Entries {
     type Entry;
     fn into_entries(self) -> Vec<Self::Entry>;
     fn as_entries(&self) -> &[Self::Entry];
     fn as_entries_mut(&mut self) -> &mut [Self::Entry];
+}
+
+/// Borrow entries and fix the indices afterward
+trait WithEntries: Entries {
     fn with_entries<F>(&mut self, f: F)
     where
         F: FnOnce(&mut [Self::Entry]);