Optimize sorting tree container lookup

src/main/java/twilightforest/block/SortLogCoreBlock.java

@@ -14,6 +14,7 @@
 import net.neoforged.neoforge.transfer.ResourceHandlerUtil;
 import net.neoforged.neoforge.transfer.item.ItemResource;
 import net.neoforged.neoforge.transfer.transaction.Transaction;
+import org.jspecify.annotations.Nullable;
 import twilightforest.config.TFConfig;
 import twilightforest.init.TFParticleType;
 import twilightforest.network.ParticlePacket;
@@ -28,7 +29,7 @@
 
 public class SortLogCoreBlock extends SpecialMagicLogBlock {
 
-	private final BlockCapabilityDirectionalCache<ResourceHandler<ItemResource>> capabilityCache = new BlockCapabilityDirectionalCache<>();
+	private final BlockCapabilityDirectionalCache<ResourceHandler<ItemResource>> capabilityCache = new BlockCapabilityDirectionalCache<>(Capabilities.Item.BLOCK);
 
 	public SortLogCoreBlock(Properties properties) {
 		super(properties);
@@ -39,31 +40,28 @@ public boolean doesCoreFunction() {
 		return !TFConfig.disableSortingCore;
 	}
 
-	//TODO fuckkkkkkkkk
 	@Override
 	void performTreeEffect(ServerLevel level, BlockPos pos, RandomSource rand) {
 		Map<List<ResourceHandler<ItemResource>>, Vec3> inputMap = new HashMap<>();
 		Map<ResourceHandler<ItemResource>, Vec3> outputMap = new HashMap<>();
 
-		for (BlockPos blockPos : WorldUtil.getAllAround(pos, TFConfig.sortingCoreRange)) { // Get every itemHandler from every block in the area
+		for (BlockEntity blockEntity : WorldUtil.getLoadedBlockEntitiesInRange(level, pos, TFConfig.sortingCoreRange)) {
+			BlockPos blockPos = blockEntity.getBlockPos();
 			if (!blockPos.equals(pos)) {
-				BlockEntity blockEntity = level.getBlockEntity(blockPos);
-				if (blockEntity != null) {
-					// Put it in the input if its within 2 blocks
-					if (Math.abs(blockPos.getX() - pos.getX()) <= 2 && Math.abs(blockPos.getY() - pos.getY()) <= 2 && Math.abs(blockPos.getZ() - pos.getZ()) <= 2) {
-						List<ResourceHandler<ItemResource>> handlers = new ArrayList<>();
-						for (Direction side : Direction.values()) {
-							ResourceHandler<ItemResource> handler = this.capabilityCache.get(Capabilities.Item.BLOCK, level, blockPos, side);
-							if (handler != null) handlers.add(handler);
-						}
-						if (!handlers.isEmpty()) {
-							inputMap.put(handlers, Vec3.upFromBottomCenterOf(blockPos, 1.9D));
-						}
-					} else { // Output if its outside that range
-						for (Direction side : Direction.values()) {
-							ResourceHandler<ItemResource> handler = this.capabilityCache.get(Capabilities.Item.BLOCK, level, blockPos, side);
-							if (handler != null) outputMap.put(handler, Vec3.upFromBottomCenterOf(blockPos, 1.9D));
-						}
+				// Put it in the input if its within 2 blocks
+				if (Math.abs(blockPos.getX() - pos.getX()) <= 2 && Math.abs(blockPos.getY() - pos.getY()) <= 2 && Math.abs(blockPos.getZ() - pos.getZ()) <= 2) {
+					List<ResourceHandler<ItemResource>> handlers = new ArrayList<>();
+					for (Direction side : Direction.values()) {
+						ResourceHandler<ItemResource> handler = this.getItemHandler(level, blockEntity, side);
+						if (handler != null) handlers.add(handler);
+					}
+					if (!handlers.isEmpty()) {
+						inputMap.put(handlers, Vec3.upFromBottomCenterOf(blockPos, 1.9D));
+					}
+				} else { // Output if its outside that range
+					for (Direction side : Direction.values()) {
+						ResourceHandler<ItemResource> handler = this.getItemHandler(level, blockEntity, side);
+						if (handler != null) outputMap.put(handler, Vec3.upFromBottomCenterOf(blockPos, 1.9D));
 					}
 				}
 			}
@@ -129,7 +127,7 @@ void performTreeEffect(ServerLevel level, BlockPos pos, RandomSource rand) {
 								double x = diff.x - 0.25D + rand.nextDouble() * 0.5D;
 								double y = diff.y - 1.75D + rand.nextDouble() * 0.5D;
 								double z = diff.z - 0.25D + rand.nextDouble() * 0.5D;
-								particlePacket.queueParticle(TFParticleType.SORTING_PARTICLE.get(), false, xyz, new Vec3(x, y, z).scale(1D / diff.length()));
+								particlePacket.queueParticle(TFParticleType.SORTING_PARTICLE.get(), false, false, xyz, new Vec3(x, y, z).scale(1D / diff.length()));
 								PacketDistributor.sendToPlayersNear(level, null, xyz.x(), xyz.y(), xyz.z(), 64.0D, particlePacket);
 								break;
 							}
@@ -141,4 +139,13 @@ void performTreeEffect(ServerLevel level, BlockPos pos, RandomSource rand) {
 			}
 		}
 	}
+
+	@Nullable
+	ResourceHandler<ItemResource> getItemHandler(ServerLevel level, BlockEntity blockEntity, Direction side) {
+		return this.capabilityCache.get(level, blockEntity, side);
+	}
+
+	public void clearCapabilityCache(ServerLevel level) {
+		this.capabilityCache.clear(level);
+	}
 }

src/main/java/twilightforest/events/SortingTreeEvents.java

@@ -0,0 +1,24 @@
+package twilightforest.events;
+
+import net.minecraft.server.level.ServerLevel;
+import net.neoforged.neoforge.common.NeoForge;
+import net.neoforged.neoforge.event.level.LevelEvent;
+import tamaized.beanification.Component;
+import tamaized.beanification.PostConstruct;
+import twilightforest.block.SortLogCoreBlock;
+import twilightforest.init.TFBlocks;
+
+@Component
+public final class SortingTreeEvents {
+
+	@PostConstruct
+	private void setup() {
+		NeoForge.EVENT_BUS.addListener(this::clearCapabilityCache);
+	}
+
+	private void clearCapabilityCache(LevelEvent.Unload event) {
+		if (event.getLevel() instanceof ServerLevel level) {
+			((SortLogCoreBlock) TFBlocks.SORTING_LOG_CORE.get()).clearCapabilityCache(level);
+		}
+	}
+}

src/main/java/twilightforest/util/BlockCapabilityDirectionalCache.java

@@ -3,29 +3,55 @@
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Direction;
 import net.minecraft.server.level.ServerLevel;
+import net.minecraft.world.level.block.entity.BlockEntity;
 import net.neoforged.neoforge.capabilities.BlockCapability;
 import net.neoforged.neoforge.capabilities.BlockCapabilityCache;
+import org.jspecify.annotations.Nullable;
 
-import javax.annotation.Nullable;
-import java.util.HashMap;
+import java.util.EnumMap;
+import java.util.IdentityHashMap;
 import java.util.Map;
+import java.util.WeakHashMap;
 
-// Voidscape copy
-public class BlockCapabilityDirectionalCache<R> {
+/**
+ * Caches sided block capabilities for loaded block entities, separated by level.
+ *
+ * <p>Entries within a level are weakly keyed by the target block entity rather than globally keyed by position.
+ * This keeps identical positions in different levels separate and allows the directional caches to be collected
+ * after NeoForge invalidates them on a normal block or chunk unload.</p>
+ *
+ * <p>{@link BlockCapabilityCache} retains its {@link ServerLevel}, and level unload does not fire the usual chunk
+ * invalidations. Call {@link #clear(ServerLevel)} from {@code LevelEvent.Unload} to release the level explicitly.</p>
+ */
+public final class BlockCapabilityDirectionalCache<R> {
 
-	private final Map<BlockPosAndDirection, BlockCapabilityCache<R, Direction>> data = new HashMap<>();
+	private final BlockCapability<R, @Nullable Direction> capability;
+	private final Map<ServerLevel, Map<BlockEntity, EnumMap<Direction, BlockCapabilityCache<R, @Nullable Direction>>>> data = new IdentityHashMap<>();
+
+	public BlockCapabilityDirectionalCache(BlockCapability<R, @Nullable Direction> capability) {
+		this.capability = capability;
+	}
 
 	@Nullable
-	public R get(BlockCapability<R, Direction> capability, ServerLevel level, BlockPos pos, Direction direction) {
-		BlockCapabilityCache<R, Direction> cache = this.data.get(new BlockPosAndDirection(pos, direction));
-		if (cache == null) {
-			cache = BlockCapabilityCache.create(capability, level, pos, direction);
-			this.data.put(new BlockPosAndDirection(pos, direction), cache);
+	public R get(ServerLevel level, BlockEntity blockEntity, Direction direction) {
+		Map<BlockEntity, EnumMap<Direction, BlockCapabilityCache<R, @Nullable Direction>>> levelCaches = this.data.computeIfAbsent(
+			level,
+			ignored -> new WeakHashMap<>()
+		);
+		EnumMap<Direction, BlockCapabilityCache<R, @Nullable Direction>> directionalCaches = levelCaches.computeIfAbsent(
+			blockEntity,
+			ignored -> new EnumMap<>(Direction.class)
+		);
+		BlockPos blockPos = blockEntity.getBlockPos();
+		BlockCapabilityCache<R, @Nullable Direction> cache = directionalCaches.get(direction);
+		if (cache == null || cache.level() != level || !cache.pos().equals(blockPos)) {
+			cache = BlockCapabilityCache.create(this.capability, level, blockPos, direction);
+			directionalCaches.put(direction, cache);
 		}
 		return cache.getCapability();
 	}
 
-	private record BlockPosAndDirection(BlockPos pos, Direction direction) {
-
+	public void clear(ServerLevel level) {
+		this.data.remove(level);
 	}
 }

src/main/java/twilightforest/util/WorldUtil.java

@@ -10,6 +10,7 @@
 import net.minecraft.core.Holder;
 import net.minecraft.core.HolderSet;
 import net.minecraft.core.RegistryAccess;
+import net.minecraft.core.SectionPos;
 import net.minecraft.server.level.ServerChunkCache;
 import net.minecraft.server.level.ServerLevel;
 import net.minecraft.util.Mth;
@@ -20,6 +21,7 @@
 import net.minecraft.world.level.LevelAccessor;
 import net.minecraft.world.level.StructureManager;
 import net.minecraft.world.level.biome.Biome;
+import net.minecraft.world.level.block.entity.BlockEntity;
 import net.minecraft.world.level.chunk.ChunkGeneratorStructureState;
 import net.minecraft.world.level.levelgen.Heightmap;
 import net.minecraft.world.level.levelgen.structure.Structure;
@@ -61,6 +63,35 @@ public static Iterable<BlockPos> getAllAround(BlockPos center, int range) {
 		return BlockPos.betweenClosed(center.offset(-range, -range, -range), center.offset(range, range, range));
 	}
 
+	/**
+	 * Finds block entities in loaded chunks within a cube around {@code center}, inclusive of its edges.
+	 * This avoids loading chunks or scanning every block position in the cube.
+	 */
+	public static List<BlockEntity> getLoadedBlockEntitiesInRange(ServerLevel level, BlockPos center, int range) {
+		int minX = center.getX() - range;
+		int minY = center.getY() - range;
+		int minZ = center.getZ() - range;
+		int maxX = center.getX() + range;
+		int maxY = center.getY() + range;
+		int maxZ = center.getZ() + range;
+		ChunkPos minChunk = new ChunkPos(SectionPos.blockToSectionCoord(minX), SectionPos.blockToSectionCoord(minZ));
+		ChunkPos maxChunk = new ChunkPos(SectionPos.blockToSectionCoord(maxX), SectionPos.blockToSectionCoord(maxZ));
+		ServerChunkCache chunkSource = level.getChunkSource();
+
+		return ChunkPos.rangeClosed(minChunk, maxChunk)
+			.map(chunkPos -> chunkSource.getChunkNow(chunkPos.x(), chunkPos.z()))
+			.filter(Objects::nonNull)
+			.flatMap(chunk -> chunk.getBlockEntities().values().stream())
+			.filter(blockEntity -> !blockEntity.isRemoved())
+			.filter(blockEntity -> {
+				BlockPos blockPos = blockEntity.getBlockPos();
+				return blockPos.getX() >= minX && blockPos.getX() <= maxX
+					&& blockPos.getY() >= minY && blockPos.getY() <= maxY
+					&& blockPos.getZ() >= minZ && blockPos.getZ() <= maxZ;
+			})
+			.toList();
+	}
+
 	/**
 	 * Floors both corners of the bounding box to integers
 	 * Inclusive of edges