BlockESP.java

package io.github.ensamisten.client.module.render;

import com.google.gson.JsonObject;
import com.mojang.blaze3d.vertex.PoseStack;
import io.github.ensamisten.client.event.EventRegistry;
import io.github.ensamisten.client.event.impl.block.BlockEvent;
import io.github.ensamisten.client.event.impl.render.RenderEvent;
import io.github.ensamisten.client.event.impl.tick.TickEvent;
import io.github.ensamisten.client.gui.color.Color;
import io.github.ensamisten.client.module.Extension;
import io.github.ensamisten.client.module.ModuleEntrypoint;
import io.github.ensamisten.client.setting.Setting;
import io.github.ensamisten.client.setting.SettingFactory;
import io.github.ensamisten.client.util.Render3D;
import net.minecraft.client.Camera;
import net.minecraft.client.Minecraft;
import net.minecraft.core.BlockPos;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.resources.Identifier;
import net.minecraft.world.level.ChunkPos;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.Blocks;
import net.minecraft.world.level.chunk.ChunkAccess;
import net.minecraft.world.level.chunk.LevelChunk;
import net.minecraft.world.level.chunk.status.ChunkStatus;
import net.minecraft.world.phys.AABB;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.function.Consumer;

public class BlockESP extends Extension implements ModuleEntrypoint {

    private static final Minecraft MC = Minecraft.getInstance();

    // fields
    private static final int MAX_SCAN_RADIUS = 8;   // hard safety cap
    private static final int CHUNKS_PER_TICK = 2;   // small budget
    private final Setting<Float> colorRed   = SettingFactory.floatSetting("blockesp_color_red",   "Red",   0.0f);
    private final Setting<Float> colorGreen = SettingFactory.floatSetting("blockesp_color_green", "Green", 1.0f);
    private final Setting<Float> colorBlue  = SettingFactory.floatSetting("blockesp_color_blue",  "Blue",  1.0f);
    private final Setting<Float> colorAlpha = SettingFactory.floatSetting("blockesp_color_alpha", "Alpha", 0.3f);

    private final Setting<Boolean> visibilityToggle = SettingFactory.booleanSetting("blockesp_visibility", "Visibility", true);
    private final Setting<Float>   lineThickness    = SettingFactory.floatSetting("blockesp_line_thickness", "Line Thickness", 2.0f);
    private final Set<BlockPos> cachedSet = new HashSet<>();
    private final Setting<String> targetBlocks =
            SettingFactory.stringSetting("blockesp_blocks", "Blocks", "minecraft:spawner");

    // Parsed selection
    private final Set<Block> targetSet = new HashSet<>();
    private String lastParsed = null;

    // Aoba-style chunk scanning
    private final LinkedHashSet<ChunkPos> chunkQueue = new LinkedHashSet<>();
    private final Map<ChunkPos, ArrayList<BlockPos>> blockPositions = new HashMap<>();

    private final Consumer<TickEvent>   tickListener   = this::onTick;
    private final Consumer<RenderEvent> renderListener = this::onRender;
    private final Consumer<BlockEvent> blockListener = this::onBlock;
    public BlockESP() {
        super("BlockESP", Categories.Render);
        setDescription("Allows the player to see blocks with an ESP.");

        addSetting(colorRed);
        addSetting(colorGreen);
        addSetting(colorBlue);
        addSetting(colorAlpha);
        addSetting(visibilityToggle);
        addSetting(lineThickness);
        addSetting(targetBlocks);
    }

    @Override
    public void onEnable() {
        EventRegistry.TICK_EVENT.register(tickListener);
        EventRegistry.RENDER_EVENT.register(renderListener);
        EventRegistry.BLOCK_EVENT.register(blockListener);
        lastParsed = null;
        rebuildTargetsIfNeeded(); // builds targetSet + queues
        warmScanVisibleChunks(64); // budgeted immediate population
    }

    @Override
    public void onDisable() {
        EventRegistry.TICK_EVENT.unregister(tickListener);
        EventRegistry.RENDER_EVENT.unregister(renderListener);
        EventRegistry.BLOCK_EVENT.unregister(blockListener);
        chunkQueue.clear();
        blockPositions.clear();
    }

    private void onBlock(BlockEvent e) {
        if (MC.level == null || targetSet.isEmpty()) return;

        BlockPos pos = e.pos();
        ChunkPos cp = ChunkPos.containing(pos);

        boolean wasTarget = targetSet.contains(e.oldState().getBlock());
        boolean isTarget  = targetSet.contains(e.newState().getBlock());

        if (!wasTarget && !isTarget) return;

        ArrayList<BlockPos> list = blockPositions.computeIfAbsent(cp, k -> new ArrayList<>());

        if (isTarget) {
            BlockPos p = pos.immutable();
            if (cachedSet.add(p)) list.add(p);
        } else {
            list.remove(pos);
            cachedSet.remove(pos);
            if (list.isEmpty()) blockPositions.remove(cp);
        }
    }

    // ── Render (cheap: just iterate cached positions) ─────────────
    // ── Render (cheap: just iterate cached positions) ─────────────
    public void onRender(RenderEvent event) {
        if (event.phase() != RenderEvent.Phase.DRAW) return;
        if (!visibilityToggle.getValue()) return;
        if (MC.player == null || MC.level == null) return;

        PoseStack matrixStack = event.poseStack();
        Camera camera = event.camera();
        Color boxColor = getColor();
        float thickness = Math.clamp(lineThickness.getValue(), 0f, 5f);

        var chunkIt = blockPositions.entrySet().iterator();
        while (chunkIt.hasNext()) {
            var entry = chunkIt.next();
            ArrayList<BlockPos> positions = entry.getValue();

            var posIt = positions.iterator();
            while (posIt.hasNext()) {
                BlockPos pos = posIt.next();

                // instant stale cleanup
                if (!targetSet.contains(MC.level.getBlockState(pos).getBlock())) {
                    posIt.remove();
                    cachedSet.remove(pos);
                    continue;
                }
                Render3D.draw3DBox(matrixStack, camera, new AABB(pos), boxColor, thickness);
            }

            if (positions.isEmpty()) {
                chunkIt.remove();
            }
        }
    }

    // ── Tick: scan ONE chunk per tick (Aoba's approach) ───────────
    public void onTick(TickEvent event) {
        if (event.phase() != TickEvent.Phase.PRE) return;
        if (MC.level == null) return;

        rebuildTargetsIfNeeded();
        discoverNearbyInstant();
        if (targetSet.isEmpty() || chunkQueue.isEmpty()) return;

        int budget = CHUNKS_PER_TICK;
        while (budget-- > 0 && !chunkQueue.isEmpty()) {
            ChunkPos pos = chunkQueue.getFirst();
            chunkQueue.remove(pos);

            ChunkAccess access = MC.level.getChunk(pos.x(), pos.z(), ChunkStatus.FULL, false);            if (!(access instanceof LevelChunk chunk) || chunk.isEmpty()) continue;

            ArrayList<BlockPos> list = new ArrayList<>();
            chunk.findBlocks(
                    state -> targetSet.contains(state.getBlock()),
                    (bp, state) -> list.add(bp.immutable())
            );
            ArrayList<BlockPos> old = blockPositions.remove(pos);
            if (old != null) {
                for (BlockPos p : old) cachedSet.remove(p);
            }
            if (!list.isEmpty()) {
                blockPositions.put(pos, list);
                cachedSet.addAll(list);
            }
        }
    }
    // and call chunkQueue.add(pos) / remove(pos)+blockPositions.remove(pos),
    // mirroring Aoba's onChunkLoaded/onChunkUnloaded. Otherwise, we queue all
    // currently-loaded chunks on enable / when the block set changes.
    private void queueAllLoadedChunks() {
        if (MC.level == null || MC.player == null) return;

        int cx = MC.player.chunkPosition().x();   // field, not x()
        int cz = MC.player.chunkPosition().z();   // field, not z()

        int viewDist = Math.min(MC.options.getEffectiveRenderDistance(), MAX_SCAN_RADIUS);

        chunkQueue.clear();

        for (int dx = -viewDist; dx <= viewDist; dx++) {
            for (int dz = -viewDist; dz <= viewDist; dz++) {
                int x = cx + dx;
                int z = cz + dz;
                ChunkAccess access = MC.level.getChunk(x, z, ChunkStatus.FULL, false);
                if (access instanceof LevelChunk) {
                    chunkQueue.add(new ChunkPos(x, z));
                }
            }
        }
    }

    private void warmScanVisibleChunks(int budget) {
        if (MC.level == null || MC.player == null || targetSet.isEmpty()) return;

        ChunkPos center = MC.player.chunkPosition();
        int r = Math.min(2, MAX_SCAN_RADIUS); // 5x5 around player for instant feel

        for (int dx = -r; dx <= r && budget > 0; dx++) {
            for (int dz = -r; dz <= r && budget > 0; dz++) {
                int x = center.x() + dx;
                int z = center.z() + dz;

                ChunkAccess access = MC.level.getChunk(x, z, ChunkStatus.FULL, false);
                if (!(access instanceof LevelChunk chunk) || chunk.isEmpty()) continue;

                ChunkPos cp = new ChunkPos(x, z);
                ArrayList<BlockPos> list = new ArrayList<>();

                chunk.findBlocks(
                        state -> targetSet.contains(state.getBlock()),
                        (bp, state) -> {
                            BlockPos p = bp.immutable();
                            list.add(p);
                            cachedSet.add(p);
                        }
                );

                if (!list.isEmpty()) blockPositions.put(cp, list);
                budget--;
            }
        }
    }

    private void onBlocksChanged() {
        chunkQueue.clear();
        blockPositions.clear();
        cachedSet.clear();
        queueAllLoadedChunks();
    }
    private void rebuildTargetsIfNeeded() {
        String cur = targetBlocks.getValue();
        if (cur != null && cur.equals(lastParsed)) return;   // unchanged → no work

        lastParsed = cur;
        targetSet.clear();
        if (cur != null && !cur.isBlank()) {
            for (String id : cur.split(",")) {
                Block b = resolveBlock(id.trim());
                if (b != null) targetSet.add(b);
            }
        }
        onBlocksChanged();   // re-queue ONLY on real change
    }

    // ── Config-screen helpers ─────────────────────────────────────
    // ── Config-screen helpers ─────────────────────────────────────
    // Read-only — NO side effects. Safe for per-frame grid rendering.
    public boolean isSelected(Block b) {
        return targetSet.contains(b);
    }

    public void toggleBlock(Block b) {
        rebuildTargetsIfNeeded();   // sync targetSet to current string

        LinkedHashSet<Block> set = new LinkedHashSet<>(targetSet);
        if (!set.remove(b)) set.add(b);

        StringBuilder sb = new StringBuilder();
        for (Block blk : set) {
            Identifier id = BuiltInRegistries.BLOCK.getKey(blk);
            if (id == null) continue;
            if (!sb.isEmpty()) sb.append(',');
            sb.append(id);
        }
        targetBlocks.setValue(sb.toString());
        lastParsed = null;          // one reparse + re-queue next tick
    }

    public Setting<String> getTargetBlocks() { return targetBlocks; }

    private Block resolveBlock(String id) {
        if (id == null || id.isBlank()) return null;
        try {
            Identifier rl = Identifier.parse(id.trim());
            Block b = BuiltInRegistries.BLOCK.getValue(rl);
            System.out.println("[BlockESP] resolve '" + id + "' -> " + b
                    + " (isAir=" + (b == Blocks.AIR) + ")");
            if (b == Blocks.AIR && !id.equals("minecraft:air")) return null;
            return b;
        } catch (Exception e) {
            System.out.println("[BlockESP] resolve '" + id + "' threw " + e);
            return null;
        }
    }

    private void discoverNearbyInstant() {
        if (MC.level == null || MC.player == null || targetSet.isEmpty()) return;

        ChunkPos center = MC.player.chunkPosition();
        int r = 1; // 3x3 chunks around player = instant feel

        for (int dx = -r; dx <= r; dx++) {
            for (int dz = -r; dz <= r; dz++) {
                int x = center.x() + dx;
                int z = center.z() + dz;

                ChunkAccess access = MC.level.getChunk(x, z, ChunkStatus.FULL, false);
                if (!(access instanceof LevelChunk chunk) || chunk.isEmpty()) continue;

                ChunkPos cp = new ChunkPos(x, z);
                ArrayList<BlockPos> list = blockPositions.computeIfAbsent(cp, k -> new ArrayList<>());

                chunk.findBlocks(
                        state -> targetSet.contains(state.getBlock()),
                        (bp, state) -> {
                            BlockPos p = bp.immutable();
                            if (cachedSet.add(p)) list.add(p); // add only if new
                        }
                );
            }
        }
    }

    private Color getColor() {
        return new Color(colorRed.getValue(), colorGreen.getValue(),
                colorBlue.getValue(), colorAlpha.getValue());
    }

    @Override
    public void saveConfig(JsonObject obj) {
        super.saveConfig(obj);
        for (Setting<?> setting : getSettings()) {
            Object val = setting.getValue();
            if (val instanceof Float f) obj.addProperty(setting.getId(), f);
            else if (val instanceof Boolean b) obj.addProperty(setting.getId(), b);
            else if (val instanceof Integer i) obj.addProperty(setting.getId(), i);
            else if (val instanceof String s) obj.addProperty(setting.getId(), s);
            else if (val instanceof Enum<?> e) obj.addProperty(setting.getId(), e.name());
        }
    }

    @Override
    @SuppressWarnings({"unchecked", "rawtypes"})
    public void loadConfig(JsonObject obj) {
        super.loadConfig(obj);
        for (Setting<?> setting : getSettings()) {
            if (!obj.has(setting.getId())) continue;
            Object val = setting.getValue();
            if (val instanceof Float) ((Setting<Float>) setting).setValue(obj.get(setting.getId()).getAsFloat());
            else if (val instanceof Boolean) ((Setting<Boolean>) setting).setValue(obj.get(setting.getId()).getAsBoolean());
            else if (val instanceof Integer) ((Setting<Integer>) setting).setValue(obj.get(setting.getId()).getAsInt());
            else if (val instanceof String) ((Setting<String>) setting).setValue(obj.get(setting.getId()).getAsString());
            else if (val instanceof Enum) {
                String enumName = obj.get(setting.getId()).getAsString();
                Class<? extends Enum> enumClass = (Class<? extends Enum>) val.getClass();
                ((Setting<Enum>) setting).setValue(Enum.valueOf(enumClass, enumName));
            }
        }
        lastParsed = null;
    }

    public Setting<Float>   getColorRed()        { return colorRed; }
    public Setting<Float>   getColorGreen()      { return colorGreen; }
    public Setting<Float>   getColorBlue()       { return colorBlue; }
    public Setting<Float>   getColorAlpha()      { return colorAlpha; }
    public Setting<Boolean> getVisibilityToggle(){ return visibilityToggle; }

    @Override
    public Extension createModule() {
        return this;
    }
}

Render3D.java

package io.github.ensamisten.client.util;

import com.mojang.blaze3d.vertex.PoseStack;
import com.mojang.blaze3d.vertex.VertexConsumer;
import com.mojang.math.Axis;
import io.github.ensamisten.client.gui.color.Color;
import net.minecraft.client.Camera;
import net.minecraft.client.Minecraft;
import net.minecraft.client.model.EntityModel;
import net.minecraft.client.renderer.MultiBufferSource;
import net.minecraft.client.renderer.ShapeRenderer;
import net.minecraft.client.renderer.entity.EntityRenderer;
import net.minecraft.client.renderer.entity.LivingEntityRenderer;
import net.minecraft.client.renderer.entity.state.LivingEntityRenderState;
import net.minecraft.client.renderer.rendertype.RenderType;
import net.minecraft.client.renderer.rendertype.RenderTypes;
import net.minecraft.core.Direction;
import net.minecraft.util.Mth;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.LivingEntity;
import net.minecraft.world.entity.Pose;
import net.minecraft.world.phys.AABB;
import net.minecraft.world.phys.Vec3;
import org.joml.Matrix4f;
import org.joml.Vector3f;

public class Render3D {
    private static final Minecraft MC = Minecraft.getInstance();

    public static void draw3DBox(PoseStack matrixStack, Camera camera, AABB box, Color color, float lineThickness) {
        Vec3 cam = camera.position();
        MultiBufferSource.BufferSource buffers = MC.renderBuffers().bufferSource();
        VertexConsumer lines = buffers.getBuffer(RenderLayers.LINES);

        int argb = ((int)(color.getAlpha() * 255.0f) << 24)
                | ((int)(color.getRed()   * 255.0f) << 16)
                | ((int)(color.getGreen() * 255.0f) << 8)
                |  (int)(color.getBlue()  * 255.0f);

        // renderShape expects local coords + xyz offsets
        ShapeRenderer.renderShape(
                matrixStack,
                lines,
                net.minecraft.world.phys.shapes.Shapes.create(box),
                -cam.x, -cam.y, -cam.z,
                argb,
                Math.max(1.0f, lineThickness)
        );

    }

}

RenderLayers.java

package io.github.ensamisten.client.util;

import java.util.function.Function;

import com.mojang.blaze3d.pipeline.BlendFunction;
import com.mojang.blaze3d.pipeline.ColorTargetState;
import com.mojang.blaze3d.pipeline.DepthStencilState;
import com.mojang.blaze3d.pipeline.RenderPipeline;
import com.mojang.blaze3d.platform.CompareOp;
import net.minecraft.client.renderer.RenderPipelines;
import net.minecraft.client.renderer.rendertype.RenderSetup;
import net.minecraft.client.renderer.rendertype.RenderType;
import net.minecraft.resources.Identifier;
import net.minecraft.util.Util;

public class RenderLayers {


    public static final RenderPipeline LINES_NO_DEPTH_PIPELINE = RenderPipelines.register(
            RenderPipeline.builder(RenderPipelines.LINES_SNIPPET)
                    .withLocation("pipeline/allyship_lines_no_depth")
                    .withDepthStencilState(new DepthStencilState(CompareOp.ALWAYS_PASS, false))
                    .build()
    );

    public static final RenderPipeline TRANSLUCENT_QUADS_PIPELINE = RenderPipelines.register(
            RenderPipeline.builder(RenderPipelines.DEBUG_FILLED_SNIPPET)
                    .withLocation("pipeline/allyship_translucent_quads")
                    .withColorTargetState(new ColorTargetState(BlendFunction.TRANSLUCENT))
                    .withCull(false)
                    .build()
    );

    public static final RenderType LINES = RenderType.create(
            "allyship_lines",
            RenderSetup.builder(LINES_NO_DEPTH_PIPELINE)
                    .bufferSize(RenderType.BIG_BUFFER_SIZE)   // 4,194,304 bytes — was 1536
                    .createRenderSetup());

    public static final RenderType QUADS = RenderType.create(
            "allyship_quads",
            RenderSetup.builder(TRANSLUCENT_QUADS_PIPELINE)
                    .bufferSize(RenderType.BIG_BUFFER_SIZE)
                    .createRenderSetup());

}