Palette specialization (#638)

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

@@ -0,0 +1,160 @@
+package net.minestom.server.instance.palette;
+
+import net.minestom.server.utils.binary.BinaryWriter;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.function.IntUnaryOperator;
+
+/**
+ * Palette that switches between its backend based on the use case.
+ */
+final class AdaptivePalette implements Palette {
+    final int dimension;
+    final int dimensionBitCount;
+    final int maxBitsPerEntry;
+    final int defaultBitsPerEntry;
+    final int bitsIncrement;
+
+    private SpecializedPalette palette;
+
+    AdaptivePalette(int dimension, int maxBitsPerEntry, int bitsPerEntry, int bitsIncrement) {
+        this.dimensionBitCount = validateDimension(dimension);
+
+        this.dimension = dimension;
+        this.maxBitsPerEntry = maxBitsPerEntry;
+        this.defaultBitsPerEntry = bitsPerEntry;
+        this.bitsIncrement = bitsIncrement;
+
+        this.palette = new FilledPalette(dimension, 0);
+    }
+
+    @Override
+    public int get(int x, int y, int z) {
+        if (x < 0 || y < 0 || z < 0) {
+            throw new IllegalArgumentException("Coordinates must be positive");
+        }
+        return palette.get(x, y, z);
+    }
+
+    @Override
+    public void getAll(@NotNull EntryConsumer consumer) {
+        this.palette.getAll(consumer);
+    }
+
+    @Override
+    public void getAllPresent(@NotNull EntryConsumer consumer) {
+        this.palette.getAllPresent(consumer);
+    }
+
+    @Override
+    public void set(int x, int y, int z, int value) {
+        if (x < 0 || y < 0 || z < 0) {
+            throw new IllegalArgumentException("Coordinates must be positive");
+        }
+        Palette palette = switchFlexible();
+        palette.set(x, y, z, value);
+    }
+
+    @Override
+    public void fill(int value) {
+        this.palette = new FilledPalette(dimension, value);
+    }
+
+    @Override
+    public void setAll(@NotNull EntrySupplier supplier) {
+        SpecializedPalette newPalette = new FlexiblePalette(this);
+        newPalette.setAll(supplier);
+        this.palette = newPalette;
+    }
+
+    @Override
+    public void replace(int x, int y, int z, @NotNull IntUnaryOperator operator) {
+        if (x < 0 || y < 0 || z < 0) {
+            throw new IllegalArgumentException("Coordinates must be positive");
+        }
+        Palette palette = switchFlexible();
+        palette.replace(x, y, z, operator);
+    }
+
+    @Override
+    public void replaceAll(@NotNull EntryFunction function) {
+        Palette palette = switchFlexible();
+        palette.replaceAll(function);
+    }
+
+    @Override
+    public int count() {
+        return palette.count();
+    }
+
+    @Override
+    public int bitsPerEntry() {
+        return palette.bitsPerEntry();
+    }
+
+    @Override
+    public int maxBitsPerEntry() {
+        return maxBitsPerEntry;
+    }
+
+    @Override
+    public int dimension() {
+        return dimension;
+    }
+
+    @Override
+    public @NotNull Palette clone() {
+        try {
+            AdaptivePalette adaptivePalette = (AdaptivePalette) super.clone();
+            adaptivePalette.palette = palette.clone();
+            return adaptivePalette;
+        } catch (CloneNotSupportedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public void write(@NotNull BinaryWriter writer) {
+        optimizedPalette().write(writer);
+    }
+
+    Palette optimizedPalette() {
+        var currentPalette = this.palette;
+        if (currentPalette instanceof FlexiblePalette flexiblePalette) {
+            final int count = flexiblePalette.count();
+            if (count == 0) {
+                currentPalette = new FilledPalette(dimension, 0);
+                this.palette = currentPalette;
+            } else if (count == maxSize()) {
+                // Palette is full
+                final var values = flexiblePalette.paletteToValueList;
+                if (values.size() > 0) {
+                    currentPalette = new FilledPalette(dimension, values.getInt(1));
+                    this.palette = currentPalette;
+                }
+            }
+        }
+        return currentPalette;
+    }
+
+    Palette switchFlexible() {
+        var currentPalette = this.palette;
+        if (currentPalette instanceof FilledPalette filledPalette) {
+            currentPalette = new FlexiblePalette(this);
+            currentPalette.fill(filledPalette.value());
+            this.palette = currentPalette;
+        }
+        return currentPalette;
+    }
+
+    private static int validateDimension(int dimension) {
+        if (dimension <= 1) {
+            throw new IllegalArgumentException("Dimension must be greater 1");
+        }
+        double log2 = Math.log(dimension) / Math.log(2);
+        if ((int) Math.ceil(log2) != (int) Math.floor(log2)) {
+            throw new IllegalArgumentException("Dimension must be a power of 2");
+        }
+        return (int) log2;
+    }
+}

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

@@ -0,0 +1,51 @@
+package net.minestom.server.instance.palette;
+
+import net.minestom.server.utils.binary.BinaryWriter;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Palette containing a single value. Useful for both empty and full palettes.
+ */
+record FilledPalette(int dimension, int value) implements SpecializedPalette.Immutable {
+    @Override
+    public int get(int x, int y, int z) {
+        return value;
+    }
+
+    @Override
+    public void getAll(@NotNull EntryConsumer consumer) {
+        final int dimension = dimension();
+        for (int y = 0; y < dimension; y++)
+            for (int z = 0; z < dimension; z++)
+                for (int x = 0; x < dimension; x++)
+                    consumer.accept(x, y, z, value);
+    }
+
+    @Override
+    public void getAllPresent(@NotNull EntryConsumer consumer) {
+        if (value != 0) getAll(consumer);
+    }
+
+    @Override
+    public int count() {
+        return value != 0 ? maxSize() : 0;
+    }
+
+    @Override
+    public @NotNull SpecializedPalette clone() {
+        return this;
+    }
+
+    @Override
+    public void write(@NotNull BinaryWriter writer) {
+        writer.writeByte((byte) 1); // bitsPerEntry
+        // Palette
+        writer.writeVarInt(1);
+        writer.writeVarInt(value);
+        // Data
+        final int length = maxSize() / 64;
+        writer.writeVarInt(length);
+        // TODO: may be possible to write everything in one call instead of a loop
+        for (int i = 0; i < length; i++) writer.writeLong(0);
+    }
+}

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

@@ -10,7 +10,10 @@ import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.IntUnaryOperator;
 
-final class PaletteImpl implements Palette, Cloneable {
+/**
+ * Palette able to take any value anywhere. May consume more memory than required.
+ */
+final class FlexiblePalette implements SpecializedPalette, Cloneable {
     private static final ThreadLocal<int[]> WRITE_CACHE = ThreadLocal.withInitial(() -> new int[4096]);
     private static final int[] MAGIC_MASKS;
     private static final int[] VALUES_PER_LONG;
@@ -26,34 +29,24 @@ final class PaletteImpl implements Palette, Cloneable {
     }
 
     // Specific to this palette type
-    private final int dimension;
+    private final AdaptivePalette adaptivePalette;
-    private final int dimensionBitCount;
-    private final int maxBitsPerEntry;
-    private final int bitsIncrement;
-
     private int bitsPerEntry;
 
     private boolean hasPalette;
     private int lastPaletteIndex = 1; // First index is air
-
     private int count = 0;
 
     private long[] values;
     // palette index = value
-    private IntArrayList paletteToValueList;
+    IntArrayList paletteToValueList;
     // value = palette index
     private Int2IntOpenHashMap valueToPaletteMap;
 
-    PaletteImpl(int dimension, int maxBitsPerEntry, int bitsPerEntry, int bitsIncrement) {
+    FlexiblePalette(AdaptivePalette adaptivePalette) {
-        this.dimensionBitCount = validateDimension(dimension);
+        this.adaptivePalette = adaptivePalette;
 
-        this.dimension = dimension;
+        this.bitsPerEntry = adaptivePalette.defaultBitsPerEntry;
-        this.maxBitsPerEntry = maxBitsPerEntry;
+        this.hasPalette = bitsPerEntry <= maxBitsPerEntry();
-        this.bitsIncrement = bitsIncrement;
-
-        this.bitsPerEntry = bitsPerEntry;
-
-        this.hasPalette = bitsPerEntry <= maxBitsPerEntry;
 
         this.paletteToValueList = new IntArrayList(1);
         this.paletteToValueList.add(0);
@@ -64,9 +57,6 @@ final class PaletteImpl implements Palette, Cloneable {
 
     @Override
     public int get(int x, int y, int z) {
-        if (x < 0 || y < 0 || z < 0) {
-            throw new IllegalArgumentException("Coordinates must be positive");
-        }
         final long[] values = this.values;
         if (values == null) {
             // Section is not loaded, return default value
@@ -74,6 +64,7 @@ final class PaletteImpl implements Palette, Cloneable {
         }
         final int bitsPerEntry = this.bitsPerEntry;
         final int valuesPerLong = VALUES_PER_LONG[bitsPerEntry];
+        final int dimension = dimension();
 
         final int sectionIdentifier = getSectionIndex(x % dimension, y % dimension, z % dimension);
         final int index = sectionIdentifier / valuesPerLong;
@@ -95,13 +86,11 @@ final class PaletteImpl implements Palette, Cloneable {
 
     @Override
     public void set(int x, int y, int z, int value) {
-        if (x < 0 || y < 0 || z < 0) {
-            throw new IllegalArgumentException("Coordinates must be positive");
-        }
         final boolean placedAir = value == 0;
         if (!placedAir) value = getPaletteIndex(value);
         final int bitsPerEntry = this.bitsPerEntry;
         final int valuesPerLong = VALUES_PER_LONG[bitsPerEntry];
+        final int dimension = dimension();
         long[] values = this.values;
         if (values == null) {
             if (placedAir) {
@@ -161,7 +150,7 @@ final class PaletteImpl implements Palette, Cloneable {
     @Override
     public void setAll(@NotNull EntrySupplier supplier) {
         int[] cache = sizeCache(maxSize());
-        final int dimension = this.dimension;
+        final int dimension = dimension();
         // Fill cache with values
         int fillValue = -1;
         int count = 0;
@@ -232,18 +221,18 @@ final class PaletteImpl implements Palette, Cloneable {
 
     @Override
     public int maxBitsPerEntry() {
-        return maxBitsPerEntry;
+        return adaptivePalette.maxBitsPerEntry();
     }
 
     @Override
     public int dimension() {
-        return dimension;
+        return adaptivePalette.dimension();
     }
 
     @Override
-    public @NotNull Palette clone() {
+    public @NotNull SpecializedPalette clone() {
         try {
-            PaletteImpl palette = (PaletteImpl) super.clone();
+            FlexiblePalette palette = (FlexiblePalette) super.clone();
             palette.values = values != null ? values.clone() : null;
             palette.paletteToValueList = paletteToValueList.clone();
             palette.valueToPaletteMap = valueToPaletteMap.clone();
@@ -258,19 +247,19 @@ final class PaletteImpl implements Palette, Cloneable {
     @Override
     public void write(@NotNull BinaryWriter writer) {
         writer.writeByte((byte) bitsPerEntry);
-        if (bitsPerEntry <= maxBitsPerEntry) { // Palette index
+        if (bitsPerEntry <= maxBitsPerEntry()) { // Palette index
             writer.writeVarIntList(paletteToValueList, BinaryWriter::writeVarInt);
         }
         writer.writeLongArray(values);
     }
 
     private int fixBitsPerEntry(int bitsPerEntry) {
-        return bitsPerEntry > maxBitsPerEntry ? 15 : bitsPerEntry;
+        return bitsPerEntry > maxBitsPerEntry() ? 15 : bitsPerEntry;
     }
 
     private void retrieveAll(@NotNull EntryConsumer consumer, boolean consumeEmpty) {
         final long[] values = this.values;
-        final int dimension = this.dimension;
+        final int dimension = this.dimension();
         if (values == null) {
             if (consumeEmpty) {
                 // No values, give all 0 to make the consumer happy
@@ -287,7 +276,7 @@ final class PaletteImpl implements Palette, Cloneable {
         final int size = maxSize();
         final int dimensionMinus = dimension - 1;
         final int[] ids = hasPalette ? paletteToValueList.elements() : null;
-        final int dimensionBitCount = this.dimensionBitCount;
+        final int dimensionBitCount = adaptivePalette.dimensionBitCount;
         final int shiftedDimensionBitCount = dimensionBitCount << 1;
         for (int i = 0; i < values.length; i++) {
             final long value = values[i];
@@ -345,8 +334,8 @@ final class PaletteImpl implements Palette, Cloneable {
     }
 
     private void resize(int newBitsPerEntry) {
-        newBitsPerEntry = fixBitsPerEntry(newBitsPerEntry);
+        FlexiblePalette palette = new FlexiblePalette(adaptivePalette);
-        PaletteImpl palette = new PaletteImpl(dimension, maxBitsPerEntry, newBitsPerEntry, bitsIncrement);
+        palette.bitsPerEntry = fixBitsPerEntry(newBitsPerEntry);
         palette.lastPaletteIndex = lastPaletteIndex;
         palette.paletteToValueList = paletteToValueList;
         palette.valueToPaletteMap = valueToPaletteMap;
@@ -363,7 +352,7 @@ final class PaletteImpl implements Palette, Cloneable {
         final int lastPaletteIndex = this.lastPaletteIndex;
         if (lastPaletteIndex >= maxPaletteSize(bitsPerEntry)) {
             // Palette is full, must resize
-            resize(bitsPerEntry + bitsIncrement);
+            resize(bitsPerEntry + adaptivePalette.bitsIncrement);
             return getPaletteIndex(value);
         }
         final int lookup = valueToPaletteMap.putIfAbsent(value, lastPaletteIndex);
@@ -374,6 +363,7 @@ final class PaletteImpl implements Palette, Cloneable {
     }
 
     int getSectionIndex(int x, int y, int z) {
+        final int dimensionBitCount = adaptivePalette.dimensionBitCount;
         return y << (dimensionBitCount << 1) | z << dimensionBitCount | x;
     }
 
@@ -394,15 +384,4 @@ final class PaletteImpl implements Palette, Cloneable {
     static int maxPaletteSize(int bitsPerEntry) {
         return 1 << bitsPerEntry;
     }
-
-    private static int validateDimension(int dimension) {
-        if (dimension <= 1) {
-            throw new IllegalArgumentException("Dimension must be greater 1");
-        }
-        double log2 = Math.log(dimension) / Math.log(2);
-        if ((int) Math.ceil(log2) != (int) Math.floor(log2)) {
-            throw new IllegalArgumentException("Dimension must be a power of 2");
-        }
-        return (int) log2;
-    }
 }

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

@@ -10,7 +10,7 @@ import java.util.function.IntUnaryOperator;
  * <p>
  * 0 is the default value.
  */
-public sealed interface Palette extends Writeable permits PaletteImpl {
+public interface Palette extends Writeable {
     static Palette blocks() {
         return newPalette(16, 8, 6, 1);
     }
@@ -20,7 +20,7 @@ public sealed interface Palette extends Writeable permits PaletteImpl {
     }
 
     static Palette newPalette(int dimension, int maxBitsPerEntry, int bitsPerEntry, int bitIncrement) {
-        return new PaletteImpl(dimension, maxBitsPerEntry, bitsPerEntry, bitIncrement);
+        return new AdaptivePalette(dimension, maxBitsPerEntry, bitsPerEntry, bitIncrement);
     }
 
     int get(int x, int y, int z);

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

@@ -0,0 +1,47 @@
+package net.minestom.server.instance.palette;
+
+import org.jetbrains.annotations.NotNull;
+
+import java.util.function.IntUnaryOperator;
+
+interface SpecializedPalette extends Palette {
+    @Override
+    default int bitsPerEntry() {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    default int maxBitsPerEntry() {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    @NotNull SpecializedPalette clone();
+
+    interface Immutable extends SpecializedPalette {
+        @Override
+        default void set(int x, int y, int z, int value) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        default void fill(int value) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        default void setAll(@NotNull EntrySupplier supplier) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        default void replace(int x, int y, int z, @NotNull IntUnaryOperator operator) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        default void replaceAll(@NotNull EntryFunction function) {
+            throw new UnsupportedOperationException();
+        }
+    }
+}