Auto merge of #28497 - apasel422:issue-28493, r=Gankro

When both the key and value types were zero-sized, `BTreeMap` previously
called `heap::allocate` with `size == 0` for leaf nodes, which is
undefined behavior, and jemalloc would attempt to read invalid memory,
crashing the process.

This avoids undefined behavior by allocating enough space to store one
edge in leaf nodes that would otherwise have `size == 0`. Although this
uses extra memory, maps with zero-sized key types that have sensible
implementations of the ordering traits can only contain a single
key-value pair (and therefore only a single leaf node), and maps with
key and value types that are both zero-sized have few uses, if any.

Furthermore, this is a temporary fix that will likely be unnecessary
once the `BTreeMap` implementation is rewritten to use parent pointers.

Closes #28493.

2 files changed, 59 insertions, 1 deletions

diff --git a/src/libcollections/btree/node.rs b/src/libcollections/btree/node.rs
index d8f8ca6eae5..f5088bf4646 100644
--- a/src/libcollections/btree/node.rs
+++ b/src/libcollections/btree/node.rs
@@ -164,7 +164,12 @@ fn calculate_allocation_generic<K, V>(capacity: usize, is_leaf: bool) -> (usize,
     let (keys_size, keys_align) = (capacity * mem::size_of::<K>(), mem::align_of::<K>());
     let (vals_size, vals_align) = (capacity * mem::size_of::<V>(), mem::align_of::<V>());
     let (edges_size, edges_align) = if is_leaf {
-        (0, 1)
+        // allocate one edge to ensure that we don't pass size 0 to `heap::allocate`
+        if mem::size_of::<K>() == 0 && mem::size_of::<V>() == 0 {
+            (1, mem::align_of::<Node<K, V>>())
+        } else {
+            (0, 1)
+        }
     } else {
         ((capacity + 1) * mem::size_of::<Node<K, V>>(), mem::align_of::<Node<K, V>>())
     };
diff --git a/src/libcollectionstest/btree/map.rs b/src/libcollectionstest/btree/map.rs
index 62b46433da9..846353cc4e7 100644
--- a/src/libcollectionstest/btree/map.rs
+++ b/src/libcollectionstest/btree/map.rs
@@ -294,6 +294,59 @@ fn test_extend_ref() {
     assert_eq!(a[&3], "three");
 }
 
+#[test]
+fn test_zst() {
+    let mut m = BTreeMap::new();
+    assert_eq!(m.len(), 0);
+
+    assert_eq!(m.insert((), ()), None);
+    assert_eq!(m.len(), 1);
+
+    assert_eq!(m.insert((), ()), Some(()));
+    assert_eq!(m.len(), 1);
+    assert_eq!(m.iter().count(), 1);
+
+    m.clear();
+    assert_eq!(m.len(), 0);
+
+    for _ in 0..100 {
+        m.insert((), ());
+    }
+
+    assert_eq!(m.len(), 1);
+    assert_eq!(m.iter().count(), 1);
+}
+
+// This test's only purpose is to ensure that zero-sized keys with nonsensical orderings
+// do not cause segfaults when used with zero-sized values. All other map behavior is
+// undefined.
+#[test]
+fn test_bad_zst() {
+    use std::cmp::Ordering;
+
+    struct Bad;
+
+    impl PartialEq for Bad {
+        fn eq(&self, _: &Self) -> bool { false }
+    }
+
+    impl Eq for Bad {}
+
+    impl PartialOrd for Bad {
+        fn partial_cmp(&self, _: &Self) -> Option<Ordering> { Some(Ordering::Less) }
+    }
+
+    impl Ord for Bad {
+        fn cmp(&self, _: &Self) -> Ordering { Ordering::Less }
+    }
+
+    let mut m = BTreeMap::new();
+
+    for _ in 0..100 {
+        m.insert(Bad, Bad);
+    }
+}
+
 mod bench {
     use std::collections::BTreeMap;
     use std::__rand::{Rng, thread_rng};