From cdcfff5b102e4974d931e082d706c97c37a9797d Mon Sep 17 00:00:00 2001
From: DaPorkchop_ <daporkchop@daporkchop.net>
Date: Sun, 18 Oct 2020 21:41:56 +0200
Subject: [PATCH] allow duplicate entries in palette

---
 .../mc/protocol/data/game/chunk/Chunk.java     | 18 +++++++++++-------
 .../data/game/chunk/palette/ListPalette.java   | 13 +++++++++++++
 .../data/game/chunk/palette/MapPalette.java    | 15 +++++++++++++++
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/Chunk.java b/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/Chunk.java
index 9bb21522..76e2d89e 100644
--- a/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/Chunk.java
+++ b/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/Chunk.java
@@ -33,13 +33,7 @@ public class Chunk {
         int blockCount = in.readShort();
         int bitsPerEntry = in.readUnsignedByte();
 
-        Palette palette = createPalette(bitsPerEntry);
-        if(!(palette instanceof GlobalPalette)) {
-            int paletteLength = in.readVarInt();
-            for(int i = 0; i < paletteLength; i++) {
-                palette.stateToId(in.readVarInt());
-            }
-        }
+        Palette palette = readPalette(bitsPerEntry, in);
 
         BitStorage storage = new BitStorage(bitsPerEntry, CHUNK_SIZE, in.readLongs(in.readVarInt()));
         return new Chunk(blockCount, palette, storage);
@@ -120,6 +114,16 @@ public class Chunk {
         }
     }
 
+    private static Palette readPalette(int bitsPerEntry, NetInput in) throws IOException {
+        if(bitsPerEntry <= MIN_PALETTE_BITS_PER_ENTRY) {
+            return new ListPalette(bitsPerEntry, in);
+        } else if(bitsPerEntry <= MAX_PALETTE_BITS_PER_ENTRY) {
+            return new MapPalette(bitsPerEntry, in);
+        } else {
+            return new GlobalPalette();
+        }
+    }
+
     private static int index(int x, int y, int z) {
         return y << 8 | z << 4 | x;
     }
diff --git a/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/palette/ListPalette.java b/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/palette/ListPalette.java
index e82b33ec..8518be24 100644
--- a/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/palette/ListPalette.java
+++ b/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/palette/ListPalette.java
@@ -1,7 +1,10 @@
 package com.github.steveice10.mc.protocol.data.game.chunk.palette;
 
+import com.github.steveice10.packetlib.io.NetInput;
 import lombok.EqualsAndHashCode;
 
+import java.io.IOException;
+
 /**
  * A palette backed by a List.
  */
@@ -18,6 +21,16 @@ public class ListPalette implements Palette {
         this.data = new int[this.maxId + 1];
     }
 
+    public ListPalette(int bitsPerEntry, NetInput in) throws IOException {
+        this(bitsPerEntry);
+
+        int paletteLength = in.readVarInt();
+        for(int i = 0; i < paletteLength; i++) {
+            this.data[i] = in.readVarInt();
+        }
+        this.nextId = paletteLength;
+    }
+
     @Override
     public int size() {
         return this.nextId;
diff --git a/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/palette/MapPalette.java b/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/palette/MapPalette.java
index 92a43bdc..cc16f97f 100644
--- a/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/palette/MapPalette.java
+++ b/src/main/java/com/github/steveice10/mc/protocol/data/game/chunk/palette/MapPalette.java
@@ -1,9 +1,12 @@
 package com.github.steveice10.mc.protocol.data.game.chunk.palette;
 
+import com.github.steveice10.packetlib.io.NetInput;
 import io.netty.util.collection.IntObjectHashMap;
 import io.netty.util.collection.IntObjectMap;
 import lombok.EqualsAndHashCode;
 
+import java.io.IOException;
+
 /**
  * A palette backed by a map.
  */
@@ -21,6 +24,18 @@ public class MapPalette implements Palette {
         this.idToState = new int[this.maxId + 1];
     }
 
+    public MapPalette(int bitsPerEntry, NetInput in) throws IOException {
+        this(bitsPerEntry);
+
+        int paletteLength = in.readVarInt();
+        for(int i = 0; i < paletteLength; i++) {
+            int state = in.readVarInt();
+            this.idToState[i] = state;
+            this.stateToId.putIfAbsent(state, i);
+        }
+        this.nextId = paletteLength;
+    }
+
     @Override
     public int size() {
         return this.nextId;