BTree: avoid most unsafe code in iterators

ssomers · ssomers · commit dc3c8dfa36d4 · 2021-07-08T22:21:35.000+02:00
diff --git 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<K, V> {
 unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap<K, V> {
     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<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IntoIter<K, V> {
 /// 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<K, V> {
-    front: Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
+    front: Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, 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<K, V> Drop for Dropper<K, V> {
         ) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, 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<K, V> Drop for Dropper<K, V> {
 #[stable(feature = "btree_drop", since = "1.7.0")]
 impl<K, V> Drop for IntoIter<K, V> {
     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<K, V> Iterator for IntoIter<K, V> {
         } 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<K, V> DoubleEndedIterator for IntoIter<K, V> {
         } 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
@@ -35,46 +35,42 @@ impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
 }
 
 impl<'a, K, V> LeafRange<marker::Immut<'a>, 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<marker::ValMut<'a>, 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<K, V> LeafRange<marker::Dying, K, V> {
     #[inline]
     pub unsafe fn deallocating_next_unchecked(
         &mut self,
-    ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
-        debug_assert!(self.front.is_some());
-        let front = self.front.as_mut().unwrap();
-        unsafe { front.deallocating_next_unchecked() }
+    ) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>> {
+        unsafe { deallocating_next_unchecked(&mut self.front) }
     }
 
     #[inline]
     pub unsafe fn deallocating_next_back_unchecked(
         &mut self,
-    ) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
-        debug_assert!(self.back.is_some());
-        let back = self.back.as_mut().unwrap();
-        unsafe { back.deallocating_next_back_unchecked() }
+    ) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>> {
+        unsafe { deallocating_next_back_unchecked(&mut self.back) }
     }
 }
 
@@ -388,100 +380,80 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
     }
 }
 
-impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, 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<BorrowType: marker::BorrowType, K, V, F, R>(
+    front: &mut Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
+    f: F,
+) -> Option<R>
+where
+    F: Fn(&Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, 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<NodeRef<marker::ValMut<'a>, 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<BorrowType: marker::BorrowType, K, V, F, R>(
+    back: &mut Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
+    f: F,
+) -> Option<R>
+where
+    F: Fn(&Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, 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<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, 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<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, 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<K, V>(
+    opt_hndl: &mut Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
+) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, 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<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, 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<K, V>(
+    opt_hndl: &mut Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
+) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, 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<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {

Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ use core::ops::{Index, RangeBounds};`
`9`	`9`	`use core::ptr;`
`10`	`10`
`11`	`11`	`use super::borrow::DormantMutRef;`
`12`		`-use super::navigate::LeafRange;`
	`12`	`+use super::navigate::{deallocating_next_unchecked, LeafRange};`
`13`	`13`	`use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};`
`14`	`14`	`use super::search::SearchResult::*;`
`15`	`15`
`@@ -149,7 +149,10 @@ pub struct BTreeMap<K, V> {`
`149`	`149`	`unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap<K, V> {`
`150`	`150`	`fn drop(&mut self) {`
`151`	`151`	`if let Some(root) = self.root.take() {`
`152`		`- Dropper { front: root.into_dying().first_leaf_edge(), remaining_length: self.length };`
	`152`	`+ Dropper {`
	`153`	`+ front: Some(root.into_dying().first_leaf_edge()),`
	`154`	`+ remaining_length: self.length,`
	`155`	`+ };`
`153`	`156`	`}`
`154`	`157`	`}`
`155`	`158`	`}`
`@@ -334,7 +337,7 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IntoIter<K, V> {`
`334`	`337`	/// purpose: to drop the remainder of an `IntoIter`. Therefore it also serves to
`335`	`338`	/// drop an entire tree without the need to first look up a `back` leaf edge.
`336`	`339`	`struct Dropper<K, V> {`
`337`		`- front: Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,`
	`340`	`+ front: Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,`
`338`	`341`	`remaining_length: usize,`
`339`	`342`	`}`
`340`	`343`
`@@ -1298,7 +1301,9 @@ impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> {`
`1298`	`1301`	`None`
`1299`	`1302`	`} else {`
`1300`	`1303`	`self.length -= 1;`
`1301`		`- Some(unsafe { self.range.inner.next_unchecked() })`
	`1304`	`+ let result = self.range.inner.next_checked();`
	`1305`	`+ debug_assert!(result.is_some());`
	`1306`	`+ result`
`1302`	`1307`	`}`
`1303`	`1308`	`}`
`1304`	`1309`
`@@ -1329,7 +1334,9 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Iter<'a, K, V> {`
`1329`	`1334`	`None`
`1330`	`1335`	`} else {`
`1331`	`1336`	`self.length -= 1;`
`1332`		`- Some(unsafe { self.range.inner.next_back_unchecked() })`
	`1337`	`+ let result = self.range.inner.next_back_checked();`
	`1338`	`+ debug_assert!(result.is_some());`
	`1339`	`+ result`
`1333`	`1340`	`}`
`1334`	`1341`	`}`
`1335`	`1342`	`}`
`@@ -1367,7 +1374,9 @@ impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {`
`1367`	`1374`	`None`
`1368`	`1375`	`} else {`
`1369`	`1376`	`self.length -= 1;`
`1370`		`- Some(unsafe { self.range.inner.next_unchecked() })`
	`1377`	`+ let result = self.range.inner.next_checked();`
	`1378`	`+ debug_assert!(result.is_some());`
	`1379`	`+ result`
`1371`	`1380`	`}`
`1372`	`1381`	`}`
`1373`	`1382`
`@@ -1395,7 +1404,9 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for IterMut<'a, K, V> {`
`1395`	`1404`	`None`
`1396`	`1405`	`} else {`
`1397`	`1406`	`self.length -= 1;`
`1398`		`- Some(unsafe { self.range.inner.next_back_unchecked() })`
	`1407`	`+ let result = self.range.inner.next_back_checked();`
	`1408`	`+ debug_assert!(result.is_some());`
	`1409`	`+ result`
`1399`	`1410`	`}`
`1400`	`1411`	`}`
`1401`	`1412`	`}`
`@@ -1443,11 +1454,13 @@ impl<K, V> Drop for Dropper<K, V> {`
`1443`	`1454`	`) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>>`
`1444`	`1455`	`{`
`1445`	`1456`	`if this.remaining_length == 0 {`
`1446`		`- unsafe { ptr::read(&this.front).deallocating_end() }`
	`1457`	`+ if let Some(front) = this.front.take() {`
	`1458`	`+ front.deallocating_end();`
	`1459`	`+ }`
`1447`	`1460`	`None`
`1448`	`1461`	`} else {`
`1449`	`1462`	`this.remaining_length -= 1;`
`1450`		`- Some(unsafe { this.front.deallocating_next_unchecked() })`
	`1463`	`+ unsafe { deallocating_next_unchecked(&mut this.front) }`
`1451`	`1464`	`}`
`1452`	`1465`	`}`
`1453`	`1466`
`@@ -1474,9 +1487,7 @@ impl<K, V> Drop for Dropper<K, V> {`
`1474`	`1487`	`#[stable(feature = "btree_drop", since = "1.7.0")]`
`1475`	`1488`	`impl<K, V> Drop for IntoIter<K, V> {`
`1476`	`1489`	`fn drop(&mut self) {`
`1477`		`- if let Some(front) = self.range.take_front() {`
`1478`		`- Dropper { front, remaining_length: self.length };`
`1479`		`- }`
	`1490`	`+ Dropper { front: self.range.take_front(), remaining_length: self.length };`
`1480`	`1491`	`}`
`1481`	`1492`	`}`
`1482`	`1493`
`@@ -1490,7 +1501,7 @@ impl<K, V> Iterator for IntoIter<K, V> {`
`1490`	`1501`	`} else {`
`1491`	`1502`	`self.length -= 1;`
`1492`	`1503`	`let kv = unsafe { self.range.deallocating_next_unchecked() };`
`1493`		`- Some(kv.into_key_val())`
	`1504`	`+ Some(kv.unwrap().into_key_val())`
`1494`	`1505`	`}`
`1495`	`1506`	`}`
`1496`	`1507`
`@@ -1507,7 +1518,7 @@ impl<K, V> DoubleEndedIterator for IntoIter<K, V> {`
`1507`	`1518`	`} else {`
`1508`	`1519`	`self.length -= 1;`
`1509`	`1520`	`let kv = unsafe { self.range.deallocating_next_back_unchecked() };`
`1510`		`- Some(kv.into_key_val())`
	`1521`	`+ Some(kv.unwrap().into_key_val())`
`1511`	`1522`	`}`
`1512`	`1523`	`}`
`1513`	`1524`	`}`