From 68845bd24d843670cd45e6decc60d01ac7266f40 Mon Sep 17 00:00:00 2001
From: kawaemon <me@kawaemon.dev>
Date: Sat, 11 Jan 2025 03:26:52 +0900
Subject: [PATCH] fix: HeaderMap::reserve allocates insufficient capacity
 (#741)

This bug caused additional allocation when attempted to insert
the requested number of entries.
This commit fix that by converting capacity to raw capacity before
allocation.
---
 src/header/map.rs   | 16 +++++++++-------
 tests/header_map.rs | 24 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/header/map.rs b/src/header/map.rs
index 07b4554a..ebbc5937 100644
--- a/src/header/map.rs
+++ b/src/header/map.rs
@@ -707,20 +707,22 @@ impl<T> HeaderMap<T> {
             .checked_add(additional)
             .ok_or_else(MaxSizeReached::new)?;
 
-        if cap > self.indices.len() {
-            let cap = cap
+        let raw_cap = to_raw_capacity(cap);
+
+        if raw_cap > self.indices.len() {
+            let raw_cap = raw_cap
                 .checked_next_power_of_two()
                 .ok_or_else(MaxSizeReached::new)?;
-            if cap > MAX_SIZE {
+            if raw_cap > MAX_SIZE {
                 return Err(MaxSizeReached::new());
             }
 
             if self.entries.is_empty() {
-                self.mask = cap as Size - 1;
-                self.indices = vec![Pos::none(); cap].into_boxed_slice();
-                self.entries = Vec::with_capacity(usable_capacity(cap));
+                self.mask = raw_cap as Size - 1;
+                self.indices = vec![Pos::none(); raw_cap].into_boxed_slice();
+                self.entries = Vec::with_capacity(usable_capacity(raw_cap));
             } else {
-                self.try_grow(cap)?;
+                self.try_grow(raw_cap)?;
             }
         }
 
diff --git a/tests/header_map.rs b/tests/header_map.rs
index 9859b0a8..9a9d7e12 100644
--- a/tests/header_map.rs
+++ b/tests/header_map.rs
@@ -63,6 +63,30 @@ fn reserve_overflow() {
     headers.reserve(std::usize::MAX); // next_power_of_two overflows
 }
 
+#[test]
+fn reserve() {
+    let mut headers = HeaderMap::<usize>::default();
+    assert_eq!(headers.capacity(), 0);
+
+    let requested_cap = 8;
+    headers.reserve(requested_cap);
+
+    let reserved_cap = headers.capacity();
+    assert!(
+        reserved_cap >= requested_cap,
+        "requested {} capacity, but it reserved only {} entries",
+        requested_cap,
+        reserved_cap,
+    );
+
+    for i in 0..requested_cap {
+        let name = format!("h{i}").parse::<HeaderName>().unwrap();
+        headers.insert(name, i);
+    }
+
+    assert_eq!(headers.capacity(), reserved_cap, "unexpected reallocation");
+}
+
 #[test]
 fn drain() {
     let mut headers = HeaderMap::new();