this is pure insanity

common/src/main/java/net/lerariemann/infinity/access/BlockEntityTypeAccess.java

···
+package net.lerariemann.infinity.access;
+
+import net.minecraft.block.Block;
+import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.block.entity.BlockEntityType;
+
+import java.util.Set;
+
+public interface BlockEntityTypeAccess<T extends BlockEntity> {
+    BlockEntityType.BlockEntityFactory<? extends T> infinity$getFactory();
+    Set<Block> infinity$getBlocks();
+}

common/src/main/java/net/lerariemann/infinity/dimensions/RandomFeaturesList.java

···
 import net.lerariemann.infinity.dimensions.features.local_modifications.RandomGeode;
 import net.lerariemann.infinity.dimensions.features.local_modifications.RandomIceberg;
 import net.lerariemann.infinity.dimensions.features.local_modifications.RandomRock;
-import net.lerariemann.infinity.dimensions.features.raw_generation.RandomBonusChest;
+import net.lerariemann.infinity.dimensions.features.top_layer.RandomBonusChest;
 import net.lerariemann.infinity.dimensions.features.raw_generation.RandomEndIsland;
 import net.lerariemann.infinity.dimensions.features.raw_generation.RandomShape;
 import net.lerariemann.infinity.dimensions.features.surface_structures.*;
···
         data.add(undergroundDecoration());
         data.add(fluidSprings());
         data.add(vegetation());
-        data.add(getAllElements(ConfigType.TOP_LAYER));
+        data.add(topLayer());
     }
 
     NbtList getAllElements(ConfigType name) {
···
     NbtList rawGeneration() {
         NbtList res = new NbtList();
         addRandomFeature("end_island", res, RandomEndIsland::new);
-        addRandomFeature("bonus_chest", res, RandomBonusChest::new);
         if (roll("shape")) addRandomFeature(res, new RandomShape(this, PROVIDER.randomName(random, ConfigType.SHAPE_TYPES)));
         return res;
     }
···
             addRandomFeature("water_plants", res, RandomPickle::new);
             addRandomFeature("water_plants", res, RandomKelp::new);
         }
+        return res;
+    }
+
+    NbtList topLayer() {
+        NbtList res = getAllElements(ConfigType.TOP_LAYER);
+        addRandomFeature("bonus_chest", res, RandomBonusChest::new);
         return res;
     }
 

common/src/main/java/net/lerariemann/infinity/dimensions/features/raw_generation/RandomBonusChest.java

···
-package net.lerariemann.infinity.dimensions.features.raw_generation;
-
-import net.lerariemann.infinity.dimensions.RandomFeaturesList;
-import net.lerariemann.infinity.dimensions.features.Placement;
-import net.lerariemann.infinity.dimensions.features.RandomisedFeature;
-import net.lerariemann.infinity.util.core.ConfigType;
-import net.minecraft.nbt.NbtCompound;
-import net.minecraft.nbt.NbtList;
-
-public class RandomBonusChest extends RandomisedFeature {
-    public RandomBonusChest(RandomFeaturesList lst) {
-        super(lst, "chest");
-        id = "infinity:bonus_chest";
-        savePlacement();
-    }
-
-    @Override
-    public NbtCompound feature() {
-        NbtCompound config = new NbtCompound();
-        config.putString("loot", PROVIDER.randomName(random, ConfigType.LOOT_TABLES));
-        return feature(config);
-    }
-
-    @Override
-    public NbtList placement() {
-        Placement res = new Placement();
-        res.addCount(random.nextInt(1, 16));
-        res.addBiome();
-        return res.data;
-    }
-}

common/src/main/java/net/lerariemann/infinity/dimensions/features/top_layer/RandomBonusChest.java

···
+package net.lerariemann.infinity.dimensions.features.top_layer;
+
+import net.lerariemann.infinity.dimensions.RandomFeaturesList;
+import net.lerariemann.infinity.dimensions.features.Placement;
+import net.lerariemann.infinity.dimensions.features.RandomisedFeature;
+import net.lerariemann.infinity.util.core.ConfigType;
+import net.lerariemann.infinity.util.core.NbtUtils;
+import net.minecraft.loot.LootTables;
+import net.minecraft.nbt.NbtCompound;
+import net.minecraft.nbt.NbtList;
+
+public class RandomBonusChest extends RandomisedFeature {
+    boolean isFrequent;
+
+    public RandomBonusChest(RandomFeaturesList lst) {
+        super(lst, "chest");
+        id = "infinity:bonus_chest";
+        isFrequent = random.nextBoolean();
+        savePlacement();
+    }
+
+    @Override
+    public NbtCompound feature() {
+        NbtCompound config = new NbtCompound();
+        config.putString("loot", random.nextBoolean() ? LootTables.SPAWN_BONUS_CHEST.getValue().toString() :
+                random.nextBoolean() ? "infinity:random" : PROVIDER.randomName(random, ConfigType.LOOT_TABLES));
+        config.put("block", NbtUtils.nameToElement(random.nextBoolean() ? "minecraft:chest" :
+                isFrequent ? "minecraft:barrel" : PROVIDER.randomName(random, ConfigType.CHESTS)));
+        return feature(config);
+    }
+
+    @Override
+    public NbtList placement() {
+        Placement res = new Placement();
+        if (isFrequent) {
+            res.addCountEveryLayer(random.nextInt(1, 8));
+        } else {
+            res.addCountEveryLayer(1);
+            res.addRarityFilter(random.nextInt(1, 8));
+        }
+        res.addBiome();
+        return res.data;
+    }
+}

common/src/main/java/net/lerariemann/infinity/features/RandomBonusChestFeature.java

···
 
 import com.mojang.serialization.Codec;
 import com.mojang.serialization.codecs.RecordCodecBuilder;
-import it.unimi.dsi.fastutil.ints.IntArrayList;
+import net.lerariemann.infinity.InfinityMod;
+import net.lerariemann.infinity.util.core.ConfigType;
 import net.minecraft.block.BlockState;
 import net.minecraft.block.Blocks;
 import net.minecraft.inventory.LootableInventory;
 import net.minecraft.registry.RegistryKey;
 import net.minecraft.registry.RegistryKeys;
 import net.minecraft.util.Identifier;
-import net.minecraft.util.Util;
 import net.minecraft.util.math.BlockPos;
-import net.minecraft.util.math.ChunkPos;
 import net.minecraft.util.math.Direction;
 import net.minecraft.util.math.random.Random;
-import net.minecraft.world.Heightmap;
 import net.minecraft.world.StructureWorldAccess;
 import net.minecraft.world.gen.feature.Feature;
 import net.minecraft.world.gen.feature.FeatureConfig;
 import net.minecraft.world.gen.feature.util.FeatureContext;
 
-import java.util.stream.IntStream;
+import java.util.Optional;
 
 public class RandomBonusChestFeature extends Feature<RandomBonusChestFeature.Config> {
     public RandomBonusChestFeature(Codec<Config> codec) {
···
     public boolean generate(FeatureContext<Config> context) {
         Random random = context.getRandom();
         StructureWorldAccess structureWorldAccess = context.getWorld();
-        ChunkPos chunkPos = new ChunkPos(context.getOrigin());
-        IntArrayList intArrayList = Util.shuffle(IntStream.rangeClosed(chunkPos.getStartX(), chunkPos.getEndX()), random);
-        IntArrayList intArrayList2 = Util.shuffle(IntStream.rangeClosed(chunkPos.getStartZ(), chunkPos.getEndZ()), random);
-        BlockPos.Mutable mutable = new BlockPos.Mutable();
 
-        for (Integer integer : intArrayList) {
-            for (Integer integer2 : intArrayList2) {
-                mutable.set(integer, 0, integer2);
-                BlockPos blockPos = structureWorldAccess.getTopPosition(Heightmap.Type.MOTION_BLOCKING_NO_LEAVES, mutable);
-                if (structureWorldAccess.isAir(blockPos) || structureWorldAccess.getBlockState(blockPos).getCollisionShape(structureWorldAccess, blockPos).isEmpty()) {
-                    structureWorldAccess.setBlockState(blockPos, Blocks.CHEST.getDefaultState(), 2);
-                    LootableInventory.setLootTable(structureWorldAccess, random, blockPos, RegistryKey.of(RegistryKeys.LOOT_TABLE, context.getConfig().loot));
-                    BlockState blockState = Blocks.TORCH.getDefaultState();
+        BlockPos blockPos = context.getOrigin();
+        if (structureWorldAccess.isAir(blockPos) || structureWorldAccess.getBlockState(blockPos).getCollisionShape(structureWorldAccess, blockPos).isEmpty()) {
+            BlockState state = context.getConfig().block.orElseGet(Blocks.CHEST::getDefaultState);
+            structureWorldAccess.setBlockState(blockPos, state, 2);
+            Identifier id = context.getConfig().loot;
+            if (id.toString().equals("infinity:random")) id = Identifier.of(InfinityMod.provider.randomName(context.getRandom(), ConfigType.LOOT_TABLES));
+            LootableInventory.setLootTable(structureWorldAccess, random, blockPos, RegistryKey.of(RegistryKeys.LOOT_TABLE, id));
+            BlockState blockState = Blocks.TORCH.getDefaultState();
 
-                    for (Direction direction : Direction.Type.HORIZONTAL) {
-                        BlockPos blockPos2 = blockPos.offset(direction);
-                        if (blockState.canPlaceAt(structureWorldAccess, blockPos2)) {
-                            structureWorldAccess.setBlockState(blockPos2, blockState, 2);
-                        }
-                    }
-
-                    return true;
+            for (Direction direction : Direction.Type.HORIZONTAL) {
+                BlockPos blockPos2 = blockPos.offset(direction);
+                if (blockState.canPlaceAt(structureWorldAccess, blockPos2)) {
+                    structureWorldAccess.setBlockState(blockPos2, blockState, 2);
                 }
             }
+
+            return true;
         }
         return false;
     }
 
-    public record Config(Identifier loot) implements FeatureConfig {
+    public record Config(Identifier loot, Optional<BlockState> block) implements FeatureConfig {
         public static final Codec<Config> CODEC = RecordCodecBuilder.create(instance -> instance.group(
-                (Identifier.CODEC.fieldOf("loot")).forGetter(a -> a.loot)).apply(
+                (Identifier.CODEC.fieldOf("loot")).forGetter(a -> a.loot),
+                (BlockState.CODEC.optionalFieldOf("block")).forGetter(a -> a.block)).apply(
                 instance, Config::new));
     }
 }

common/src/main/java/net/lerariemann/infinity/mixin/qol/BlockEntityTypeMixin.java

···
+package net.lerariemann.infinity.mixin.qol;
+
+import net.lerariemann.infinity.access.BlockEntityTypeAccess;
+import net.minecraft.block.Block;
+import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.block.entity.BlockEntityType;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+
+import java.util.Set;
+
+@Mixin(BlockEntityType.class)
+public class BlockEntityTypeMixin<T extends BlockEntity> implements BlockEntityTypeAccess<T> {
+    @Shadow @Final private BlockEntityType.BlockEntityFactory<? extends T> factory;
+
+    @Shadow @Final private Set<Block> blocks;
+
+    @Override
+    public BlockEntityType.BlockEntityFactory<? extends T> infinity$getFactory() {
+        return factory;
+    }
+
+    @Override
+    public Set<Block> infinity$getBlocks() {
+        return blocks;
+    }
+}

common/src/main/java/net/lerariemann/infinity/util/config/ConfigGenerator.java

···
 package net.lerariemann.infinity.util.config;
 
 import net.lerariemann.infinity.InfinityMod;
+import net.lerariemann.infinity.access.BlockEntityTypeAccess;
 import net.lerariemann.infinity.block.entity.CosmicAltarBlockEntity;
 import net.lerariemann.infinity.util.core.ConfigType;
 import net.lerariemann.infinity.util.var.ColorLogic;
 import net.minecraft.block.Block;
 import net.minecraft.block.BlockState;
 import net.minecraft.block.FallingBlock;
+import net.minecraft.block.entity.BlockEntity;
+import net.minecraft.block.entity.BlockEntityType;
+import net.minecraft.block.entity.LootableContainerBlockEntity;
+import net.minecraft.block.entity.ShulkerBoxBlockEntity;
 import net.minecraft.block.enums.BlockFace;
 import net.minecraft.entity.EntityType;
 import net.minecraft.entity.SpawnGroup;
···
         ConfigFactory.of(Registries.STATUS_EFFECT, ConfigGenerator::extractEffect).generate(ConfigType.EFFECTS);
         generateSounds();
         generateBlockTags();
+        generateBEs();
         SurfaceRuleScanner.scan(server);
         DynamicRegistryManager manager = server.getRegistryManager();
         ConfigFactory.of(manager.get(RegistryKeys.JUKEBOX_SONG)).generate(ConfigType.JUKEBOXES);
···
 
     static void generateSounds() {
         Registry<SoundEvent> r = Registries.SOUND_EVENT;
-        DataCollection.Logged music = new DataCollection.Logged(ConfigType.MUSIC, "music tracks");
-        DataCollection.Logged sounds = new DataCollection.Logged(ConfigType.SOUNDS);
+        DataCollection music = new DataCollection.Logged(ConfigType.MUSIC, "music tracks");
+        DataCollection sounds = new DataCollection.Logged(ConfigType.SOUNDS);
         r.getKeys().forEach(a -> {
             Identifier id = a.getValue();
             if (id.toString().contains("music")) music.addIdentifier(id);
···
     }
 
     static void generateBlockTags() {
-        DataCollection.Logged tagMap = new DataCollection.Logged(ConfigType.TAGS, "block tags");
+        DataCollection tagMap = new DataCollection.Logged(ConfigType.TAGS, "block tags");
         Registries.BLOCK.streamTags().forEach(tagKey -> tagMap.addIdentifier(tagKey.id()));
         tagMap.save();
     }
 
     static Set<String> generateFluids() {
-        DataCollection.Logged fluidMap = new DataCollection.Logged(ConfigType.FLUIDS);
+        DataCollection fluidMap = new DataCollection.Logged(ConfigType.FLUIDS);
         Registry<Fluid> r = Registries.FLUID;
         Set<String> fluidBlockNames = new HashSet<>();
         r.getKeys().forEach(key -> {
···
     }
 
     static void generateBlocks(ServerWorld serverWorld, BlockPos inAir, BlockPos onAltar, Set<String> excludedBlockNames) {
-        DataCollection.Logged blockMap = new DataCollection.Logged(ConfigType.BLOCKS);
+        DataCollection blockMap = new DataCollection.Logged(ConfigType.BLOCKS);
         DataCollection colorPresetMap = new DataCollection(ConfigType.COLOR_PRESETS);
         DataCollection flowerMap = new DataCollection(ConfigType.FLOWERS);
         Registry<Block> r = Registries.BLOCK;
···
         blockMap.save();
         colorPresetMap.save();
         flowerMap.save();
+    }
+
+    static void generateBEs() {
+        Registry<BlockEntityType<?>> r = Registries.BLOCK_ENTITY_TYPE;
+        DataCollection chests = new DataCollection.Logged(ConfigType.CHESTS);
+        Set<Identifier> allBlocks = new HashSet<>();
+        r.getKeys().forEach(key -> {
+            BlockEntityTypeAccess<? extends BlockEntity> bet = (BlockEntityTypeAccess<?>)(r.get(key));
+            assert bet != null;
+            Set<Block> blocks = bet.infinity$getBlocks();
+
+            BlockEntity be = bet.infinity$getFactory().create(BlockPos.ORIGIN, ((Block)blocks.toArray()[0]).getDefaultState());
+            if (be instanceof LootableContainerBlockEntity && !(be instanceof ShulkerBoxBlockEntity)) {
+                for (Block b : blocks) allBlocks.add(Registries.BLOCK.getId(b));
+            }
+        });
+        allBlocks.forEach(chests::addIdentifier);
+        chests.save();
     }
 
     static boolean checkColorSet(String block) {

common/src/main/java/net/lerariemann/infinity/util/core/ConfigType.java

···
     FULL_BLOCKS("full_blocks", "minecraft:stone"),
     FULL_BLOCKS_WG("full_blocks_worldgen", "minecraft:stone"),
     FLOWERS("flowers", "minecraft:poppy"),
+    CHESTS("chests", "minecraft:chest"),
     FLUIDS("fluids", "minecraft:water", NbtUtils::nameToFluid),
     ITEMS("items", "minecraft:stick"),
     SOUNDS("sounds", "minecraft:block.stone.step"),

common/src/main/resources/config/hardcoded/features/toplayermodification.json

···
 {
+  "infinity_version": 2004006,
   "elements": [
     {
       "key": "minecraft:freeze_top_layer",
···
     {
       "key": "minecraft:void_start_platform",
       "weight": 0.5
-    },
-    {
-      "key": "infinity:bonus_chest",
-      "weight": 0.2
     }
   ]
 }

common/src/main/resources/infinity.mixins.json

···
     "options.EntityMixin",
     "options.PlayerManagerMixin",
     "options.ServerWorldMixin",
+    "qol.BlockEntityTypeMixin",
     "qol.EnderDragonEntityMixin",
     "qol.FireBlockMixin",
     "qol.LavaFluidMixin"