From 935e83a49005935b044c0608cb12c37a97e77ef4 Mon Sep 17 00:00:00 2001
From: Elad <18193363+elad335@users.noreply.github.com>
Date: Fri, 15 Nov 2024 19:10:59 +0200
Subject: [PATCH] Utilities\lockless.h: Flatten recusrion in lf_array

Add assert for sane index access.
---
 Utilities/lockless.h | 45 +++++++++++++++++++++++++++++++++-----------
 1 file changed, 34 insertions(+), 11 deletions(-)

diff --git a/Utilities/lockless.h b/Utilities/lockless.h
index 08c8cb2df5..28d5d65246 100644
--- a/Utilities/lockless.h
+++ b/Utilities/lockless.h
@@ -31,22 +31,45 @@ public:
 
 	T& operator [](usz index)
 	{
-		if (index < N) [[likely]]
+		lf_array* _this = this;
+
+		T* result{};
+		bool installed = false;
+
+		for (usz i = 0;; i += N)
 		{
-			return m_data[index];
-		}
-		else if (!m_next) [[unlikely]]
-		{
-			// Create new array block. It's not a full-fledged once-synchronization, unlikely needed.
-			for (auto _new = new lf_array, ptr = this; ptr;)
+			if (index - i < N)
 			{
-				// Install the pointer. If failed, go deeper.
-				ptr = ptr->m_next.compare_and_swap(nullptr, _new);
+				result = std::addressof(m_data[index - i]);
+				break;
 			}
+
+			lf_array* next = m_next;
+
+			if (!next)
+			{
+				// Do not allow access beyond one element more at a time 
+				ensure(!installed && index - i == N);
+
+				installed = true;
+
+				for (auto _new = new lf_array, ptr = _this; ptr;)
+				{
+					// Install the pointer. If failed, go deeper.
+					ptr = ptr->m_next.compare_and_swap(nullptr, _new);
+
+					if (!next)
+					{
+						// Determine the next pointer (if null then the new memory has been installed)
+						next = ptr ? ptr : _new;
+					}
+				}
+			}
+
+			_this = next;
 		}
 
-		// Access recursively
-		return (*m_next)[index - N];
+		return *result;
 	}
 
 	u64 size() const