From 1befb9ba6fa6f38ef127cb77f3617b4de13c4444 Mon Sep 17 00:00:00 2001
From: Tim Wojtulewicz <timwoj@gmail.com>
Date: Wed, 31 Aug 2022 20:48:18 +0000
Subject: [PATCH] Dict: Reduce some math required in repeatedly-called
 functions

---
 src/Dict.h | 57 ++++++++++++++++++++++++++++--------------------------
 1 file changed, 30 insertions(+), 27 deletions(-)

diff --git a/src/Dict.h b/src/Dict.h
index 36f5f68667..c3fe5a893b 100644
--- a/src/Dict.h
+++ b/src/Dict.h
@@ -457,7 +457,7 @@ public:
 			{
 			// If an initial size is speicified, init the table right away. Otherwise wait until the
 			// first insertion to init.
-			log2_buckets = static_cast<unsigned char>(std::log2(initial_size));
+			SetLog2Buckets(static_cast<uint16_t>(std::log2(initial_size)));
 			Init();
 			}
 
@@ -727,13 +727,8 @@ public:
 		}
 
 	/// The capacity of the table, Buckets + Overflow Size.
-	int Capacity(bool expected = false) const
-		{
-		int capacity = (1 << log2_buckets) + (log2_buckets + 0);
-		if ( expected )
-			return capacity;
-		return table ? capacity : 0;
-		}
+	int Capacity() const { return table ? bucket_capacity : 0; }
+	int ExpectedCapacity() const { return bucket_capacity; }
 
 	// Debugging
 #define DUMPIF(f)                                                                                  \
@@ -932,17 +927,18 @@ private:
 	friend zeek::DictIterator<T>;
 	friend zeek::RobustDictIterator<T>;
 
-	/// Buckets of the table, not including overflow size.
-	int Buckets(bool expected = false) const
+	void SetLog2Buckets(int value)
 		{
-		int buckets = (1 << log2_buckets);
-		if ( expected )
-			return buckets;
-		return table ? buckets : 0;
+		log2_buckets = value;
+		bucket_count = 1 << log2_buckets;
+		bucket_capacity = (1 << log2_buckets) + log2_buckets;
 		}
 
+	/// Buckets of the table, not including overflow size.
+	int Buckets() const { return table ? bucket_count : 0; }
+
 	// bucket math
-	int ThresholdEntries() const
+	uint32_t ThresholdEntries() const
 		{
 		// Increase the size of the dictionary when it is 75% full. However, when the dictionary
 		// is small ( <= 20 elements ), only resize it when it's 100% full. The dictionary will
@@ -998,7 +994,8 @@ private:
 		{
 		ASSERT(bucket >= 0 && bucket < Buckets());
 		int i = bucket;
-		while ( i < Capacity() && ! table[i].Empty() && BucketByPosition(i) <= bucket )
+		int current_cap = Capacity();
+		while ( i < current_cap && ! table[i].Empty() && BucketByPosition(i) <= bucket )
 			i++;
 		return i;
 		}
@@ -1025,7 +1022,8 @@ private:
 
 		int bucket = BucketByPosition(position);
 		int i = position;
-		while ( i < Capacity() && ! table[i].Empty() && BucketByPosition(i) == bucket )
+		int current_cap = Capacity();
+		while ( i < current_cap && ! table[i].Empty() && BucketByPosition(i) == bucket )
 			i++; // stop just over the tail.
 
 		return i - 1;
@@ -1050,10 +1048,11 @@ private:
 		{
 		ASSERT(table && -1 <= position && position < Capacity());
 
+		int current_cap = Capacity();
 		do
 			{
 			position++;
-			} while ( position < Capacity() && table[position].Empty() );
+			} while ( position < current_cap && table[position].Empty() );
 
 		return position;
 		}
@@ -1061,7 +1060,7 @@ private:
 	void Init()
 		{
 		ASSERT(! table);
-		table = (detail::DictEntry<T>*)malloc(sizeof(detail::DictEntry<T>) * Capacity(true));
+		table = (detail::DictEntry<T>*)malloc(sizeof(detail::DictEntry<T>) * ExpectedCapacity());
 		for ( int i = Capacity() - 1; i >= 0; i-- )
 			table[i].SetEmpty();
 		}
@@ -1069,7 +1068,8 @@ private:
 	// Lookup
 	int LinearLookupIndex(const void* key, int key_size, detail::hash_t hash) const
 		{
-		for ( int i = 0; i < Capacity(); i++ )
+		auto current_cap = Capacity();
+		for ( int i = 0; i < current_cap; i++ )
 			if ( ! table[i].Empty() && table[i].Equal((const char*)key, key_size, hash) )
 				return i;
 		return -1;
@@ -1381,7 +1381,8 @@ private:
 	void SizeUp()
 		{
 		int prev_capacity = Capacity();
-		log2_buckets++;
+		SetLog2Buckets(log2_buckets + 1);
+
 		int capacity = Capacity();
 		table = (detail::DictEntry<T>*)realloc(table, capacity * sizeof(detail::DictEntry<T>));
 		for ( int i = prev_capacity; i < capacity; i++ )
@@ -1488,18 +1489,20 @@ private:
 	// as it will be remapped to new dict size anyway. however, the missed count is recorded
 	// for lookup. if position not found for a key in the position of dict of current size, it
 	// still could be in the position of dict of previous N sizes.
-	unsigned char remaps = 0;
-	unsigned char log2_buckets = 0;
+	uint16_t remaps = 0;
+	uint16_t log2_buckets = 0;
+	uint32_t bucket_capacity = 1;
+	uint32_t bucket_count = 1;
 
 	// Pending number of iterators on the Dict, including both robust and non-robust.
 	// This is used to avoid remapping if there are any active iterators.
-	unsigned short num_iterators = 0;
+	uint16_t num_iterators = 0;
 
 	// The last index to be remapped.
-	int remap_end = -1;
+	int32_t remap_end = -1;
 
-	int num_entries = 0;
-	int max_entries = 0;
+	uint32_t num_entries = 0;
+	uint32_t max_entries = 0;
 	uint64_t cum_entries = 0;
 
 	dict_delete_func delete_func = nullptr;