Reduce poi memory

common/src/main/java/net/caffeinemc/mods/lithium/common/util/collections/ListeningLong2ObjectOpenHashMap.java

@@ -40,6 +40,13 @@ public V remove(long k) {
         return ret;
     }
 
+    /**
+     * Only use this if you are sure you are replicating the intended listening effect
+     */
+    public V removeSilently(long k) {
+        return super.remove(k);
+    }
+
     public interface Callback<V> {
         void apply(long key, V value);
     }

common/src/main/java/net/caffeinemc/mods/lithium/common/world/interests/PoiCheckConsistency.java

@@ -0,0 +1,11 @@
+package net.caffeinemc.mods.lithium.common.world.interests;
+
+import net.minecraft.core.SectionPos;
+import net.minecraft.world.level.chunk.LevelChunkSection;
+
+import java.util.BitSet;
+
+public interface PoiCheckConsistency<R> {
+    R lithium$createSectionAndInsert(long l);
+    void lithium$CheckConsistencyWithBlocks(SectionPos sectionPos, LevelChunkSection levelChunkSection, BitSet column);
+}

common/src/main/java/net/caffeinemc/mods/lithium/common/world/interests/PoiUnloading.java

@@ -0,0 +1,8 @@
+package net.caffeinemc.mods.lithium.common.world.interests;
+
+import net.minecraft.world.level.ChunkPos;
+
+public interface PoiUnloading {
+    boolean lithium$shouldUnloadChunkPOIs(long chunkPos);
+    void lithium$unloadChunkPOIs(long chunkPos);
+}

common/src/main/java/net/caffeinemc/mods/lithium/common/world/interests/RegionBasedStorageSectionExtended.java

@@ -3,13 +3,20 @@
 import net.caffeinemc.mods.lithium.common.util.functions.FunLongAnd5;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Holder;
+import net.minecraft.world.level.ChunkPos;
+import org.jspecify.annotations.Nullable;
 
 import java.util.BitSet;
 import java.util.Optional;
 import java.util.function.Predicate;
 
 public interface RegionBasedStorageSectionExtended<R> {
 
+    @Nullable
+    public BitSet lithium$getColumn(long chunkPos);
+
+    BitSet lithium$getOrAddColumnIfNull(long chunkPos);
+
     <S, T, U> U lithium$getFirstInRangeInChunkColumn(int chunkX, int chunkZ,
                                                      long deltaYSqMargin,
                                                      BlockPos center, long radiusSq,
@@ -27,9 +34,24 @@ public interface RegionBasedStorageSectionExtended<R> {
 
     BitSet lithium$getNonEmptyPOISections(int chunkX, int chunkZ);
 
+    /**
+     * Manually remove a chunk section column bitset - only used when unloading chunk with no POISections
+     *
+     * @param chunkPos
+     * @return
+     */
+    BitSet lithium$removeColumn(long chunkPos);
+
+    /**
+     * Remove a section from the storage without updating the columns map - column adjustment must be handled manually
+     * @param l
+     * @return
+     */
+    Optional<R> lithium$removeSectionWithoutUpdatingColumn(long l);
+
     int lithium$getChunkYMin();
 
     int lithium$getChunkYMaxInclusive();
 
-    Optional<R> lithium$getElementAt(long sectionPos);
+    Optional<R> lithium$uncheckedGetElementAt(long sectionPos);
 }

common/src/main/java/net/caffeinemc/mods/lithium/common/world/interests/iterator/NearbyPointOfInterestStream.java

@@ -250,7 +250,7 @@ public boolean tryAdvance(Consumer<? super PoiRecord> action) {
             while (!this.isSectionListEmpty() && this.minCollectedElementDistanceSq >= this.getMinimumNextPotentialDistanceSq()) {
                 final long sectionPos = queuedPOISections.get(this.queuedSectionsSearched++).sectionPos;
                 this.nextSectionDistanceSq = this.getNextSectionDistanceSq();
-                final Optional<PoiSection> poiSection = this.storage.lithium$getElementAt(sectionPos);
+                final Optional<PoiSection> poiSection = this.storage.lithium$uncheckedGetElementAt(sectionPos);
                 //noinspection OptionalIsPresent
                 if (poiSection.isPresent()) {
                     ((PointOfInterestSetExtended) poiSection.get()).lithium$collectMatchingPoints(this.typeSelector, this.occupationStatus, this);

common/src/main/java/net/caffeinemc/mods/lithium/common/world/interests/iterator/SingleChunkPointOfInterestStream.java

@@ -0,0 +1,84 @@
+package net.caffeinemc.mods.lithium.common.world.interests.iterator;
+
+import net.caffeinemc.mods.lithium.common.world.interests.PointOfInterestSetExtended;
+import net.caffeinemc.mods.lithium.common.world.interests.RegionBasedStorageSectionExtended;
+import net.minecraft.core.Holder;
+import net.minecraft.core.SectionPos;
+import net.minecraft.world.entity.ai.village.poi.PoiManager;
+import net.minecraft.world.entity.ai.village.poi.PoiRecord;
+import net.minecraft.world.entity.ai.village.poi.PoiSection;
+import net.minecraft.world.entity.ai.village.poi.PoiType;
+import net.minecraft.world.level.ChunkPos;
+
+import java.util.BitSet;
+import java.util.Iterator;
+import java.util.Spliterator;
+import java.util.Spliterators;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
+
+/**
+ * Stream-less fallback for PoiManager::getInChunk in case search is not otherwise optimized
+ * @author jcw780
+ */
+public class SingleChunkPointOfInterestStream extends Spliterators.AbstractSpliterator<PoiRecord> {
+    final int chunkX;
+    final int chunkZ;
+    final int chunkYMin;
+    RegionBasedStorageSectionExtended<PoiSection> storage;
+    final BitSet chunkPoiSections;
+    int nextPoiSectionIndex;
+    Iterator<PoiRecord> sectionIterator;
+    private final Predicate<Holder<PoiType>> typeFilter;
+    private final PoiManager.Occupancy status;
+
+    public SingleChunkPointOfInterestStream(Predicate<Holder<PoiType>> typeFilter, ChunkPos chunkPos, PoiManager.Occupancy status, RegionBasedStorageSectionExtended<PoiSection> storage) {
+        super(Long.MAX_VALUE, Spliterator.ORDERED);
+        this.chunkX = chunkPos.x;
+        this.chunkZ = chunkPos.z;
+        this.storage = storage;
+
+        this.chunkYMin = this.storage.lithium$getChunkYMin();
+
+        this.chunkPoiSections = this.storage.lithium$getNonEmptyPOISections(this.chunkX, this.chunkZ);
+        this.typeFilter = typeFilter;
+        this.status = status;
+    }
+
+    @Override
+    public boolean tryAdvance(Consumer<? super PoiRecord> action) {
+        do {
+            while (this.sectionIterator != null && this.sectionIterator.hasNext()) {
+                PoiRecord next = this.sectionIterator.next();
+                if (status.getTest().test(next)) {
+                    action.accept(next);
+                    return true;
+                }
+            }
+
+        } while (this.nextSection());
+        return false;
+    }
+
+    private boolean nextSection() {
+        int nextSectionIndex;
+        while ((nextSectionIndex = this.chunkPoiSections.nextSetBit(this.nextPoiSectionIndex)) != -1) {
+            this.nextPoiSectionIndex = nextSectionIndex + 1;
+            int chunkY = nextSectionIndex + this.chunkYMin;
+            this.sectionIterator = this.getSectionIterator(this.chunkX, chunkY, this.chunkZ);
+            if (this.sectionIterator != null && this.sectionIterator.hasNext()) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private Iterator<PoiRecord> getSectionIterator(int chunkX, int chunkY, int chunkZ) {
+        // Note: chunk is already "POI loaded" by lithium$getNonEmptyPOISections
+        PoiSection poiSection = this.storage.lithium$uncheckedGetElementAt(SectionPos.asLong(chunkX, chunkY, chunkZ)).orElse(null);
+        if (poiSection == null) {
+            return null;
+        }
+        return ((PointOfInterestSetExtended) poiSection).lithium$iterate(this.typeFilter);
+    }
+}

common/src/main/java/net/caffeinemc/mods/lithium/mixin/ai/poi/PoiManagerMixin.java

@@ -8,17 +8,20 @@
 import net.caffeinemc.mods.lithium.common.world.interests.PointOfInterestStorageExtended;
 import net.caffeinemc.mods.lithium.common.world.interests.RegionBasedStorageSectionExtended;
 import net.caffeinemc.mods.lithium.common.world.interests.iterator.NearbyPointOfInterestStream;
+import net.caffeinemc.mods.lithium.common.world.interests.iterator.SingleChunkPointOfInterestStream;
 import net.caffeinemc.mods.lithium.common.world.interests.iterator.SinglePointOfInterestTypeFilter;
 import net.caffeinemc.mods.lithium.common.world.interests.iterator.SphereChunkOrderedPoiSetSpliterator;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Holder;
 import net.minecraft.core.RegistryAccess;
 import net.minecraft.core.SectionPos;
 import net.minecraft.util.RandomSource;
+import net.minecraft.util.VisibleForDebug;
 import net.minecraft.world.entity.ai.village.poi.PoiManager;
 import net.minecraft.world.entity.ai.village.poi.PoiRecord;
 import net.minecraft.world.entity.ai.village.poi.PoiSection;
 import net.minecraft.world.entity.ai.village.poi.PoiType;
+import net.minecraft.world.level.ChunkPos;
 import net.minecraft.world.level.LevelHeightAccessor;
 import net.minecraft.world.level.border.WorldBorder;
 import net.minecraft.world.level.chunk.storage.ChunkIOErrorReporter;
@@ -240,6 +243,16 @@ public Stream<PoiRecord> getInRange(Predicate<Holder<PoiType>> predicate, BlockP
         return StreamSupport.stream(new SphereChunkOrderedPoiSetSpliterator(radius, center, this, predicate, status), false);
     }
 
+    /**
+     * @author jcw780
+     * @reason Fall back to column checked accesses in case POI search is otherwise unoptimized
+     */
+    @VisibleForDebug
+    @Overwrite
+    public Stream<PoiRecord> getInChunk(Predicate<Holder<PoiType>> predicate, ChunkPos chunkPos, PoiManager.Occupancy occupancy) {
+        return StreamSupport.stream(new SingleChunkPointOfInterestStream(predicate, chunkPos, occupancy, this), false);
+    }
+
     @Override
     public Optional<PoiRecord> lithium$findNearestForPortalLogic(BlockPos origin, int radius, Holder<PoiType> type,
                                                                  PoiManager.Occupancy status,

common/src/main/java/net/caffeinemc/mods/lithium/mixin/ai/poi/SectionStorageMixin.java

@@ -1,14 +1,20 @@
 package net.caffeinemc.mods.lithium.mixin.ai.poi;
 
 import com.google.common.collect.AbstractIterator;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import com.mojang.serialization.Codec;
+import com.mojang.serialization.DataResult;
+import com.mojang.serialization.Dynamic;
+import com.mojang.serialization.DynamicOps;
 import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
 import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
 import net.caffeinemc.mods.lithium.common.util.Distances;
 import net.caffeinemc.mods.lithium.common.util.Pos;
 import net.caffeinemc.mods.lithium.common.util.collections.ListeningLong2ObjectOpenHashMap;
 import net.caffeinemc.mods.lithium.common.util.functions.FunLongAnd5;
 import net.caffeinemc.mods.lithium.common.world.interests.RegionBasedStorageSectionExtended;
+import net.minecraft.SharedConstants;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Holder;
 import net.minecraft.core.RegistryAccess;
@@ -18,40 +24,45 @@
 import net.minecraft.world.level.chunk.storage.ChunkIOErrorReporter;
 import net.minecraft.world.level.chunk.storage.SectionStorage;
 import net.minecraft.world.level.chunk.storage.SimpleRegionStorage;
-import org.spongepowered.asm.mixin.Final;
-import org.spongepowered.asm.mixin.Mixin;
-import org.spongepowered.asm.mixin.Mutable;
-import org.spongepowered.asm.mixin.Shadow;
+import org.jspecify.annotations.Nullable;
+import org.slf4j.Logger;
+import org.spongepowered.asm.mixin.*;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
-import java.util.BitSet;
-import java.util.Collections;
-import java.util.Objects;
-import java.util.Optional;
+import java.util.*;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Predicate;
 
 @SuppressWarnings("OptionalUsedAsFieldOrParameterType") // We don't get a choice, this is Minecraft's doing!
 @Mixin(SectionStorage.class)
-public abstract class SectionStorageMixin<R> implements RegionBasedStorageSectionExtended<R> {
+public abstract class SectionStorageMixin<R, P> implements RegionBasedStorageSectionExtended<R> {
     @Mutable
     @Shadow
     @Final
     private Long2ObjectMap<Optional<R>> storage;
 
-    @Shadow
-    protected abstract Optional<R> get(long pos);
-
     @Shadow
     @Final
     protected LevelHeightAccessor levelHeightAccessor;
 
     @Shadow
     protected abstract void unpackChunk(ChunkPos chunkPos);
 
+    @Shadow
+    @Final
+    private Codec<P> codec;
+
+    @Shadow
+    @Final
+    static Logger LOGGER;
+
+    @Shadow
+    @Final
+    private Function<R, P> packer;
+
     private Long2ObjectOpenHashMap<BitSet> columns;
 
     @SuppressWarnings("rawtypes")
@@ -105,6 +116,33 @@ private void onEntryAdded(long key, Optional<R> value) {
         flags.set(y, value.isPresent());
     }
 
+    @Override
+    @Nullable
+    public BitSet lithium$getColumn(long chunkPos) {
+        return this.columns.get(chunkPos);
+    }
+
+    @Override
+    public BitSet lithium$removeColumn(long chunkPos) {
+        return this.columns.remove(chunkPos);
+    }
+
+    @Override
+    public Optional<R> lithium$removeSectionWithoutUpdatingColumn (long sectionPos) {
+        return ((ListeningLong2ObjectOpenHashMap<Optional<R>>) this.storage).removeSilently(sectionPos);
+    }
+
+    @Override
+    public BitSet lithium$getOrAddColumnIfNull(long chunkPos) {
+        BitSet column = this.columns.get(chunkPos);
+
+        if (column == null) {
+            this.columns.put(chunkPos, column = new BitSet(Pos.SectionYIndex.getNumYSections(this.levelHeightAccessor)));
+        }
+
+        return column;
+    }
+
     @Override
     public <S, T, U> U lithium$getFirstInRangeInChunkColumn(int chunkX, int chunkZ,
                                                             long deltaYSqMargin,
@@ -182,7 +220,7 @@ protected R computeNext() {
     }
 
     @Override
-    public Optional<R> lithium$getElementAt(long sectionPos) {
+    public Optional<R> lithium$uncheckedGetElementAt(long sectionPos) {
         return this.storage.get(sectionPos);
     }
 
@@ -195,4 +233,55 @@ protected R computeNext() {
     public int lithium$getChunkYMaxInclusive() {
         return Pos.SectionYCoord.getMaxYSectionInclusive(this.levelHeightAccessor);
     }
+
+    /**
+     * These mixins are to remove Optional.empty() sections stored in the storage hashmap.
+     * This is done to mitigate a memory leak in vanilla as these are never unloaded. The long keys will eventually
+     * build up as more chunks have been loaded.
+     * The new logic utilizes the Lithium columns lookup for populated sections instead of Optional.empty() stored
+     * instead the storage hashmap.
+     */
+
+
+    /**
+     * @author jcw780
+     * @reason Use Lithium columns look up in write chunk
+     */
+    @Overwrite
+    private <T> com.mojang.serialization.Dynamic<T> writeChunk(ChunkPos chunkPos, DynamicOps<T> dynamicOps) {
+        Map<T, T> map = Maps.<T, T>newHashMap();
+        final int chunkX = chunkPos.x;
+        final int chunkZ = chunkPos.z;
+        BitSet sectionsWithPOI = this.columns.get(chunkPos.toLong());
+
+        if (sectionsWithPOI != null) {
+            int nextBit = sectionsWithPOI.nextSetBit(0);
+            while (nextBit >= 0) {
+                final int chunkY = Pos.SectionYCoord.fromSectionIndex(this.levelHeightAccessor, nextBit);
+                Optional<R> next = this.storage.get(SectionPos.asLong(chunkX, chunkY, chunkZ));
+
+                // Find and advance to the next set bit
+                nextBit = sectionsWithPOI.nextSetBit(nextBit + 1);
+
+                if (next.isPresent()) {
+                    DataResult<T> dataResult = this.codec.encodeStart(dynamicOps, (P)this.packer.apply(next.get()));
+                    String string = Integer.toString(chunkY);
+                    dataResult.resultOrPartial(LOGGER::error).ifPresent(object -> map.put(dynamicOps.createString(string), object));
+                }
+            }
+        }
+
+        return new Dynamic<>(
+                dynamicOps,
+                dynamicOps.createMap(
+                        ImmutableMap.of(
+                                dynamicOps.createString("Sections"),
+                                dynamicOps.createMap(map),
+                                dynamicOps.createString("DataVersion"),
+                                dynamicOps.createInt(SharedConstants.getCurrentVersion().dataVersion().version())
+                        )
+                )
+        );
+    }
+
 }