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.

Author: bors

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>>())
     };

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};