From ddb0316dc25e12fb0ecda55f0a4e754dd62e6e5c Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Fri, 1 May 2026 16:48:15 -0700 Subject: [PATCH] Separate the implementations of `get_disjoint_mut` At first glance this may seem worse because it's roughly duplicating the `unsafe` code. However, these are now placed in a dedicated module where the parallels are apparent, and they should be more direct for their respective use-cases. * For disjoint key lookup in a map, we don't need to produce a `Result` that we're just going to turn into panics anyway. * For disjoint indices in a slice, we don't need to deal with `Option`, and we can drop our own unsafe implementation here in favor of the standard library once we update to MSRV 1.86. --- src/map.rs | 19 ++++++-------- src/map/disjoint.rs | 60 +++++++++++++++++++++++++++++++++++++++++++++ src/map/slice.rs | 39 +++-------------------------- 3 files changed, 70 insertions(+), 48 deletions(-) create mode 100644 src/map/disjoint.rs diff --git a/src/map.rs b/src/map.rs index 1c2bb4b4..803d4107 100644 --- a/src/map.rs +++ b/src/map.rs @@ -1,6 +1,7 @@ //! [`IndexMap`] is a hash table where the iteration order of the key-value //! pairs is independent of the hash values of the keys. +mod disjoint; mod entry; mod iter; mod mutable; @@ -952,7 +953,10 @@ where /// /// ``` /// let mut map = indexmap::IndexMap::from([(1, 'a'), (3, 'b'), (2, 'c')]); - /// assert_eq!(map.get_disjoint_mut([&2, &1]), [Some(&mut 'c'), Some(&mut 'a')]); + /// assert_eq!( + /// map.get_disjoint_mut([&2, &1, &0]), + /// [Some(&mut 'c'), Some(&mut 'a'), None], + /// ); /// ``` #[track_caller] pub fn get_disjoint_mut(&mut self, keys: [&Q; N]) -> [Option<&mut V>; N] @@ -960,17 +964,8 @@ where Q: ?Sized + Hash + Equivalent, { let indices = keys.map(|key| self.get_index_of(key)); - match self.as_mut_slice().get_disjoint_opt_mut(indices) { - Err(GetDisjointMutError::IndexOutOfBounds) => { - unreachable!( - "Internal error: indices should never be OOB as we got them from get_index_of" - ); - } - Err(GetDisjointMutError::OverlappingIndices) => { - panic!("duplicate keys found"); - } - Ok(key_values) => key_values.map(|kv_opt| kv_opt.map(|kv| kv.1)), - } + disjoint::get_disjoint_opt_mut(self.as_entries_mut(), indices) + .map(|opt| opt.map(Bucket::value_mut)) } /// Remove the key-value pair equivalent to `key` and return diff --git a/src/map/disjoint.rs b/src/map/disjoint.rs new file mode 100644 index 00000000..90a0702f --- /dev/null +++ b/src/map/disjoint.rs @@ -0,0 +1,60 @@ +#![allow(unsafe_code)] + +use crate::GetDisjointMutError; + +/// Like `slice::get_disjoint_mut`, although we're not dealing with ranges (yet). +// TODO(MSRV 1.86): remove this in favor of the standard library's method. +pub(super) fn get_disjoint_mut( + entries: &mut [T], + indices: [usize; N], +) -> Result<[&mut T; N], GetDisjointMutError> { + // SAFETY: Can't allow duplicate indices as we would return mutable refs to the same data. + let len = entries.len(); + for i in 0..N { + let idx = indices[i]; + if idx >= len { + return Err(GetDisjointMutError::IndexOutOfBounds); + } else if indices[..i].contains(&idx) { + return Err(GetDisjointMutError::OverlappingIndices); + } + } + + let entries_ptr = entries.as_mut_ptr(); + Ok(indices.map(move |idx| { + // SAFETY: The base pointer is valid as it comes from a slice and the reference is always + // in-bounds & unique as we've already checked the indices above. + unsafe { &mut *(entries_ptr.add(idx)) } + })) +} + +/// Like `slice::get_disjoint_mut` but with optional indices, +/// allowing for absent keys from the user's original request. +#[track_caller] +pub(super) fn get_disjoint_opt_mut( + entries: &mut [T], + indices: [Option; N], +) -> [Option<&mut T>; N] { + // SAFETY: Can't allow duplicate indices as we would return mutable refs to the same data. + let len = entries.len(); + for i in 0..N { + if let Some(idx) = indices[i] { + if idx >= len { + unreachable!("`get_index_of` returned an out-of-bounds index"); + } else if indices[..i].contains(&Some(idx)) { + panic!("duplicate keys found"); + } + } + } + + let entries_ptr = entries.as_mut_ptr(); + indices.map(move |idx_opt| { + match idx_opt { + Some(idx) => { + // SAFETY: The base pointer is valid as it comes from a slice and the reference is always + // in-bounds & unique as we've already checked the indices above. + Some(unsafe { &mut *entries_ptr.add(idx) }) + } + None => None, + } + }) +} diff --git a/src/map/slice.rs b/src/map/slice.rs index 41a483fd..1243739e 100644 --- a/src/map/slice.rs +++ b/src/map/slice.rs @@ -350,42 +350,9 @@ impl Slice { &mut self, indices: [usize; N], ) -> Result<[(&K, &mut V); N], GetDisjointMutError> { - let indices = indices.map(Some); - let key_values = self.get_disjoint_opt_mut(indices)?; - Ok(key_values.map(Option::unwrap)) - } - - #[allow(unsafe_code)] - pub(crate) fn get_disjoint_opt_mut( - &mut self, - indices: [Option; N], - ) -> Result<[Option<(&K, &mut V)>; N], GetDisjointMutError> { - // SAFETY: Can't allow duplicate indices as we would return several mutable refs to the same data. - let len = self.len(); - for i in 0..N { - if let Some(idx) = indices[i] { - if idx >= len { - return Err(GetDisjointMutError::IndexOutOfBounds); - } else if indices[..i].contains(&Some(idx)) { - return Err(GetDisjointMutError::OverlappingIndices); - } - } - } - - let entries_ptr = self.entries.as_mut_ptr(); - let out = indices.map(|idx_opt| { - match idx_opt { - Some(idx) => { - // SAFETY: The base pointer is valid as it comes from a slice and the reference is always - // in-bounds & unique as we've already checked the indices above. - let kv = unsafe { (*(entries_ptr.add(idx))).ref_mut() }; - Some(kv) - } - None => None, - } - }); - - Ok(out) + // TODO(MSRV 1.86): use the standard library's `slice::get_disjoint_mut` + let entries = super::disjoint::get_disjoint_mut(&mut self.entries, indices)?; + Ok(entries.map(Bucket::ref_mut)) } }