From dc3c8dfa36d4c5448dafdbe723f9882bdc1c64a8 Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Mon, 5 Jul 2021 14:30:35 +0200 Subject: [PATCH] BTree: avoid most unsafe code in iterators --- library/alloc/src/collections/btree/map.rs | 39 ++-- .../alloc/src/collections/btree/navigate.rs | 214 ++++++++---------- 2 files changed, 118 insertions(+), 135 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index dfd693d13b330..166c4d5f11e41 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -9,7 +9,7 @@ use core::ops::{Index, RangeBounds}; use core::ptr; use super::borrow::DormantMutRef; -use super::navigate::LeafRange; +use super::navigate::{deallocating_next_unchecked, LeafRange}; use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root}; use super::search::SearchResult::*; @@ -149,7 +149,10 @@ pub struct BTreeMap { unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap { fn drop(&mut self) { if let Some(root) = self.root.take() { - Dropper { front: root.into_dying().first_leaf_edge(), remaining_length: self.length }; + Dropper { + front: Some(root.into_dying().first_leaf_edge()), + remaining_length: self.length, + }; } } } @@ -334,7 +337,7 @@ impl fmt::Debug for IntoIter { /// purpose: to drop the remainder of an `IntoIter`. Therefore it also serves to /// drop an entire tree without the need to first look up a `back` leaf edge. struct Dropper { - front: Handle, marker::Edge>, + front: Option, marker::Edge>>, remaining_length: usize, } @@ -1298,7 +1301,9 @@ impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> { None } else { self.length -= 1; - Some(unsafe { self.range.inner.next_unchecked() }) + let result = self.range.inner.next_checked(); + debug_assert!(result.is_some()); + result } } @@ -1329,7 +1334,9 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Iter<'a, K, V> { None } else { self.length -= 1; - Some(unsafe { self.range.inner.next_back_unchecked() }) + let result = self.range.inner.next_back_checked(); + debug_assert!(result.is_some()); + result } } } @@ -1367,7 +1374,9 @@ impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> { None } else { self.length -= 1; - Some(unsafe { self.range.inner.next_unchecked() }) + let result = self.range.inner.next_checked(); + debug_assert!(result.is_some()); + result } } @@ -1395,7 +1404,9 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for IterMut<'a, K, V> { None } else { self.length -= 1; - Some(unsafe { self.range.inner.next_back_unchecked() }) + let result = self.range.inner.next_back_checked(); + debug_assert!(result.is_some()); + result } } } @@ -1443,11 +1454,13 @@ impl Drop for Dropper { ) -> Option, marker::KV>> { if this.remaining_length == 0 { - unsafe { ptr::read(&this.front).deallocating_end() } + if let Some(front) = this.front.take() { + front.deallocating_end(); + } None } else { this.remaining_length -= 1; - Some(unsafe { this.front.deallocating_next_unchecked() }) + unsafe { deallocating_next_unchecked(&mut this.front) } } } @@ -1474,9 +1487,7 @@ impl Drop for Dropper { #[stable(feature = "btree_drop", since = "1.7.0")] impl Drop for IntoIter { fn drop(&mut self) { - if let Some(front) = self.range.take_front() { - Dropper { front, remaining_length: self.length }; - } + Dropper { front: self.range.take_front(), remaining_length: self.length }; } } @@ -1490,7 +1501,7 @@ impl Iterator for IntoIter { } else { self.length -= 1; let kv = unsafe { self.range.deallocating_next_unchecked() }; - Some(kv.into_key_val()) + Some(kv.unwrap().into_key_val()) } } @@ -1507,7 +1518,7 @@ impl DoubleEndedIterator for IntoIter { } else { self.length -= 1; let kv = unsafe { self.range.deallocating_next_back_unchecked() }; - Some(kv.into_key_val()) + Some(kv.unwrap().into_key_val()) } } } diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index bf3542b384d78..0188dda60230a 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -35,46 +35,42 @@ impl LeafRange { } impl<'a, K, V> LeafRange, K, V> { + /// Moves the leaf edge handle to the next leaf edge and returns + /// references to the key and value in between. #[inline] pub fn next_checked(&mut self) -> Option<(&'a K, &'a V)> { - if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) } + if self.is_empty() { None } else { perform_next(&mut self.front, |kv| kv.into_kv()) } } + /// Moves the leaf edge handle to the previous leaf edge and returns + /// references to the key and value in between. #[inline] pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a V)> { - if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) } - } - - #[inline] - pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) { - unsafe { self.front.as_mut().unwrap().next_unchecked() } - } - - #[inline] - pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) { - unsafe { self.back.as_mut().unwrap().next_back_unchecked() } + if self.is_empty() { None } else { perform_next_back(&mut self.back, |kv| kv.into_kv()) } } } impl<'a, K, V> LeafRange, K, V> { + /// Moves the leaf edge handle to the next leaf edge and returns + /// references to the key and value in between. #[inline] pub fn next_checked(&mut self) -> Option<(&'a K, &'a mut V)> { - if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) } + if self.is_empty() { + None + } else { + perform_next(&mut self.front, |kv| unsafe { ptr::read(kv).into_kv_valmut() }) + } } + /// Moves the leaf edge handle to the previous leaf edge and returns + /// references to the key and value in between. #[inline] pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a mut V)> { - if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) } - } - - #[inline] - pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) { - unsafe { self.front.as_mut().unwrap().next_unchecked() } - } - - #[inline] - pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) { - unsafe { self.back.as_mut().unwrap().next_back_unchecked() } + if self.is_empty() { + None + } else { + perform_next_back(&mut self.back, |kv| unsafe { ptr::read(kv).into_kv_valmut() }) + } } } @@ -89,19 +85,15 @@ impl LeafRange { #[inline] pub unsafe fn deallocating_next_unchecked( &mut self, - ) -> Handle, marker::KV> { - debug_assert!(self.front.is_some()); - let front = self.front.as_mut().unwrap(); - unsafe { front.deallocating_next_unchecked() } + ) -> Option, marker::KV>> { + unsafe { deallocating_next_unchecked(&mut self.front) } } #[inline] pub unsafe fn deallocating_next_back_unchecked( &mut self, - ) -> Handle, marker::KV> { - debug_assert!(self.back.is_some()); - let back = self.back.as_mut().unwrap(); - unsafe { back.deallocating_next_back_unchecked() } + ) -> Option, marker::KV>> { + unsafe { deallocating_next_back_unchecked(&mut self.back) } } } @@ -388,100 +380,80 @@ impl Handle, marker::Edge> { } } -impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Edge> { - /// Moves the leaf edge handle to the next leaf edge and returns references to the - /// key and value in between. - /// - /// # Safety - /// There must be another KV in the direction travelled. - unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) { - super::mem::replace(self, |leaf_edge| { - let kv = leaf_edge.next_kv().ok().unwrap(); - (kv.next_leaf_edge(), kv.into_kv()) - }) - } - - /// Moves the leaf edge handle to the previous leaf edge and returns references to the - /// key and value in between. - /// - /// # Safety - /// There must be another KV in the direction travelled. - unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) { - super::mem::replace(self, |leaf_edge| { - let kv = leaf_edge.next_back_kv().ok().unwrap(); - (kv.next_back_leaf_edge(), kv.into_kv()) - }) - } +/// If possible, extract some result from the next KV and move to the edge beyond it. +fn perform_next( + front: &mut Option, marker::Edge>>, + f: F, +) -> Option +where + F: Fn(&Handle, marker::KV>) -> R, +{ + super::mem::replace(front, |edge| { + if let Some(kv) = edge.and_then(|edge| edge.next_kv().ok()) { + let result = f(&kv); + (Some(kv.next_leaf_edge()), Some(result)) + } else { + (None, None) + } + }) } -impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Edge> { - /// Moves the leaf edge handle to the next leaf edge and returns references to the - /// key and value in between. - /// - /// # Safety - /// There must be another KV in the direction travelled. - unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) { - let kv = super::mem::replace(self, |leaf_edge| { - let kv = leaf_edge.next_kv().ok().unwrap(); - (unsafe { ptr::read(&kv) }.next_leaf_edge(), kv) - }); - // Doing this last is faster, according to benchmarks. - kv.into_kv_valmut() - } - - /// Moves the leaf edge handle to the previous leaf and returns references to the - /// key and value in between. - /// - /// # Safety - /// There must be another KV in the direction travelled. - unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) { - let kv = super::mem::replace(self, |leaf_edge| { - let kv = leaf_edge.next_back_kv().ok().unwrap(); - (unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv) - }); - // Doing this last is faster, according to benchmarks. - kv.into_kv_valmut() - } +/// If possible, extract some result from the previous KV and move to the edge beyond it. +fn perform_next_back( + back: &mut Option, marker::Edge>>, + f: F, +) -> Option +where + F: Fn(&Handle, marker::KV>) -> R, +{ + super::mem::replace(back, |edge| { + if let Some(kv) = edge.and_then(|edge| edge.next_back_kv().ok()) { + let result = f(&kv); + (Some(kv.next_back_leaf_edge()), Some(result)) + } else { + (None, None) + } + }) } -impl Handle, marker::Edge> { - /// Moves the leaf edge handle to the next leaf edge and returns the key and value - /// in between, deallocating any node left behind while leaving the corresponding - /// edge in its parent node dangling. - /// - /// # Safety - /// - There must be another KV in the direction travelled. - /// - That KV was not previously returned by counterpart - /// `deallocating_next_back_unchecked` on any copy of the handles - /// being used to traverse the tree. - /// - /// The only safe way to proceed with the updated handle is to compare it, drop it, - /// or call this method or counterpart `deallocating_next_back_unchecked` again. - pub unsafe fn deallocating_next_unchecked( - &mut self, - ) -> Handle, marker::KV> { - super::mem::replace(self, |leaf_edge| unsafe { leaf_edge.deallocating_next().unwrap() }) - } +/// Moves the leaf edge handle to the next leaf edge and returns the key and value +/// in between, deallocating any node left behind while leaving the corresponding +/// edge in its parent node dangling. +/// +/// # Safety +/// - The next KV, if any, was not previously returned by counterpart +/// `deallocating_next_back_unchecked` on any copy of the handles +/// being used to traverse the tree. +/// +/// The only safe way to proceed with the updated handle is to compare it, drop it, +/// or call this method or counterpart `deallocating_next_back_unchecked` again. +pub unsafe fn deallocating_next_unchecked( + opt_hndl: &mut Option, marker::Edge>>, +) -> Option, marker::KV>> { + super::mem::replace(opt_hndl, |edge| { + edge.and_then(|edge| unsafe { edge.deallocating_next() }) + .map_or((None, None), |p| (Some(p.0), Some(p.1))) + }) +} - /// Moves the leaf edge handle to the previous leaf edge and returns the key and value - /// in between, deallocating any node left behind while leaving the corresponding - /// edge in its parent node dangling. - /// - /// # Safety - /// - There must be another KV in the direction travelled. - /// - That leaf edge was not previously returned by counterpart - /// `deallocating_next_unchecked` on any copy of the handles - /// being used to traverse the tree. - /// - /// The only safe way to proceed with the updated handle is to compare it, drop it, - /// or call this method or counterpart `deallocating_next_unchecked` again. - unsafe fn deallocating_next_back_unchecked( - &mut self, - ) -> Handle, marker::KV> { - super::mem::replace(self, |leaf_edge| unsafe { - leaf_edge.deallocating_next_back().unwrap() - }) - } +/// Moves the leaf edge handle to the previous leaf edge and returns the key and value +/// in between, deallocating any node left behind while leaving the corresponding +/// edge in its parent node dangling. +/// +/// # Safety +/// - The previous KV, if any, was not previously returned by counterpart +/// `deallocating_next_unchecked` on any copy of the handles +/// being used to traverse the tree. +/// +/// The only safe way to proceed with the updated handle is to compare it, drop it, +/// or call this method or counterpart `deallocating_next_unchecked` again. +unsafe fn deallocating_next_back_unchecked( + opt_hndl: &mut Option, marker::Edge>>, +) -> Option, marker::KV>> { + super::mem::replace(opt_hndl, |edge| { + edge.and_then(|edge| unsafe { edge.deallocating_next_back() }) + .map_or((None, None), |p| (Some(p.0), Some(p.1))) + }) } impl NodeRef {