From 2afdaa1ac94c8acaa749637baaf019223d16bde2 Mon Sep 17 00:00:00 2001
From: Brady <zeromemesdev@gmail.com>
Date: Fri, 14 Dec 2018 18:50:39 -0600
Subject: [PATCH] Add individual chunk scanning to IWorldScanner

---
 .../baritone/api/cache/IWorldScanner.java     | 15 +++++++
 .../java/baritone/cache/WorldScanner.java     | 43 +++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/src/api/java/baritone/api/cache/IWorldScanner.java b/src/api/java/baritone/api/cache/IWorldScanner.java
index 6d6f49ef..caa44cbc 100644
--- a/src/api/java/baritone/api/cache/IWorldScanner.java
+++ b/src/api/java/baritone/api/cache/IWorldScanner.java
@@ -20,6 +20,7 @@ package baritone.api.cache;
 import baritone.api.utils.IPlayerContext;
 import net.minecraft.block.Block;
 import net.minecraft.util.math.BlockPos;
+import net.minecraft.util.math.ChunkPos;
 
 import java.util.List;
 
@@ -42,4 +43,18 @@ public interface IWorldScanner {
      * @return The matching block positions
      */
     List<BlockPos> scanChunkRadius(IPlayerContext ctx, List<Block> blocks, int max, int yLevelThreshold, int maxSearchRadius);
+
+    /**
+     * Scans a single chunk for the specified blocks.
+     *
+     * @param ctx             The {@link IPlayerContext} containing player and world info that the
+     *                        scan is based upon
+     * @param blocks          The blocks to scan for
+     * @param pos             The position of the target chunk
+     * @param max             The maximum number of blocks to scan before cutoff
+     * @param yLevelThreshold If a block is found within this Y level, the current result will be
+     *                        returned, if the value is negative, then this condition doesn't apply.
+     * @return The matching block positions
+     */
+    List<BlockPos> scanChunk(IPlayerContext ctx, List<Block> blocks, ChunkPos pos, int max, int yLevelThreshold);
 }
diff --git a/src/main/java/baritone/cache/WorldScanner.java b/src/main/java/baritone/cache/WorldScanner.java
index 1d220188..30b67ded 100644
--- a/src/main/java/baritone/cache/WorldScanner.java
+++ b/src/main/java/baritone/cache/WorldScanner.java
@@ -23,10 +23,12 @@ import net.minecraft.block.Block;
 import net.minecraft.block.state.IBlockState;
 import net.minecraft.client.multiplayer.ChunkProviderClient;
 import net.minecraft.util.math.BlockPos;
+import net.minecraft.util.math.ChunkPos;
 import net.minecraft.world.chunk.BlockStateContainer;
 import net.minecraft.world.chunk.Chunk;
 import net.minecraft.world.chunk.storage.ExtendedBlockStorage;
 
+import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
 
@@ -107,4 +109,45 @@ public enum WorldScanner implements IWorldScanner {
             searchRadiusSq++;
         }
     }
+
+    @Override
+    public List<BlockPos> scanChunk(IPlayerContext ctx, List<Block> blocks, ChunkPos pos, int max, int yLevelThreshold) {
+        if (blocks.isEmpty()) {
+            return Collections.emptyList();
+        }
+
+        ChunkProviderClient chunkProvider = (ChunkProviderClient) ctx.world().getChunkProvider();
+        Chunk chunk = chunkProvider.getLoadedChunk(pos.x, pos.z);
+        int playerY = ctx.playerFeet().getY();
+
+        if (chunk == null || chunk.isEmpty()) {
+            return Collections.emptyList();
+        }
+
+        LinkedList<BlockPos> res = new LinkedList<>();
+        ExtendedBlockStorage[] chunkInternalStorageArray = chunk.getBlockStorageArray();
+        for (int y0 = 0; y0 < 16; y0++) {
+            ExtendedBlockStorage extendedblockstorage = chunkInternalStorageArray[y0];
+            if (extendedblockstorage == null) {
+                continue;
+            }
+            int yReal = y0 << 4;
+            BlockStateContainer bsc = extendedblockstorage.getData();
+            for (int y = 0; y < 16; y++) {
+                for (int z = 0; z < 16; z++) {
+                    for (int x = 0; x < 16; x++) {
+                        IBlockState state = bsc.get(x, y, z);
+                        if (blocks.contains(state.getBlock())) {
+                            int yy = yReal | y;
+                            res.add(new BlockPos((pos.x << 4) | x, yy, (pos.z << 4) | z));
+                            if (res.size() >= max || Math.abs(yy - playerY) < yLevelThreshold) {
+                                return res;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        return res;
+    }
 }