minecraft-kajetans-mod/src/main/java/me/km/snuviscript/Limits.java

package me.km.snuviscript;

import java.util.HashMap;
import net.minecraft.entity.EntityType;
import net.minecraft.world.server.ServerWorld;

public class Limits {
    private final HashMap<EntityType, Integer> limits = new HashMap<>();
    private final HashMap<EntityType, Integer> currentAmount = new HashMap<>();

    public void setLimit(EntityType type, int limit) {
        limits.put(type, limit);
    }

    public void removeLimit(EntityType type) {
        limits.remove(type);
    }

    public void clearLimits() {
        limits.clear();
    }

    public void tick(Iterable<ServerWorld> worlds) {
        currentAmount.replaceAll((type, amount) -> 0);
        for(ServerWorld w : worlds) {
            w.getEntities().forEach(ent -> add(ent.getType()));
        }
    }

    private void add(EntityType type) {
        currentAmount.compute(type, (entType, amount) -> amount == null ? 1 : amount + 1);
    }

    public boolean isAllowedToSpawn(EntityType type) {
        Integer limit = limits.get(type);
        if(limit == null) {
            return true;
        }
        Integer amount = currentAmount.get(type);
        if(amount == null) {
            return true;
        }
        return amount < limit;
    }
}