From 552818607ad6f36b954c2d8b00f7d403210ed759 Mon Sep 17 00:00:00 2001
From: junderw <jonathan.underwood4649@gmail.com>
Date: Wed, 28 Jun 2023 12:42:33 -0700
Subject: [PATCH] Better initial capacity

---
 backend/rust-gbt/src/gbt.rs | 15 +++++++++------
 backend/rust-gbt/src/lib.rs |  7 ++-----
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/backend/rust-gbt/src/gbt.rs b/backend/rust-gbt/src/gbt.rs
index 1899650d9..bab19ac89 100644
--- a/backend/rust-gbt/src/gbt.rs
+++ b/backend/rust-gbt/src/gbt.rs
@@ -10,7 +10,7 @@ use crate::{
     u32_hasher_types::{
         u32hashmap_with_capacity, u32hashset_new, u32priority_queue_with_capacity, U32HasherState,
     },
-    GbtResult, ThreadTransactionsMap, STARTING_CAPACITY,
+    GbtResult, ThreadTransactionsMap,
 };
 
 const MAX_BLOCK_WEIGHT_UNITS: u32 = 4_000_000 - 4_000;
@@ -59,8 +59,9 @@ impl Ord for TxPriority {
 #[allow(clippy::too_many_lines)]
 #[allow(clippy::cognitive_complexity)]
 pub fn gbt(mempool: &mut ThreadTransactionsMap) -> GbtResult {
-    let mut audit_pool: AuditPool = u32hashmap_with_capacity(STARTING_CAPACITY);
-    let mut mempool_stack: Vec<u32> = Vec::with_capacity(STARTING_CAPACITY);
+    let mempool_len = mempool.len();
+    let mut audit_pool: AuditPool = u32hashmap_with_capacity(mempool_len);
+    let mut mempool_stack: Vec<u32> = Vec::with_capacity(mempool_len);
     let mut clusters: Vec<Vec<u32>> = Vec::new();
     let mut block_weights: Vec<u32> = Vec::new();
 
@@ -96,8 +97,10 @@ pub fn gbt(mempool: &mut ThreadTransactionsMap) -> GbtResult {
     let mut blocks: Vec<Vec<u32>> = Vec::new();
     let mut block_weight: u32 = BLOCK_RESERVED_WEIGHT;
     let mut block_sigops: u32 = BLOCK_RESERVED_SIGOPS;
-    let mut transactions: Vec<u32> = Vec::with_capacity(STARTING_CAPACITY);
-    let mut modified: ModifiedQueue = u32priority_queue_with_capacity(STARTING_CAPACITY);
+    // No need to be bigger than 4096 transactions for the per-block transaction Vec.
+    let initial_txes_per_block: usize = 4096.min(mempool_len);
+    let mut transactions: Vec<u32> = Vec::with_capacity(initial_txes_per_block);
+    let mut modified: ModifiedQueue = u32priority_queue_with_capacity(mempool_len);
     let mut overflow: Vec<u32> = Vec::new();
     let mut failures = 0;
     while !mempool_stack.is_empty() || !modified.is_empty() {
@@ -196,7 +199,7 @@ pub fn gbt(mempool: &mut ThreadTransactionsMap) -> GbtResult {
                 block_weights.push(block_weight);
             }
             // reset for the next block
-            transactions = Vec::with_capacity(STARTING_CAPACITY);
+            transactions = Vec::with_capacity(initial_txes_per_block);
             block_weight = BLOCK_RESERVED_WEIGHT;
             block_sigops = BLOCK_RESERVED_SIGOPS;
             failures = 0;
diff --git a/backend/rust-gbt/src/lib.rs b/backend/rust-gbt/src/lib.rs
index 6a0e3ea5e..0f7c0f7ba 100644
--- a/backend/rust-gbt/src/lib.rs
+++ b/backend/rust-gbt/src/lib.rs
@@ -23,13 +23,10 @@ mod u32_hasher_types;
 
 use u32_hasher_types::{u32hashmap_with_capacity, U32HasherState};
 
-/// This is the starting capacity for HashMap/Vec/etc. that deal with transactions.
-/// `HashMap` doubles capacity when it hits it, so 2048 is a decent tradeoff between
-/// not wasting too much memory when it's below this, and allowing for less re-allocations
-/// by virtue of starting with such a large capacity.
+/// This is the initial capacity of the GbtGenerator struct's inner HashMap.
 ///
 /// Note: This doesn't *have* to be a power of 2. (uwu)
-const STARTING_CAPACITY: usize = 32768;
+const STARTING_CAPACITY: usize = 1_048_576;
 
 type ThreadTransactionsMap = HashMap<u32, ThreadTransaction, U32HasherState>;