Decode saved section palettes directly

common/src/main/java/xyz/jpenilla/squaremap/common/chunksnapshot/ChunkSnapshotFactory.java

@@ -1,6 +1,5 @@
 package xyz.jpenilla.squaremap.common.chunksnapshot;
 
-import com.mojang.serialization.Codec;
 import java.util.Arrays;
 import java.util.EnumMap;
 import java.util.Map;
@@ -9,7 +8,6 @@
 import net.minecraft.core.Holder;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.nbt.ListTag;
-import net.minecraft.nbt.NbtOps;
 import net.minecraft.nbt.NbtUtils;
 import net.minecraft.util.Mth;
 import net.minecraft.world.level.ChunkPos;
@@ -19,9 +17,6 @@
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraft.world.level.chunk.ChunkAccess;
 import net.minecraft.world.level.chunk.LevelChunkSection;
-import net.minecraft.world.level.chunk.PalettedContainer;
-import net.minecraft.world.level.chunk.PalettedContainerFactory;
-import net.minecraft.world.level.chunk.PalettedContainerRO;
 import net.minecraft.world.level.chunk.storage.SerializableChunkData;
 import net.minecraft.world.level.dimension.DimensionType;
 import net.minecraft.world.level.levelgen.Heightmap;
@@ -51,8 +46,8 @@ public static ChunkSnapshot snapshotLiveChunk(final Level level, final ChunkAcce
         final LevelHeightAccessor heightAccessor = LevelHeightAccessor.create(chunk.getMinY(), chunk.getHeight());
         final int sectionCount = heightAccessor.getSectionsCount();
         final LevelChunkSection[] sections = chunk.getSections();
-        final PalettedContainer<BlockState>[] states = new PalettedContainer[sectionCount];
-        final PalettedContainerRO<Holder<Biome>>[] biomes = new PalettedContainerRO[sectionCount];
+        final SectionData<BlockState>[] states = new SectionData[sectionCount];
+        final SectionData<Holder<Biome>>[] biomes = new SectionData[sectionCount];
 
         final boolean[] empty = new boolean[sectionCount];
         if (!chunk.hasPrimedHeightmap(Heightmap.Types.WORLD_SURFACE)) {
@@ -68,10 +63,10 @@ public static ChunkSnapshot snapshotLiveChunk(final Level level, final ChunkAcce
             if (sectionEmpty) {
                 states[i] = EmptySectionHolder.getEmptySectionBlockStates();
             } else {
-                states[i] = sections[i].getStates().copy();
+                states[i] = SectionData.of(sections[i].getStates().copy());
             }
 
-            biomes[i] = sections[i].getBiomes().copy();
+            biomes[i] = SectionData.of(sections[i].getBiomes().copy());
         }
 
         return new ChunkSnapshotImpl(
@@ -91,12 +86,12 @@ public static ChunkSnapshot snapshotLiveChunk(final Level level, final ChunkAcce
      *
      * <p>The caller must supply NBT already upgraded to the current data version and classified
      * as {@link ChunkSnapshotEligibility#ELIGIBLE}. The data is only read, so it may be shared
-     * with other readers, but must not be mutated while this runs. This method does not check
-     * generation status or complete retrogen and may run off-thread.</p>
+     * with other readers, but must not be mutated afterward; the snapshot keeps its data arrays.
+     * This method does not check generation status or complete retrogen and may run off-thread.</p>
      *
      * @param levelHeight the level's vertical bounds
      * @param dimensionType the level's dimension type
-     * @param containerFactory the level's palette container factory
+     * @param sectionDecoder the level's section decoder
      * @param chunkData the upgraded, snapshot-ready chunk data
      * @return the decoded snapshot
      * @throws SerializableChunkData.ChunkReadException if a section palette cannot be decoded
@@ -105,36 +100,37 @@ public static ChunkSnapshot snapshotLiveChunk(final Level level, final ChunkAcce
     public static ChunkSnapshot snapshotFromChunkData(
         final LevelHeightAccessor levelHeight,
         final DimensionType dimensionType,
-        final PalettedContainerFactory containerFactory,
+        final SectionDecoder sectionDecoder,
         final CompoundTag chunkData
     ) {
         final ChunkPos chunkPos = readChunkPos(chunkData);
         final ListTag sectionTags = chunkData.getListOrEmpty(SerializableChunkData.SECTIONS_TAG);
         final int sectionCount = levelHeight.getSectionsCount();
 
-        final PalettedContainer<BlockState>[] states = new PalettedContainer[sectionCount];
-        final PalettedContainerRO<Holder<Biome>>[] biomes = new PalettedContainerRO[sectionCount];
+        final SectionData<BlockState>[] states = new SectionData[sectionCount];
+        final SectionData<Holder<Biome>>[] biomes = new SectionData[sectionCount];
         final boolean[] empty = new boolean[sectionCount];
         Arrays.fill(states, EmptySectionHolder.getEmptySectionBlockStates());
         Arrays.fill(biomes, EmptySectionHolder.getEmptySectionBiomes());
         Arrays.fill(empty, true);
 
-        final Codec<PalettedContainerRO<Holder<Biome>>> biomesCodec = containerFactory.biomeContainerCodec();
-        final Codec<PalettedContainer<BlockState>> blockStatesCodec = containerFactory.blockStatesContainerCodec();
-
         for (int i = 0; i < sectionTags.size(); ++i) {
             final Optional<CompoundTag> maybeSectionTag = sectionTags.getCompound(i);
             if (maybeSectionTag.isPresent()) {
                 final CompoundTag sectionTag = maybeSectionTag.get();
                 final int y = sectionTag.getByteOr(ChunkDataKeys.SECTION_Y, (byte) 0);
                 final int index = y - levelHeight.getMinSectionY();
                 if (index >= 0 && index < sectionCount) {
-                    final Optional<PalettedContainer<BlockState>> maybeBlocks = sectionTag.getCompound(ChunkDataKeys.BLOCK_STATES)
-                        .map((container) -> blockStatesCodec.parse(NbtOps.INSTANCE, container).promotePartial((msg) -> Logging.logger().warn("Failed to decode chunk {} section {}: {}", chunkPos, y, msg)).getOrThrow(SerializableChunkData.ChunkReadException::new));
-                    final PalettedContainer<BlockState> blocks = maybeBlocks.orElse(EmptySectionHolder.getEmptySectionBlockStates());
-                    final Optional<PalettedContainerRO<Holder<Biome>>> maybeBiomes = sectionTag.getCompound(ChunkDataKeys.BIOMES)
-                        .map((container) -> biomesCodec.parse(NbtOps.INSTANCE, container).promotePartial((msg) -> Logging.logger().warn("Failed to decode chunk {} section {}: {}", chunkPos, y, msg)).getOrThrow(SerializableChunkData.ChunkReadException::new));
-                    final PalettedContainerRO<Holder<Biome>> sectionBiomes = maybeBiomes.orElse(EmptySectionHolder.getEmptySectionBiomes());
+                    final SectionData<BlockState> blocks = sectionTag.getCompound(ChunkDataKeys.BLOCK_STATES)
+                        .map(container -> sectionDecoder.blockStates(container)
+                            .promotePartial(msg -> logSectionDecodeErrors(chunkPos, y, msg))
+                            .getOrThrow(SerializableChunkData.ChunkReadException::new))
+                        .orElseGet(EmptySectionHolder::getEmptySectionBlockStates);
+                    final SectionData<Holder<Biome>> sectionBiomes = sectionTag.getCompound(ChunkDataKeys.BIOMES)
+                        .map(container -> sectionDecoder.biomes(container)
+                            .promotePartial(msg -> logSectionDecodeErrors(chunkPos, y, msg))
+                            .getOrThrow(SerializableChunkData.ChunkReadException::new))
+                        .orElseGet(EmptySectionHolder::getEmptySectionBiomes);
                     states[index] = blocks;
                     empty[index] = !blocks.maybeHas(state -> !state.isAir());
                     biomes[index] = sectionBiomes;
@@ -188,10 +184,15 @@ public static ChunkPos readChunkPos(final CompoundTag chunkData) {
         );
     }
 
+    private static void logSectionDecodeErrors(final ChunkPos chunkPos, final int sectionY, final String message) {
+        // Same message as vanilla's SerializableChunkData.logErrors.
+        Logging.logger().error("Recoverable errors when loading section [{}, {}, {}]: {}", chunkPos.x(), sectionY, chunkPos.z(), message);
+    }
+
     private static Map<Heightmap.Types, HeightmapSnapshot> readOrComputeHeightmaps(
         final LevelHeightAccessor levelHeight,
         final CompoundTag chunkData,
-        final PalettedContainer<BlockState>[] states,
+        final SectionData<BlockState>[] states,
         final boolean[] empty
     ) {
         final int bits = Mth.ceillog2(levelHeight.getHeight() + 1);
@@ -208,10 +209,10 @@ private static Map<Heightmap.Types, HeightmapSnapshot> readOrComputeHeightmaps(
     }
 
     @SuppressWarnings({"unchecked", "rawtypes"})
-    private static @Nullable Holder<Biome>[] uniformBiomes(final PalettedContainerRO<Holder<Biome>>[] biomes) {
+    private static @Nullable Holder<Biome>[] uniformBiomes(final SectionData<Holder<Biome>>[] biomes) {
         final @Nullable Holder<Biome>[] uniform = new Holder[biomes.length];
         for (int i = 0; i < biomes.length; i++) {
-            final Holder<Biome> first = biomes[i].get(0, 0, 0);
+            final Holder<Biome> first = biomes[i].get(0);
             // Checks palette entries only, so a palette with unused entries reads as mixed.
             if (!biomes[i].maybeHas(biome -> biome != first)) {
                 uniform[i] = first;

common/src/main/java/xyz/jpenilla/squaremap/common/chunksnapshot/ChunkSnapshotImpl.java

@@ -11,8 +11,6 @@
 import net.minecraft.world.level.biome.Biome;
 import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.block.state.BlockState;
-import net.minecraft.world.level.chunk.PalettedContainer;
-import net.minecraft.world.level.chunk.PalettedContainerRO;
 import net.minecraft.world.level.dimension.DimensionType;
 import net.minecraft.world.level.levelgen.Heightmap;
 import net.minecraft.world.level.material.FluidState;
@@ -24,8 +22,8 @@
 @DefaultQualifier(NonNull.class)
 record ChunkSnapshotImpl(
     LevelHeightAccessor heightAccessor,
-    PalettedContainer<BlockState>[] states,
-    PalettedContainerRO<Holder<Biome>>[] biomes,
+    SectionData<BlockState>[] states,
+    SectionData<Holder<Biome>>[] biomes,
     @Nullable Holder<Biome>[] uniformBiomes,
     Map<Heightmap.Types, HeightmapSnapshot> heightmaps,
     boolean[] emptySections,
@@ -95,7 +93,7 @@ public Holder<Biome> getNoiseBiome(final int quartX, final int quartY, final int
         if (uniform != null) {
             return uniform;
         }
-        return this.biomes[sectionIndex].get(quartX & 3, clampedQuartY & 3, quartZ & 3);
+        return this.biomes[sectionIndex].get((clampedQuartY & 3) << 4 | (quartZ & 3) << 2 | quartX & 3);
     }
 
     @Override

common/src/main/java/xyz/jpenilla/squaremap/common/chunksnapshot/EmptySectionHolder.java

@@ -4,39 +4,35 @@
 import net.minecraft.world.level.biome.Biome;
 import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.block.state.BlockState;
-import net.minecraft.world.level.chunk.PalettedContainer;
 import net.minecraft.world.level.chunk.PalettedContainerFactory;
 import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.framework.qual.DefaultQualifier;
 
 @DefaultQualifier(NonNull.class)
 public final class EmptySectionHolder {
-    private static @MonotonicNonNull PalettedContainer<BlockState> EMPTY_SECTION_BLOCK_STATES;
-    private static @MonotonicNonNull PalettedContainer<Holder<Biome>> EMPTY_SECTION_BIOMES;
+    private static @MonotonicNonNull SectionData<BlockState> EMPTY_SECTION_BLOCK_STATES;
+    private static @MonotonicNonNull SectionData<Holder<Biome>> EMPTY_SECTION_BIOMES;
 
     private EmptySectionHolder() {
     }
 
     public static void init(final PalettedContainerFactory palettedContainerFactory) {
         if (EMPTY_SECTION_BLOCK_STATES == null) {
-            EMPTY_SECTION_BLOCK_STATES = new PalettedContainer<>(
-                Blocks.AIR.defaultBlockState(),
-                palettedContainerFactory.blockStatesStrategy()
-            );
+            EMPTY_SECTION_BLOCK_STATES = SectionData.Packed.single(Blocks.AIR.defaultBlockState());
         }
         // Always replace, biomes can change between integrated server runs
-        EMPTY_SECTION_BIOMES = palettedContainerFactory.createForBiomes();
+        EMPTY_SECTION_BIOMES = SectionData.Packed.single(palettedContainerFactory.defaultBiome());
     }
 
-    static PalettedContainer<BlockState> getEmptySectionBlockStates() {
+    static SectionData<BlockState> getEmptySectionBlockStates() {
         if (EMPTY_SECTION_BLOCK_STATES == null) {
             throw new IllegalStateException("EmptySectionHolder not initialized");
         }
         return EMPTY_SECTION_BLOCK_STATES;
     }
 
-    static PalettedContainer<Holder<Biome>> getEmptySectionBiomes() {
+    static SectionData<Holder<Biome>> getEmptySectionBiomes() {
         if (EMPTY_SECTION_BIOMES == null) {
             throw new IllegalStateException("EmptySectionHolder not initialized");
         }

common/src/main/java/xyz/jpenilla/squaremap/common/chunksnapshot/HeightmapSnapshot.java

@@ -7,7 +7,6 @@
 import net.minecraft.world.level.LevelHeightAccessor;
 import net.minecraft.world.level.block.state.BlockState;
 import net.minecraft.world.level.chunk.ChunkAccess;
-import net.minecraft.world.level.chunk.PalettedContainer;
 import net.minecraft.world.level.levelgen.Heightmap;
 
 final class HeightmapSnapshot {
@@ -36,15 +35,15 @@ final class HeightmapSnapshot {
 
     static HeightmapSnapshot computeWorldSurface(
         final LevelHeightAccessor heightAccessor,
-        final PalettedContainer<BlockState>[] sections,
+        final SectionData<BlockState>[] sections,
         final boolean[] empty
     ) {
         final SimpleBitStorage data = new SimpleBitStorage(Mth.ceillog2(heightAccessor.getHeight() + 1), 256);
         final Predicate<BlockState> isOpaque = Heightmap.Types.WORLD_SURFACE.isOpaque();
         // Rebuild directly from palette data using the heightmap type's predicate,
         // avoiding live section construction and block counting.
         for (int sectionIndex = sections.length - 1; sectionIndex >= 0; sectionIndex--) {
-            final PalettedContainer<BlockState> section = sections[sectionIndex];
+            final SectionData<BlockState> section = sections[sectionIndex];
             if (empty[sectionIndex]) {
                 continue;
             }
@@ -56,7 +55,7 @@ static HeightmapSnapshot computeWorldSurface(
                         continue;
                     }
                     for (int y = 15; y >= 0; y--) {
-                        if (isOpaque.test(section.get(x, y, z))) {
+                        if (isOpaque.test(section.get(y << 8 | z << 4 | x))) {
                             data.set(index, sectionIndex * 16 + y + 1);
                             break;
                         }

common/src/main/java/xyz/jpenilla/squaremap/common/chunksnapshot/SectionData.java

@@ -0,0 +1,80 @@
+package xyz.jpenilla.squaremap.common.chunksnapshot;
+
+import java.util.function.Predicate;
+import net.minecraft.util.SimpleBitStorage;
+import net.minecraft.world.level.chunk.MissingPaletteEntryException;
+import net.minecraft.world.level.chunk.PalettedContainer;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.checkerframework.framework.qual.DefaultQualifier;
+
+/**
+ * Read-only block states or biomes of a chunk section.
+ *
+ * @param <T> the value type
+ */
+@DefaultQualifier(NonNull.class)
+sealed interface SectionData<T> permits SectionData.Container, SectionData.Packed {
+    /**
+     * Gets a value by its index, as in {@link PalettedContainer#get(int)}.
+     *
+     * @param index the index
+     * @return the value
+     */
+    T get(int index);
+
+    /**
+     * Checks the palette, like {@link PalettedContainer#maybeHas(Predicate)}.
+     *
+     * @param predicate the predicate
+     * @return whether any palette entry matches
+     */
+    boolean maybeHas(Predicate<T> predicate);
+
+    static <T> SectionData<T> of(final PalettedContainer<T> container) {
+        return new Container<>(container);
+    }
+
+    record Container<T>(PalettedContainer<T> container) implements SectionData<T> {
+        @Override
+        public T get(final int index) {
+            return this.container.get(index);
+        }
+
+        @Override
+        public boolean maybeHas(final Predicate<T> predicate) {
+            return this.container.maybeHas(predicate);
+        }
+    }
+
+    // Storage is null for a single-value palette.
+    record Packed<T>(T[] palette, @Nullable SimpleBitStorage storage) implements SectionData<T> {
+        @SuppressWarnings("unchecked")
+        static <T> Packed<T> single(final T value) {
+            return new Packed<>((T[]) new Object[]{value}, null);
+        }
+
+        @Override
+        public T get(final int index) {
+            if (this.storage == null) {
+                return this.palette[0];
+            }
+            final int id = this.storage.get(index);
+            if (id >= this.palette.length) {
+                // Like vanilla's linear and hash map palettes, which also fail on read.
+                throw new MissingPaletteEntryException(id);
+            }
+            return this.palette[id];
+        }
+
+        @Override
+        public boolean maybeHas(final Predicate<T> predicate) {
+            for (final T value : this.palette) {
+                if (predicate.test(value)) {
+                    return true;
+                }
+            }
+            return false;
+        }
+    }
+}