Clean up or comment every unwrap in BTreeMap's main code.

ssomers · ssomers · commit ca253cab3689 · 2020-07-16T12:53:01.000+02:00
diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
@@ -151,7 +151,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
                     let mut out_tree = BTreeMap { root: Some(node::Root::new_leaf()), length: 0 };
 
                     {
-                        let root = out_tree.root.as_mut().unwrap();
+                        let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped
                         let mut out_node = match root.as_mut().force() {
                             Leaf(leaf) => leaf,
                             Internal(_) => unreachable!(),
@@ -171,14 +171,10 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
                 }
                 Internal(internal) => {
                     let mut out_tree = clone_subtree(internal.first_edge().descend());
-                    out_tree.ensure_root_is_owned();
 
                     {
-                        // Ideally we'd use the return of ensure_root_is_owned
-                        // instead of re-unwrapping here but unfortunately that
-                        // borrows all of out_tree and we need access to the
-                        // length below.
-                        let mut out_node = out_tree.root.as_mut().unwrap().push_level();
+                        let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
+                        let mut out_node = out_root.push_level();
                         let mut in_edge = internal.first_edge();
                         while let Ok(kv) = in_edge.right_kv() {
                             let (k, v) = kv.into_kv();
@@ -212,7 +208,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
             // Ord` constraint, which this method lacks.
             BTreeMap { root: None, length: 0 }
         } else {
-            clone_subtree(self.root.as_ref().unwrap().as_ref())
+            clone_subtree(self.root.as_ref().unwrap().as_ref()) // unwrap succeeds because not empty
         }
     }
 }
@@ -243,8 +239,8 @@ where
     }
 
     fn replace(&mut self, key: K) -> Option<K> {
-        self.ensure_root_is_owned();
-        match search::search_tree::<marker::Mut<'_>, K, (), K>(self.root.as_mut()?.as_mut(), &key) {
+        let root = Self::ensure_is_owned(&mut self.root);
+        match search::search_tree::<marker::Mut<'_>, K, (), K>(root.as_mut(), &key) {
             Found(handle) => Some(mem::replace(handle.into_kv_mut().0, key)),
             GoDown(handle) => {
                 VacantEntry { key, handle, length: &mut self.length, _marker: PhantomData }
@@ -943,7 +939,6 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
         // Second, we build a tree from the sorted sequence in linear time.
         self.from_sorted_iter(iter);
-        self.fix_right_edge();
     }
 
     /// Constructs a double-ended iterator over a sub-range of elements in the map.
@@ -1058,8 +1053,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
         // FIXME(@porglezomp) Avoid allocating if we don't insert
-        self.ensure_root_is_owned();
-        match search::search_tree(self.root.as_mut().unwrap().as_mut(), &key) {
+        let root = Self::ensure_is_owned(&mut self.root);
+        match search::search_tree(root.as_mut(), &key) {
             Found(handle) => {
                 Occupied(OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData })
             }
@@ -1070,8 +1065,8 @@ impl<K: Ord, V> BTreeMap<K, V> {
     }
 
     fn from_sorted_iter<I: Iterator<Item = (K, V)>>(&mut self, iter: I) {
-        self.ensure_root_is_owned();
-        let mut cur_node = self.root.as_mut().unwrap().as_mut().last_leaf_edge().into_node();
+        let root = Self::ensure_is_owned(&mut self.root);
+        let mut cur_node = root.as_mut().last_leaf_edge().into_node();
         // Iterate through all key-value pairs, pushing them into nodes at the right level.
         for (key, value) in iter {
             // Try to push key-value pair into the current leaf node.
@@ -1116,11 +1111,12 @@ impl<K: Ord, V> BTreeMap<K, V> {
 
             self.length += 1;
         }
+        Self::fix_right_edge(root)
     }
 
-    fn fix_right_edge(&mut self) {
+    fn fix_right_edge(root: &mut node::Root<K, V>) {
         // Handle underfull nodes, start from the top.
-        let mut cur_node = self.root.as_mut().unwrap().as_mut();
+        let mut cur_node = root.as_mut();
         while let Internal(internal) = cur_node.force() {
             // Check if right-most child is underfull.
             let mut last_edge = internal.last_edge();
@@ -1179,16 +1175,17 @@ impl<K: Ord, V> BTreeMap<K, V> {
         }
 
         let total_num = self.len();
+        let left_root = self.root.as_mut().unwrap(); // unwrap succeeds because not empty
 
         let mut right = Self::new();
-        let right_root = right.ensure_root_is_owned();
-        for _ in 0..(self.root.as_ref().unwrap().as_ref().height()) {
+        let right_root = Self::ensure_is_owned(&mut right.root);
+        for _ in 0..left_root.height() {
             right_root.push_level();
         }
 
         {
-            let mut left_node = self.root.as_mut().unwrap().as_mut();
-            let mut right_node = right.root.as_mut().unwrap().as_mut();
+            let mut left_node = left_root.as_mut();
+            let mut right_node = right_root.as_mut();
 
             loop {
                 let mut split_edge = match search::search_node(left_node, key) {
@@ -1214,12 +1211,10 @@ impl<K: Ord, V> BTreeMap<K, V> {
             }
         }
 
-        self.fix_right_border();
-        right.fix_left_border();
+        Self::fix_right_border(left_root);
+        Self::fix_left_border(right_root);
 
-        if self.root.as_ref().unwrap().as_ref().height()
-            < right.root.as_ref().unwrap().as_ref().height()
-        {
+        if left_root.height() < right_root.height() {
             self.recalc_length();
             right.length = total_num - self.len();
         } else {
@@ -1303,23 +1298,17 @@ impl<K: Ord, V> BTreeMap<K, V> {
     }
 
     /// Removes empty levels on the top.
-    fn fix_top(&mut self) {
-        loop {
-            {
-                let node = self.root.as_ref().unwrap().as_ref();
-                if node.height() == 0 || node.len() > 0 {
-                    break;
-                }
-            }
-            self.root.as_mut().unwrap().pop_level();
+    fn fix_top(root: &mut node::Root<K, V>) {
+        while root.height() > 0 && root.as_ref().len() == 0 {
+            root.pop_level();
         }
     }
 
-    fn fix_right_border(&mut self) {
-        self.fix_top();
+    fn fix_right_border(root: &mut node::Root<K, V>) {
+        Self::fix_top(root);
 
         {
-            let mut cur_node = self.root.as_mut().unwrap().as_mut();
+            let mut cur_node = root.as_mut();
 
             while let Internal(node) = cur_node.force() {
                 let mut last_kv = node.last_kv();
@@ -1337,15 +1326,15 @@ impl<K: Ord, V> BTreeMap<K, V> {
             }
         }
 
-        self.fix_top();
+        Self::fix_top(root);
     }
 
     /// The symmetric clone of `fix_right_border`.
-    fn fix_left_border(&mut self) {
-        self.fix_top();
+    fn fix_left_border(root: &mut node::Root<K, V>) {
+        Self::fix_top(root);
 
         {
-            let mut cur_node = self.root.as_mut().unwrap().as_mut();
+            let mut cur_node = root.as_mut();
 
             while let Internal(node) = cur_node.force() {
                 let mut first_kv = node.first_kv();
@@ -1362,7 +1351,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
             }
         }
 
-        self.fix_top();
+        Self::fix_top(root);
     }
 }
 
@@ -2321,9 +2310,9 @@ impl<K, V> BTreeMap<K, V> {
     }
 
     /// If the root node is the empty (non-allocated) root node, allocate our
-    /// own node.
-    fn ensure_root_is_owned(&mut self) -> &mut node::Root<K, V> {
-        self.root.get_or_insert_with(node::Root::new_leaf)
+    /// own node. Is an associated function to avoid borrowing the entire BTreeMap.
+    fn ensure_is_owned(root: &mut Option<node::Root<K, V>>) -> &mut node::Root<K, V> {
+        root.get_or_insert_with(node::Root::new_leaf)
     }
 }
 
diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs
@@ -153,6 +153,11 @@ unsafe impl<K: Sync, V: Sync> Sync for Root<K, V> {}
 unsafe impl<K: Send, V: Send> Send for Root<K, V> {}
 
 impl<K, V> Root<K, V> {
+    /// Returns the number of levels below the root.
+    pub fn height(&self) -> usize {
+        self.height
+    }
+
     /// Returns a new owned tree, with its own root node that is initially empty.
     pub fn new_leaf() -> Self {
         Root { node: BoxedNode::from_leaf(Box::new(unsafe { LeafNode::new() })), height: 0 }

Original file line number	Diff line number	Diff line change
`@@ -151,7 +151,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {`
`151`	`151`	`let mut out_tree = BTreeMap { root: Some(node::Root::new_leaf()), length: 0 };`
`152`	`152`
`153`	`153`	`{`
`154`		`- let root = out_tree.root.as_mut().unwrap();`
	`154`	`+ let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped`
`155`	`155`	`let mut out_node = match root.as_mut().force() {`
`156`	`156`	`Leaf(leaf) => leaf,`
`157`	`157`	`Internal(_) => unreachable!(),`
`@@ -171,14 +171,10 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {`
`171`	`171`	`}`
`172`	`172`	`Internal(internal) => {`
`173`	`173`	`let mut out_tree = clone_subtree(internal.first_edge().descend());`
`174`		`- out_tree.ensure_root_is_owned();`
`175`	`174`
`176`	`175`	`{`
`177`		`- // Ideally we'd use the return of ensure_root_is_owned`
`178`		`- // instead of re-unwrapping here but unfortunately that`
`179`		`- // borrows all of out_tree and we need access to the`
`180`		`- // length below.`
`181`		`- let mut out_node = out_tree.root.as_mut().unwrap().push_level();`
	`176`	`+ let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);`
	`177`	`+ let mut out_node = out_root.push_level();`
`182`	`178`	`let mut in_edge = internal.first_edge();`
`183`	`179`	`while let Ok(kv) = in_edge.right_kv() {`
`184`	`180`	`let (k, v) = kv.into_kv();`
`@@ -212,7 +208,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {`
`212`	`208`	// Ord` constraint, which this method lacks.
`213`	`209`	`BTreeMap { root: None, length: 0 }`
`214`	`210`	`} else {`
`215`		`- clone_subtree(self.root.as_ref().unwrap().as_ref())`
	`211`	`+ clone_subtree(self.root.as_ref().unwrap().as_ref()) // unwrap succeeds because not empty`
`216`	`212`	`}`
`217`	`213`	`}`
`218`	`214`	`}`
`@@ -243,8 +239,8 @@ where`
`243`	`239`	`}`
`244`	`240`
`245`	`241`	`fn replace(&mut self, key: K) -> Option<K> {`
`246`		`- self.ensure_root_is_owned();`
`247`		`- match search::search_tree::<marker::Mut<'_>, K, (), K>(self.root.as_mut()?.as_mut(), &key) {`
	`242`	`+ let root = Self::ensure_is_owned(&mut self.root);`
	`243`	`+ match search::search_tree::<marker::Mut<'_>, K, (), K>(root.as_mut(), &key) {`
`248`	`244`	`Found(handle) => Some(mem::replace(handle.into_kv_mut().0, key)),`
`249`	`245`	`GoDown(handle) => {`
`250`	`246`	`VacantEntry { key, handle, length: &mut self.length, _marker: PhantomData }`
`@@ -943,7 +939,6 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`943`	`939`
`944`	`940`	`// Second, we build a tree from the sorted sequence in linear time.`
`945`	`941`	`self.from_sorted_iter(iter);`
`946`		`- self.fix_right_edge();`
`947`	`942`	`}`
`948`	`943`
`949`	`944`	`/// Constructs a double-ended iterator over a sub-range of elements in the map.`
`@@ -1058,8 +1053,8 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`1058`	`1053`	`#[stable(feature = "rust1", since = "1.0.0")]`
`1059`	`1054`	`pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {`
`1060`	`1055`	`// FIXME(@porglezomp) Avoid allocating if we don't insert`
`1061`		`- self.ensure_root_is_owned();`
`1062`		`- match search::search_tree(self.root.as_mut().unwrap().as_mut(), &key) {`
	`1056`	`+ let root = Self::ensure_is_owned(&mut self.root);`
	`1057`	`+ match search::search_tree(root.as_mut(), &key) {`
`1063`	`1058`	`Found(handle) => {`
`1064`	`1059`	`Occupied(OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData })`
`1065`	`1060`	`}`
`@@ -1070,8 +1065,8 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`1070`	`1065`	`}`
`1071`	`1066`
`1072`	`1067`	`fn from_sorted_iter<I: Iterator<Item = (K, V)>>(&mut self, iter: I) {`
`1073`		`- self.ensure_root_is_owned();`
`1074`		`- let mut cur_node = self.root.as_mut().unwrap().as_mut().last_leaf_edge().into_node();`
	`1068`	`+ let root = Self::ensure_is_owned(&mut self.root);`
	`1069`	`+ let mut cur_node = root.as_mut().last_leaf_edge().into_node();`
`1075`	`1070`	`// Iterate through all key-value pairs, pushing them into nodes at the right level.`
`1076`	`1071`	`for (key, value) in iter {`
`1077`	`1072`	`// Try to push key-value pair into the current leaf node.`
`@@ -1116,11 +1111,12 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`1116`	`1111`
`1117`	`1112`	`self.length += 1;`
`1118`	`1113`	`}`
	`1114`	`+ Self::fix_right_edge(root)`
`1119`	`1115`	`}`
`1120`	`1116`
`1121`		`- fn fix_right_edge(&mut self) {`
	`1117`	`+ fn fix_right_edge(root: &mut node::Root<K, V>) {`
`1122`	`1118`	`// Handle underfull nodes, start from the top.`
`1123`		`- let mut cur_node = self.root.as_mut().unwrap().as_mut();`
	`1119`	`+ let mut cur_node = root.as_mut();`
`1124`	`1120`	`while let Internal(internal) = cur_node.force() {`
`1125`	`1121`	`// Check if right-most child is underfull.`
`1126`	`1122`	`let mut last_edge = internal.last_edge();`
`@@ -1179,16 +1175,17 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`1179`	`1175`	`}`
`1180`	`1176`
`1181`	`1177`	`let total_num = self.len();`
	`1178`	`+ let left_root = self.root.as_mut().unwrap(); // unwrap succeeds because not empty`
`1182`	`1179`
`1183`	`1180`	`let mut right = Self::new();`
`1184`		`- let right_root = right.ensure_root_is_owned();`
`1185`		`- for _ in 0..(self.root.as_ref().unwrap().as_ref().height()) {`
	`1181`	`+ let right_root = Self::ensure_is_owned(&mut right.root);`
	`1182`	`+ for _ in 0..left_root.height() {`
`1186`	`1183`	`right_root.push_level();`
`1187`	`1184`	`}`
`1188`	`1185`
`1189`	`1186`	`{`
`1190`		`- let mut left_node = self.root.as_mut().unwrap().as_mut();`
`1191`		`- let mut right_node = right.root.as_mut().unwrap().as_mut();`
	`1187`	`+ let mut left_node = left_root.as_mut();`
	`1188`	`+ let mut right_node = right_root.as_mut();`
`1192`	`1189`
`1193`	`1190`	`loop {`
`1194`	`1191`	`let mut split_edge = match search::search_node(left_node, key) {`
`@@ -1214,12 +1211,10 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`1214`	`1211`	`}`
`1215`	`1212`	`}`
`1216`	`1213`
`1217`		`- self.fix_right_border();`
`1218`		`- right.fix_left_border();`
	`1214`	`+ Self::fix_right_border(left_root);`
	`1215`	`+ Self::fix_left_border(right_root);`
`1219`	`1216`
`1220`		`- if self.root.as_ref().unwrap().as_ref().height()`
`1221`		`- < right.root.as_ref().unwrap().as_ref().height()`
`1222`		`- {`
	`1217`	`+ if left_root.height() < right_root.height() {`
`1223`	`1218`	`self.recalc_length();`
`1224`	`1219`	`right.length = total_num - self.len();`
`1225`	`1220`	`} else {`
`@@ -1303,23 +1298,17 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`1303`	`1298`	`}`
`1304`	`1299`
`1305`	`1300`	`/// Removes empty levels on the top.`
`1306`		`- fn fix_top(&mut self) {`
`1307`		`- loop {`
`1308`		`- {`
`1309`		`- let node = self.root.as_ref().unwrap().as_ref();`
`1310`		`- if node.height() == 0 \|\| node.len() > 0 {`
`1311`		`- break;`
`1312`		`- }`
`1313`		`- }`
`1314`		`- self.root.as_mut().unwrap().pop_level();`
	`1301`	`+ fn fix_top(root: &mut node::Root<K, V>) {`
	`1302`	`+ while root.height() > 0 && root.as_ref().len() == 0 {`
	`1303`	`+ root.pop_level();`
`1315`	`1304`	`}`
`1316`	`1305`	`}`
`1317`	`1306`
`1318`		`- fn fix_right_border(&mut self) {`
`1319`		`- self.fix_top();`
	`1307`	`+ fn fix_right_border(root: &mut node::Root<K, V>) {`
	`1308`	`+ Self::fix_top(root);`
`1320`	`1309`
`1321`	`1310`	`{`
`1322`		`- let mut cur_node = self.root.as_mut().unwrap().as_mut();`
	`1311`	`+ let mut cur_node = root.as_mut();`
`1323`	`1312`
`1324`	`1313`	`while let Internal(node) = cur_node.force() {`
`1325`	`1314`	`let mut last_kv = node.last_kv();`
`@@ -1337,15 +1326,15 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`1337`	`1326`	`}`
`1338`	`1327`	`}`
`1339`	`1328`
`1340`		`- self.fix_top();`
	`1329`	`+ Self::fix_top(root);`
`1341`	`1330`	`}`
`1342`	`1331`
`1343`	`1332`	/// The symmetric clone of `fix_right_border`.
`1344`		`- fn fix_left_border(&mut self) {`
`1345`		`- self.fix_top();`
	`1333`	`+ fn fix_left_border(root: &mut node::Root<K, V>) {`
	`1334`	`+ Self::fix_top(root);`
`1346`	`1335`
`1347`	`1336`	`{`
`1348`		`- let mut cur_node = self.root.as_mut().unwrap().as_mut();`
	`1337`	`+ let mut cur_node = root.as_mut();`
`1349`	`1338`
`1350`	`1339`	`while let Internal(node) = cur_node.force() {`
`1351`	`1340`	`let mut first_kv = node.first_kv();`
`@@ -1362,7 +1351,7 @@ impl<K: Ord, V> BTreeMap<K, V> {`
`1362`	`1351`	`}`
`1363`	`1352`	`}`
`1364`	`1353`
`1365`		`- self.fix_top();`
	`1354`	`+ Self::fix_top(root);`
`1366`	`1355`	`}`
`1367`	`1356`	`}`
`1368`	`1357`
`@@ -2321,9 +2310,9 @@ impl<K, V> BTreeMap<K, V> {`
`2321`	`2310`	`}`
`2322`	`2311`
`2323`	`2312`	`/// If the root node is the empty (non-allocated) root node, allocate our`
`2324`		`- /// own node.`
`2325`		`- fn ensure_root_is_owned(&mut self) -> &mut node::Root<K, V> {`
`2326`		`- self.root.get_or_insert_with(node::Root::new_leaf)`
	`2313`	`+ /// own node. Is an associated function to avoid borrowing the entire BTreeMap.`
	`2314`	`+ fn ensure_is_owned(root: &mut Option<node::Root<K, V>>) -> &mut node::Root<K, V> {`
	`2315`	`+ root.get_or_insert_with(node::Root::new_leaf)`
`2327`	`2316`	`}`
`2328`	`2317`	`}`
`2329`	`2318`