From d3f1968e917ee594cf7c9d8619e9aa97dfc288aa Mon Sep 17 00:00:00 2001 From: Niklas Jonsson Date: Thu, 14 Jul 2022 19:52:18 +0200 Subject: [PATCH 1/6] Implement get_many_mut --- src/map.rs | 46 ++++++++++++++++++++++++++++++++++++ src/map/tests.rs | 61 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 107 insertions(+) diff --git a/src/map.rs b/src/map.rs index e71c799c..34787326 100644 --- a/src/map.rs +++ b/src/map.rs @@ -477,6 +477,52 @@ where } } + pub fn get_many_mut<'a, 'b, Q: ?Sized, const N: usize>( + &'a mut self, + keys: [&'b Q; N], + ) -> Option<[&'a mut V; N]> + where + Q: Hash + Equivalent, + { + let indices = keys.map(|key| self.get_index_of(key)); + if indices.iter().any(Option::is_none) { + return None; + } + let indices = indices.map(Option::unwrap); + + // SAFETY: Can't allow duplicate indices as we would return several mutable refs to the same data + for i in 0..N { + let idx = indices[i]; + if indices[i + 1..N].contains(&idx) { + return None; + } + } + + // Replace with MaybeUninit::uninit_array when that is stable + // SAFETY: Creating MaybeUninit from uninit is always safe + #[allow(unsafe_code)] + let mut out: [std::mem::MaybeUninit<&'a mut V>; N] = + unsafe { std::mem::MaybeUninit::uninit().assume_init() }; + + let entries = self.as_entries_mut(); + for (elem, idx) in out.iter_mut().zip(indices) { + let v: &mut V = &mut entries[idx].value; + // SAFETY: As we know that each index is unique, it is OK to discard the mutable + // borrow lifetime of v, we will never mutably borrow an element twice. + // The pointer is valid and aligned as we get it from MaybeUninit. + #[allow(unsafe_code)] + unsafe { std::ptr::write(elem.as_mut_ptr(), &mut *(v as *mut V)) }; + } + + // Can't transmute a const-sized array: + // https://github.com/rust-lang/rust/issues/61956 + // This is the workaround. + // SAFETY: This is fine as the references all are from unique entries that we own and all of + // them have been properly initialized by the above loop. + #[allow(unsafe_code)] + Some(unsafe { std::mem::transmute_copy::<_, [&'a mut V; N]>(&out) }) + } + /// Remove the key-value pair equivalent to `key` and return /// its value. /// diff --git a/src/map/tests.rs b/src/map/tests.rs index b6c6a42d..cbf856a0 100644 --- a/src/map/tests.rs +++ b/src/map/tests.rs @@ -418,3 +418,64 @@ fn from_array() { assert_eq!(map, expected) } + +#[test] +fn many_mut_empty() { + let mut map: IndexMap = IndexMap::default(); + assert!(map.get_many_mut([&0, &1, &2, &3]).is_none()); +} + +#[test] +fn many_mut_single_fail() { + let mut map: IndexMap = IndexMap::default(); + map.insert(1, 10); + assert!(map.get_many_mut([&0]).is_none()); +} + +#[test] +fn many_mut_single_success() { + let mut map: IndexMap = IndexMap::default(); + map.insert(1, 10); + assert_eq!(map.get_many_mut([&1]), Some([&mut 10])); +} + +#[test] +fn many_mut_multi_success() { + let mut map: IndexMap = IndexMap::default(); + map.insert(1, 10); + map.insert(1123, 100); + map.insert(321, 20); + map.insert(1337, 30); + assert_eq!(map.get_many_mut([&1, &1123]), Some([&mut 10, &mut 100])); + assert_eq!(map.get_many_mut([&1, &1337]), Some([&mut 10, &mut 30])); + assert_eq!( + map.get_many_mut([&1337, &321, &1, &1123]), + Some([&mut 30, &mut 20, &mut 10, &mut 100]) + ); +} + +#[test] +fn many_mut_multi_fail_missing() { + let mut map: IndexMap = IndexMap::default(); + map.insert(1, 10); + map.insert(1123, 100); + map.insert(321, 20); + map.insert(1337, 30); + assert_eq!(map.get_many_mut([&121, &1123]), None); + assert_eq!(map.get_many_mut([&1, &1337, &56]), None); + assert_eq!(map.get_many_mut([&1337, &123, &321, &1, &1123]), None); +} + +#[test] +fn many_mut_multi_fail_duplicate() { + let mut map: IndexMap = IndexMap::default(); + map.insert(1, 10); + map.insert(1123, 100); + map.insert(321, 20); + map.insert(1337, 30); + assert_eq!(map.get_many_mut([&1, &1]), None); + assert_eq!( + map.get_many_mut([&1337, &123, &321, &1337, &1, &1123]), + None + ); +} From 66108be0ed23a74b58a8c9bdfb7840a61d987cc6 Mon Sep 17 00:00:00 2001 From: Niklas Jonsson Date: Fri, 15 Jul 2022 15:26:50 +0200 Subject: [PATCH 2/6] Reduce (unsafe) code with array::map and pointers --- src/map.rs | 27 +++++++-------------------- 1 file changed, 7 insertions(+), 20 deletions(-) diff --git a/src/map.rs b/src/map.rs index 34787326..b0ceef0c 100644 --- a/src/map.rs +++ b/src/map.rs @@ -498,29 +498,16 @@ where } } - // Replace with MaybeUninit::uninit_array when that is stable - // SAFETY: Creating MaybeUninit from uninit is always safe - #[allow(unsafe_code)] - let mut out: [std::mem::MaybeUninit<&'a mut V>; N] = - unsafe { std::mem::MaybeUninit::uninit().assume_init() }; - let entries = self.as_entries_mut(); - for (elem, idx) in out.iter_mut().zip(indices) { - let v: &mut V = &mut entries[idx].value; - // SAFETY: As we know that each index is unique, it is OK to discard the mutable - // borrow lifetime of v, we will never mutably borrow an element twice. - // The pointer is valid and aligned as we get it from MaybeUninit. + let out = indices.map(|i| { + // SAFETY: OK to discard mutable borrow lifetime as each index is unique #[allow(unsafe_code)] - unsafe { std::ptr::write(elem.as_mut_ptr(), &mut *(v as *mut V)) }; - } + unsafe { + &mut *(&mut entries[i].value as *mut V) + } + }); - // Can't transmute a const-sized array: - // https://github.com/rust-lang/rust/issues/61956 - // This is the workaround. - // SAFETY: This is fine as the references all are from unique entries that we own and all of - // them have been properly initialized by the above loop. - #[allow(unsafe_code)] - Some(unsafe { std::mem::transmute_copy::<_, [&'a mut V; N]>(&out) }) + Some(out) } /// Remove the key-value pair equivalent to `key` and return From 72f0146c64e5fd457ccddf1ab237c8874fd5b103 Mon Sep 17 00:00:00 2001 From: Niklas Jonsson Date: Fri, 15 Jul 2022 20:30:10 +0200 Subject: [PATCH 3/6] Implement get_many_index_mut & and add docs Avoid mutable aliasing by using a pointer to entries --- src/map.rs | 77 ++++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 57 insertions(+), 20 deletions(-) diff --git a/src/map.rs b/src/map.rs index b0ceef0c..1241ae4c 100644 --- a/src/map.rs +++ b/src/map.rs @@ -477,6 +477,15 @@ where } } + /// Return the values for `N` keys. If any key is missing a value, or there + /// are duplicate keys, `None` is returned. + /// + /// # Examples + /// + /// ``` + /// let mut map = indexmap::IndexMap::from([(1, 'a'), (3, 'b'), (2, 'c')]); + /// assert_eq!(map.get_many_mut([&2, &1]), Some([&mut 'c', &mut 'a'])); + /// ``` pub fn get_many_mut<'a, 'b, Q: ?Sized, const N: usize>( &'a mut self, keys: [&'b Q; N], @@ -484,30 +493,19 @@ where where Q: Hash + Equivalent, { + let len = self.len(); let indices = keys.map(|key| self.get_index_of(key)); - if indices.iter().any(Option::is_none) { - return None; - } - let indices = indices.map(Option::unwrap); - // SAFETY: Can't allow duplicate indices as we would return several mutable refs to the same data - for i in 0..N { - let idx = indices[i]; - if indices[i + 1..N].contains(&idx) { - return None; + // Handle out-of-bounds indices with panic as this is an internal error in get_index_of. + for idx in indices { + let idx = idx?; + if idx >= len { + panic!("Index is out of range! Got '{idx}' but length is '{len}'") } } - - let entries = self.as_entries_mut(); - let out = indices.map(|i| { - // SAFETY: OK to discard mutable borrow lifetime as each index is unique - #[allow(unsafe_code)] - unsafe { - &mut *(&mut entries[i].value as *mut V) - } - }); - - Some(out) + let indices = indices.map(Option::unwrap); + let entries = self.get_many_index_mut(indices)?; + Some(entries.map(|(_key, value)| value)) } /// Remove the key-value pair equivalent to `key` and return @@ -817,6 +815,45 @@ impl IndexMap { self.as_entries_mut().get_mut(index).map(Bucket::ref_mut) } + /// Get an array of `N` key-value pairs by `N` indices + /// + /// Valid indices are *0 <= index < self.len()* and each index needs to be unique. + /// + /// Computes in **O(1)** time. + /// + /// # Examples + /// + /// ``` + /// let mut map = indexmap::IndexMap::from([(1, 'a'), (3, 'b'), (2, 'c')]); + /// assert_eq!(map.get_many_index_mut([2, 0]), Some([(&2, &mut 'c'), (&1, &mut 'a')])); + /// ``` + pub fn get_many_index_mut( + &mut self, + indices: [usize; N], + ) -> Option<[(&K, &mut V); N]> { + // SAFETY: Can't allow duplicate indices as we would return several mutable refs to the same data. + // Additionally, handle out-of-bounds indices (internal error in get_index_of) with panic. + let len = self.len(); + for i in 0..N { + let idx = indices[i]; + if idx >= len || indices[i + 1..N].contains(&idx) { + return None; + } + } + + let entries_ptr = self.as_entries_mut().as_mut_ptr(); + let out = indices.map(|i| { + // SAFETY: The base pointer is valid as it comes from a slice and the deref is always + // in-bounds as we've already checked the indices above. + #[allow(unsafe_code)] + unsafe { + (*(entries_ptr.add(i))).ref_mut() + } + }); + + Some(out) + } + /// Returns a slice of key-value pairs in the given range of indices. /// /// Valid indices are *0 <= index < self.len()* From 5795fe327434411793de0c1658a667871525571f Mon Sep 17 00:00:00 2001 From: Niklas Jonsson Date: Fri, 15 Jul 2022 21:30:26 +0200 Subject: [PATCH 4/6] fix comment and add test for many_index oob --- src/map.rs | 1 - src/map/tests.rs | 8 ++++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/src/map.rs b/src/map.rs index 1241ae4c..6e0ebbb5 100644 --- a/src/map.rs +++ b/src/map.rs @@ -832,7 +832,6 @@ impl IndexMap { indices: [usize; N], ) -> Option<[(&K, &mut V); N]> { // SAFETY: Can't allow duplicate indices as we would return several mutable refs to the same data. - // Additionally, handle out-of-bounds indices (internal error in get_index_of) with panic. let len = self.len(); for i in 0..N { let idx = indices[i]; diff --git a/src/map/tests.rs b/src/map/tests.rs index cbf856a0..c2ae8ede 100644 --- a/src/map/tests.rs +++ b/src/map/tests.rs @@ -479,3 +479,11 @@ fn many_mut_multi_fail_duplicate() { None ); } + +#[test] +fn many_index_mut_fail_oob() { + let mut map: IndexMap = IndexMap::default(); + map.insert(1, 10); + map.insert(321, 20); + assert_eq!(map.get_many_index_mut([1, 3]), None); +} From dfa303b334c83ab8931806df15ba18d743e30b8e Mon Sep 17 00:00:00 2001 From: Niklas Jonsson Date: Fri, 15 Jul 2022 21:51:12 +0200 Subject: [PATCH 5/6] change oob panic to debug assert --- src/map.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/map.rs b/src/map.rs index 6e0ebbb5..e79c896f 100644 --- a/src/map.rs +++ b/src/map.rs @@ -499,9 +499,10 @@ where // Handle out-of-bounds indices with panic as this is an internal error in get_index_of. for idx in indices { let idx = idx?; - if idx >= len { - panic!("Index is out of range! Got '{idx}' but length is '{len}'") - } + debug_assert!( + idx < len, + "Index is out of range! Got '{idx}' but length is '{len}'" + ); } let indices = indices.map(Option::unwrap); let entries = self.get_many_index_mut(indices)?; From 9f800a0bf8047e73b05d9f537084543c4d0608f4 Mon Sep 17 00:00:00 2001 From: Niklas Jonsson Date: Sat, 16 Jul 2022 18:42:08 +0200 Subject: [PATCH 6/6] Avoid using format features not preset on MSRV --- src/map.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/map.rs b/src/map.rs index e79c896f..bb915e30 100644 --- a/src/map.rs +++ b/src/map.rs @@ -501,7 +501,9 @@ where let idx = idx?; debug_assert!( idx < len, - "Index is out of range! Got '{idx}' but length is '{len}'" + "Index is out of range! Got '{}' but length is '{}'", + idx, + len ); } let indices = indices.map(Option::unwrap);