Optimize #replaceAll

jmh-benchmarks/src/jmh/java/net/minestom/jmh/palette/PaletteReplaceBenchmark.java

@@ -0,0 +1,45 @@
+package net.minestom.jmh.palette;
+
+import net.minestom.server.instance.palette.Palette;
+import org.openjdk.jmh.annotations.*;
+
+import java.util.concurrent.TimeUnit;
+
+@Warmup(iterations = 5, time = 1000, timeUnit = TimeUnit.MILLISECONDS)
+@Measurement(iterations = 10, time = 1000, timeUnit = TimeUnit.MILLISECONDS)
+@Fork(3)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@State(Scope.Benchmark)
+public class PaletteReplaceBenchmark {
+
+    //@Param({"4", "16"})
+    //public int dimension;
+
+    private Palette palette;
+
+    @Setup
+    public void setup() {
+        // FIXME: StackOverflowError
+        // palette = Palette.newPalette(dimension, 15, 4, 1);
+        palette = Palette.blocks();
+        palette.setAll((x, y, z) -> x + y + z + 1);
+    }
+
+    @Benchmark
+    public void replaceAll() {
+        palette.replaceAll((x, y, z, value) -> value + 1);
+    }
+
+    @Benchmark
+    public void replaceLoop() {
+        final int dimension = palette.dimension();
+        for (int x = 0; x < dimension; x++) {
+            for (int y = 0; y < dimension; y++) {
+                for (int z = 0; z < dimension; z++) {
+                    palette.replace(x, y, z, value -> value + 1);
+                }
+            }
+        }
+    }
+}

src/main/java/net/minestom/server/instance/palette/PaletteImpl.java

@@ -7,6 +7,7 @@ import net.minestom.server.utils.binary.BinaryWriter;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.IntUnaryOperator;
 
 final class PaletteImpl implements Palette, Cloneable {
@@ -191,11 +192,7 @@ final class PaletteImpl implements Palette, Cloneable {
 
     @Override
     public void setAll(@NotNull EntrySupplier supplier) {
-        int[] cache = WRITE_CACHE.get();
-        if (cache.length < maxSize()) {
-            cache = new int[maxSize()];
-            WRITE_CACHE.set(cache);
-        }
+        int[] cache = sizeCache(maxSize());
         // Fill cache with values
         final int dimensionMinus = dimension - 1;
         int count = 0;
@@ -244,14 +241,22 @@ final class PaletteImpl implements Palette, Cloneable {
 
     @Override
     public void replaceAll(@NotNull EntryFunction function) {
-        // TODO optimize
-        for (int y = 0; y < dimension; y++) {
-            for (int z = 0; z < dimension; z++) {
-                for (int x = 0; x < dimension; x++) {
-                    set(x, y, z, function.apply(x, y, z, get(x, y, z)));
-                }
-            }
-        }
+        int[] cache = sizeCache(maxSize());
+        AtomicInteger count = new AtomicInteger();
+        // Fill cache with values
+        getAll((x, y, z, value) -> {
+            final int newValue = function.apply(x, y, z, value);
+            final int index = count.getPlain();
+            count.setPlain(index + 1);
+            cache[index] = getPaletteIndex(newValue);
+        });
+        // Set values to final array
+        count.set(0);
+        setAll((x, y, z) -> {
+            final int index = count.getPlain();
+            count.setPlain(index + 1);
+            return cache[index];
+        });
     }
 
     @Override
@@ -358,6 +363,15 @@ final class PaletteImpl implements Palette, Cloneable {
         return y << (dimensionBitCount << 1) | z << dimensionBitCount | x;
     }
 
+    static int[] sizeCache(int size) {
+        int[] cache = WRITE_CACHE.get();
+        if (cache.length < size) {
+            cache = new int[size];
+            WRITE_CACHE.set(cache);
+        }
+        return cache;
+    }
+
     static int maxPaletteSize(int bitsPerEntry) {
         return 1 << bitsPerEntry;
     }

src/test/java/net/minestom/server/instance/PaletteTest.java

@@ -194,15 +194,44 @@ public class PaletteTest {
         }
     }
 
+    @Test
+    public void replaceAll() {
+        var palettes = testPalettes();
+        for (Palette palette : palettes) {
+            palette.setAll((x, y, z) -> x+y+z+1);
+            palette.replaceAll((x, y, z, value) -> {
+                assertEquals(x+y+z+1, value);
+                return x+y+z+2;
+            });
+            palette.getAll((x, y, z, value) -> assertEquals(x+y+z+2, value));
+        }
+    }
+
     @Test
     public void replace() {
-        var palette = Palette.blocks();
-        palette.set(0, 0, 0, 1);
-        palette.replace(0, 0, 0, operand -> {
-            assertEquals(1, operand);
-            return 2;
-        });
-        assertEquals(2, palette.get(0, 0, 0));
+        var palettes = testPalettes();
+        for (Palette palette : palettes) {
+            palette.set(0, 0, 0, 1);
+            palette.replace(0, 0, 0, operand -> {
+                assertEquals(1, operand);
+                return 2;
+            });
+            assertEquals(2, palette.get(0, 0, 0));
+        }
+    }
+
+    @Test
+    public void replaceLoop(){
+        var palette = Palette.newPalette(2, 15, 4, 1);
+        palette.setAll((x, y, z) -> x+y+z);
+        final int dimension = palette.dimension();
+        for (int x = 0; x < dimension; x++) {
+            for (int y = 0; y < dimension; y++) {
+                for (int z = 0; z < dimension; z++) {
+                    palette.replace(x, y, z, value -> value + 1);
+                }
+            }
+        }
     }
 
     @Test