From 957d4ddd3576d74aadad745ff4be50134206107e Mon Sep 17 00:00:00 2001
From: Ignacio Vera <ignacio.vera@elastic.co>
Date: Wed, 9 Oct 2024 09:22:53 +0200
Subject: [PATCH] Improve performance of Int3Hash#removeAndAdd

---
 .../elasticsearch/common/util/Int3Hash.java   | 33 +++++++++----------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/common/util/Int3Hash.java b/server/src/main/java/org/elasticsearch/common/util/Int3Hash.java
index dc49b39a031a1..c32da943ea02e 100644
--- a/server/src/main/java/org/elasticsearch/common/util/Int3Hash.java
+++ b/server/src/main/java/org/elasticsearch/common/util/Int3Hash.java
@@ -58,8 +58,8 @@ public int getKey3(long id) {
      * Get the id associated with <code>key</code> or -1 if the key is not contained in the hash.
      */
     public long find(int key1, int key2, int key3) {
-        long index = slot(hash(key1, key2, key3), mask);
-        while (true) {
+        final long slot = slot(hash(key1, key2, key3), mask);
+        for (long index = slot;; index = nextSlot(index, mask)) {
             final long id = id(index);
             if (id == -1) {
                 return id;
@@ -69,14 +69,13 @@ public long find(int key1, int key2, int key3) {
                     return id;
                 }
             }
-            index = nextSlot(index, mask);
         }
     }
 
     private long set(int key1, int key2, int key3, long id) {
         assert size < maxSize;
-        long index = slot(hash(key1, key2, key3), mask);
-        while (true) {
+        long slot = slot(hash(key1, key2, key3), mask);
+        for (long index = slot;; index = nextSlot(index, mask)) {
             final long curId = id(index);
             if (curId == -1) { // means unset
                 setId(index, id);
@@ -84,33 +83,35 @@ private long set(int key1, int key2, int key3, long id) {
                 ++size;
                 return id;
             } else {
-                long keyOffset = 3 * curId;
+                final long keyOffset = 3 * curId;
                 if (keys.get(keyOffset) == key1 && keys.get(keyOffset + 1) == key2 && keys.get(keyOffset + 2) == key3) {
                     return -1 - curId;
                 }
             }
-            index = nextSlot(index, mask);
         }
     }
 
     private void append(long id, int key1, int key2, int key3) {
-        long keyOffset = 3 * id;
+        final long keyOffset = 3 * id;
         keys = bigArrays.grow(keys, keyOffset + 3);
         keys.set(keyOffset, key1);
         keys.set(keyOffset + 1, key2);
         keys.set(keyOffset + 2, key3);
     }
 
-    private void reset(int key1, int key2, int key3, long id) {
-        long index = slot(hash(key1, key2, key3), mask);
-        while (true) {
+    private void reset(long id) {
+        final IntArray keys = this.keys;
+        final long keyOffset = id * 3;
+        final int key1 = keys.get(keyOffset);
+        final int key2 = keys.get(keyOffset + 1);
+        final int key3 = keys.get(keyOffset + 2);
+        final long slot = slot(hash(key1, key2, key3), mask);
+        for (long index = slot;; index = nextSlot(index, mask)) {
             final long curId = id(index);
             if (curId == -1) { // means unset
                 setId(index, id);
-                append(id, key1, key2, key3);
                 break;
             }
-            index = nextSlot(index, mask);
         }
     }
 
@@ -132,11 +133,7 @@ public long add(int key1, int key2, int key3) {
     protected void removeAndAdd(long index) {
         final long id = getAndSetId(index, -1);
         assert id >= 0;
-        long keyOffset = id * 3;
-        final int key1 = keys.getAndSet(keyOffset, 0);
-        final int key2 = keys.getAndSet(keyOffset + 1, 0);
-        final int key3 = keys.getAndSet(keyOffset + 2, 0);
-        reset(key1, key2, key3, id);
+        reset(id);
     }
 
     @Override