From 0a4dc54fb940faf2a6986c97ac1dc75adc1c7790 Mon Sep 17 00:00:00 2001
From: John Platts <john_platts@hotmail.com>
Date: Mon, 19 Sep 2022 14:17:37 -0500
Subject: [PATCH] Update put and ensureCapacity routines

---
 .../src/main/cpp/HWNDMap.cpp                  | 39 ++++++++++++++-----
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/flatlaf-natives/flatlaf-natives-windows/src/main/cpp/HWNDMap.cpp b/flatlaf-natives/flatlaf-natives-windows/src/main/cpp/HWNDMap.cpp
index 8f9ae24a..cf9e2822 100644
--- a/flatlaf-natives/flatlaf-natives-windows/src/main/cpp/HWNDMap.cpp
+++ b/flatlaf-natives/flatlaf-natives-windows/src/main/cpp/HWNDMap.cpp
@@ -18,7 +18,10 @@
 #define _NO_CRT_STDIO_INLINE
 
 #include <stdio.h>
+#include <stdint.h>
+#include <limits.h>
 #include "HWNDMap.h"
+#include "AllocRoutines.h"
 
 #define DEFAULT_CAPACITY		20
 #define INCREASE_CAPACITY		10
@@ -58,7 +61,7 @@ LPVOID HWNDMap::get( HWND key ) {
 	return (index >= 0) ? table[index].value : NULL;
 }
 
-void HWNDMap::put( HWND key, LPVOID value ) {
+bool HWNDMap::put( HWND key, LPVOID value ) {
 	LOCK lock( &criticalSection );
 
 	int index = binarySearch( key );
@@ -66,9 +69,11 @@ void HWNDMap::put( HWND key, LPVOID value ) {
 	if( index >= 0 ) {
 		// key already in map --> replace
 		table[index].value = value;
+		return true;
 	} else {
 		// insert new key
-		ensureCapacity( size + 1 );
+		if(size == INT_MAX || !ensureCapacity( size + 1 ))
+			return false;
 
 		// make roor for new entry
 		index = -(index + 1);
@@ -79,6 +84,7 @@ void HWNDMap::put( HWND key, LPVOID value ) {
 		// insert entry
 		table[index].key = key;
 		table[index].value = value;
+		return true;
 	}
 
 //	dump( "put" );
@@ -121,23 +127,38 @@ int HWNDMap::binarySearch( HWND key ) {
 	return -(low + 1);
 }
 
-void HWNDMap::ensureCapacity( int minCapacity ) {
-	if( minCapacity <= capacity )
-		return;
+static constexpr size_t MaxEntryArrayLength = (SIZE_MAX >> 1) / sizeof(Entry);
+static_assert(MaxEntryArrayLength > 0, "MaxEntryArrayLength > 0 must be true");
+
+static constexpr int MaxEntryArrayIntCapacity =
+	(MaxEntryArrayLength <= INT_MAX) ? static_cast<int>(MaxEntryArrayLength) : INT_MAX;
+static_assert(MaxEntryArrayIntCapacity > 0, "MaxEntryArrayIntCapacity > 0 must be true");
+
+bool HWNDMap::ensureCapacity( int minCapacity ) {
+	if(minCapacity <= capacity)
+		return true;
+	if(minCapacity > MaxEntryArrayIntCapacity)
+		return false;
 
 	// allocate new table
-	int newCapacity = minCapacity + INCREASE_CAPACITY;
-	Entry* newTable = new Entry[newCapacity];
+	unsigned newCapacity = static_cast<unsigned>(minCapacity) + INCREASE_CAPACITY;
+	if(newCapacity > MaxEntryArrayIntCapacity)
+		newCapacity = MaxEntryArrayIntCapacity;
+	
+	Entry* newTable = new (FlatLafNoThrow) Entry[newCapacity];
+	if(newTable == NULL)
+		return false;
 
 	// copy old table to new table
 	for( int i = 0; i < capacity; i++ )
 		newTable[i] = table[i];
 
 	// delete old table
-	delete table;
+	FlatLafWin32ProcessHeapFree(table);
 
 	table = newTable;
-	capacity = newCapacity;
+	capacity = static_cast<int>(newCapacity);
+	return true;
 }
 
 /*