Add getAndSet to Objectarray (#114200)

This commit adds a getAndSet implementation to the ObjectArray API and changes the set method to return void.

Author: iverase

server/src/main/java/org/elasticsearch/common/util/BigArrays.java

@@ -452,7 +452,13 @@ public T get(long index) {
         }
 
         @Override
-        public T set(long index, T value) {
+        public void set(long index, T value) {
+            assert index >= 0 && index < size();
+            array[(int) index] = value;
+        }
+
+        @Override
+        public T getAndSet(long index, T value) {
             assert index >= 0 && index < size();
             @SuppressWarnings("unchecked")
             T ret = (T) array[(int) index];

server/src/main/java/org/elasticsearch/common/util/BigObjectArray.java

@@ -46,7 +46,15 @@ public T get(long index) {
     }
 
     @Override
-    public T set(long index, T value) {
+    public void set(long index, T value) {
+        final int pageIndex = pageIndex(index);
+        final int indexInPage = indexInPage(index);
+        final Object[] page = pages[pageIndex];
+        page[indexInPage] = value;
+    }
+
+    @Override
+    public T getAndSet(long index, T value) {
         final int pageIndex = pageIndex(index);
         final int indexInPage = indexInPage(index);
         final Object[] page = pages[pageIndex];

server/src/main/java/org/elasticsearch/common/util/LongObjectPagedHashMap.java

@@ -77,7 +77,7 @@ public T put(long key, T value) {
      */
     public T remove(long key) {
         for (long i = slot(hash(key), mask);; i = nextSlot(i, mask)) {
-            final T previous = values.set(i, null);
+            final T previous = values.getAndSet(i, null);
             if (previous == null) {
                 return null;
             } else if (keys.get(i) == key) {
@@ -98,7 +98,7 @@ private T set(long key, T value) {
             throw new IllegalArgumentException("Null values are not supported");
         }
         for (long i = slot(hash(key), mask);; i = nextSlot(i, mask)) {
-            final T previous = values.set(i, value);
+            final T previous = values.getAndSet(i, value);
             if (previous == null) {
                 // slot was free
                 keys.set(i, key);
@@ -180,7 +180,7 @@ protected boolean used(long bucket) {
     @Override
     protected void removeAndAdd(long index) {
         final long key = keys.get(index);
-        final T value = values.set(index, null);
+        final T value = values.getAndSet(index, null);
         --size;
         final T removed = set(key, value);
         assert removed == null;

server/src/main/java/org/elasticsearch/common/util/ObjectArray.java

@@ -19,9 +19,14 @@ public interface ObjectArray<T> extends BigArray {
      */
     T get(long index);
 
+    /**
+     * Set a value at the given index.
+     */
+    void set(long index, T value);
+
     /**
      * Set a value at the given index and return the previous value.
      */
-    T set(long index, T value);
+    T getAndSet(long index, T value);
 
 }

server/src/main/java/org/elasticsearch/common/util/ObjectObjectPagedHashMap.java

@@ -82,7 +82,7 @@ public V put(K key, V value) {
     public V remove(K key) {
         final long slot = slot(key.hashCode(), mask);
         for (long index = slot;; index = nextSlot(index, mask)) {
-            final V previous = values.set(index, null);
+            final V previous = values.getAndSet(index, null);
             if (previous == null) {
                 return null;
             } else if (keys.get(index).equals(key)) {
@@ -104,7 +104,7 @@ private V set(K key, int code, V value) {
         assert size < maxSize;
         final long slot = slot(code, mask);
         for (long index = slot;; index = nextSlot(index, mask)) {
-            final V previous = values.set(index, value);
+            final V previous = values.getAndSet(index, value);
             if (previous == null) {
                 // slot was free
                 keys.set(index, key);
@@ -186,7 +186,7 @@ protected boolean used(long bucket) {
     @Override
     protected void removeAndAdd(long index) {
         final K key = keys.get(index);
-        final V value = values.set(index, null);
+        final V value = values.getAndSet(index, null);
         --size;
         final V removed = set(key, key.hashCode(), value);
         assert removed == null;

server/src/test/java/org/elasticsearch/common/util/BigArraysTests.java

@@ -135,6 +135,7 @@ public void testObjectArrayGrowth() {
             ref[i] = randomFrom(pool);
             array = bigArrays.grow(array, i + 1);
             array.set(i, ref[i]);
+            assertEquals(ref[i], array.getAndSet(i, ref[i]));
         }
         for (int i = 0; i < totalLen; ++i) {
             assertSame(ref[i], array.get(i));

test/framework/src/main/java/org/elasticsearch/common/util/MockBigArrays.java

@@ -695,8 +695,13 @@ public T get(long index) {
         }
 
         @Override
-        public T set(long index, T value) {
-            return in.set(index, value);
+        public void set(long index, T value) {
+            in.set(index, value);
+        }
+
+        @Override
+        public T getAndSet(long index, T value) {
+            return in.getAndSet(index, value);
         }
 
         @Override