Optimize Sorting Tree inventory discovery

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

@@ -16,15 +16,12 @@
 import twilightforest.data.tags.EntityTagGenerator;
 import twilightforest.init.TFParticleType;
 import twilightforest.network.ParticlePacket;
-import twilightforest.util.BlockCapabilityDirectionalCache;
 import twilightforest.util.WorldUtil;
 
 import java.util.*;
 
 public class SortLogCoreBlock extends SpecialMagicLogBlock {
 
-	private final BlockCapabilityDirectionalCache<IItemHandler> capabilityCache = new BlockCapabilityDirectionalCache<>();
-
 	public SortLogCoreBlock(Properties properties) {
 		super(properties);
 	}
@@ -39,25 +36,23 @@ void performTreeEffect(ServerLevel level, BlockPos pos, RandomSource rand) {
 		Map<List<IItemHandler>, Vec3> inputMap = new HashMap<>();
 		Map<IItemHandler, 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).toList()) {
+			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<IItemHandler> handlers = new ArrayList<>();
-						for (Direction side : Direction.values()) {
-							IItemHandler handler = this.capabilityCache.get(Capabilities.ItemHandler.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()) {
-							IItemHandler handler = this.capabilityCache.get(Capabilities.ItemHandler.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<IItemHandler> handlers = new ArrayList<>();
+					for (Direction side : Direction.values()) {
+						IItemHandler handler = level.getCapability(Capabilities.ItemHandler.BLOCK, blockPos, blockEntity.getBlockState(), 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()) {
+						IItemHandler handler = level.getCapability(Capabilities.ItemHandler.BLOCK, blockPos, blockEntity.getBlockState(), blockEntity, side);
+						if (handler != null) outputMap.put(handler, Vec3.upFromBottomCenterOf(blockPos, 1.9D));
 					}
 				}
 			}

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

@@ -11,6 +11,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;
@@ -37,6 +39,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Stream;
 
 public final class WorldUtil {
 	private WorldUtil() {
@@ -61,6 +64,34 @@ 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 Stream<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;
+			});
+	}
+
 	/**
 	 * Floors both corners of the bounding box to integers
 	 * Inclusive of edges

src/test/java/twilightforest/block/SortLogCoreBlockTests.java

@@ -0,0 +1,120 @@
+package twilightforest.block;
+
+import net.minecraft.core.BlockPos;
+import net.minecraft.core.Direction;
+import net.minecraft.server.level.ServerChunkCache;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.util.RandomSource;
+import net.minecraft.world.entity.Entity;
+import net.minecraft.world.item.ItemStack;
+import net.minecraft.world.item.Items;
+import net.minecraft.world.level.block.entity.BlockEntity;
+import net.minecraft.world.level.block.state.BlockState;
+import net.minecraft.world.level.chunk.LevelChunk;
+import net.minecraft.world.phys.AABB;
+import net.neoforged.neoforge.capabilities.Capabilities;
+import net.neoforged.neoforge.items.ItemStackHandler;
+import net.neoforged.neoforge.network.PacketDistributor;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.MockedStatic;
+import tamaized.beanification.junit.MockitoFixer;
+import twilightforest.config.TFConfig;
+import twilightforest.init.TFBlocks;
+import twilightforest.network.ParticlePacket;
+
+import java.util.List;
+import java.util.Map;
+import java.util.function.Predicate;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyDouble;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoFixer.class)
+public class SortLogCoreBlockTests {
+
+	@Test
+	public void transfersItemsAcrossAChunkBoundary() {
+		int originalRange = TFConfig.sortingCoreRange;
+		TFConfig.sortingCoreRange = 4;
+
+		try {
+			ServerLevel level = mock(ServerLevel.class);
+			ServerChunkCache chunkSource = mock(ServerChunkCache.class);
+			LevelChunk inputChunk = mock(LevelChunk.class);
+			LevelChunk outputChunk = mock(LevelChunk.class);
+			BlockPos corePos = new BlockPos(15, 64, 0);
+			BlockPos inputPos = new BlockPos(16, 64, 0);
+			BlockPos outputPos = new BlockPos(11, 64, 0);
+			BlockEntity inputBlockEntity = blockEntityAt(inputPos);
+			BlockEntity outputBlockEntity = blockEntityAt(outputPos);
+			BlockState inputState = inputBlockEntity.getBlockState();
+			BlockState outputState = outputBlockEntity.getBlockState();
+			ItemStackHandler input = new ItemStackHandler(1);
+			ItemStackHandler output = new ItemStackHandler(1);
+			input.setStackInSlot(0, new ItemStack(Items.DIAMOND, 2));
+			output.setStackInSlot(0, new ItemStack(Items.DIAMOND));
+
+			when(level.getChunkSource()).thenReturn(chunkSource);
+			when(chunkSource.getChunkNow(0, 0)).thenReturn(outputChunk);
+			when(chunkSource.getChunkNow(1, 0)).thenReturn(inputChunk);
+			when(inputChunk.getBlockEntities()).thenReturn(Map.of(inputPos, inputBlockEntity));
+			when(outputChunk.getBlockEntities()).thenReturn(Map.of(outputPos, outputBlockEntity));
+			when(level.getCapability(
+				eq(Capabilities.ItemHandler.BLOCK),
+				eq(inputPos),
+				same(inputState),
+				same(inputBlockEntity),
+				any(Direction.class)
+			)).thenReturn(input);
+			when(level.getCapability(
+				eq(Capabilities.ItemHandler.BLOCK),
+				eq(outputPos),
+				same(outputState),
+				same(outputBlockEntity),
+				any(Direction.class)
+			)).thenReturn(output);
+			when(level.getEntities(
+				isNull(Entity.class),
+				any(AABB.class),
+				org.mockito.ArgumentMatchers.<Predicate<? super Entity>>any()
+			)).thenReturn(List.of());
+
+			try (MockedStatic<PacketDistributor> packets = mockStatic(PacketDistributor.class)) {
+				((SortLogCoreBlock) TFBlocks.SORTING_LOG_CORE.get()).performTreeEffect(level, corePos, RandomSource.create(0L));
+				packets.verify(() -> PacketDistributor.sendToPlayersNear(
+					same(level), isNull(), anyDouble(), anyDouble(), anyDouble(), anyDouble(), any(ParticlePacket.class)
+				));
+			}
+
+			assertEquals(1, input.getStackInSlot(0).getCount());
+			assertEquals(2, output.getStackInSlot(0).getCount());
+			verify(level, times(Direction.values().length)).getCapability(
+				eq(Capabilities.ItemHandler.BLOCK), eq(inputPos), same(inputState), same(inputBlockEntity), any(Direction.class)
+			);
+			verify(level, times(Direction.values().length)).getCapability(
+				eq(Capabilities.ItemHandler.BLOCK), eq(outputPos), same(outputState), same(outputBlockEntity), any(Direction.class)
+			);
+		} finally {
+			TFConfig.sortingCoreRange = originalRange;
+		}
+	}
+
+	private static BlockEntity blockEntityAt(BlockPos pos) {
+		BlockEntity blockEntity = mock(BlockEntity.class);
+		BlockState state = mock(BlockState.class);
+		when(blockEntity.getBlockPos()).thenReturn(pos);
+		when(blockEntity.getBlockState()).thenReturn(state);
+		when(blockEntity.isRemoved()).thenReturn(false);
+		return blockEntity;
+	}
+}