Pill in dungeons

Fixes #6

src/data/kotlin/ru/dbotthepony/mc/otm/datagen/DataGen.kt

@@ -6,10 +6,12 @@ import net.minecraft.data.recipes.ShapelessRecipeBuilder
 import net.minecraft.resources.ResourceLocation
 import net.minecraft.tags.Tag
 import net.minecraft.world.item.Item
+import net.minecraft.world.item.ItemStack
 import net.minecraft.world.level.block.Block
 import net.minecraft.world.level.block.IronBarsBlock
 import net.minecraftforge.client.model.generators.ModelFile
 import net.minecraftforge.common.Tags
+import net.minecraftforge.common.loot.LootTableIdCondition
 import net.minecraftforge.eventbus.api.SubscribeEvent
 import net.minecraftforge.fml.common.Mod
 import net.minecraftforge.forge.event.lifecycle.GatherDataEvent
@@ -22,11 +24,13 @@ import ru.dbotthepony.mc.otm.block.EnergyCounterBlock
 import ru.dbotthepony.mc.otm.block.BlockMatteryRotatable
 import ru.dbotthepony.mc.otm.block.PatternStorageBlock
 import ru.dbotthepony.mc.otm.block.entity.worker.WorkerState
+import ru.dbotthepony.mc.otm.data.LootTableBasicAppender
 import ru.dbotthepony.mc.otm.datagen.blocks.BatteryBankProvider
 import ru.dbotthepony.mc.otm.datagen.blocks.MatterBankProvider
 import ru.dbotthepony.mc.otm.datagen.blocks.MatteryBlockStateProvider
+import ru.dbotthepony.mc.otm.datagen.loot.LootModifiers
 import ru.dbotthepony.mc.otm.datagen.items.MatteryItemModelProvider
-import ru.dbotthepony.mc.otm.datagen.loot.MatteryLootTableProvider
+import ru.dbotthepony.mc.otm.datagen.loot.LootTables
 import ru.dbotthepony.mc.otm.datagen.loot.TileNbtCopy
 import ru.dbotthepony.mc.otm.datagen.models.BlockMatteryModelProvider
 import ru.dbotthepony.mc.otm.datagen.recipes.MatteryRecipeProvider
@@ -40,8 +44,9 @@ object DataGen {
 	private lateinit var blockModelProvider: BlockMatteryModelProvider
 	private lateinit var itemModelProvider: MatteryItemModelProvider
 	private lateinit var blockStateProvider: MatteryBlockStateProvider
-	private lateinit var lootTableProvider: MatteryLootTableProvider
+	private lateinit var lootTableProvider: LootTables
 	private lateinit var recipeProvider: MatteryRecipeProvider
+	private lateinit var lootModifier: LootModifiers
 
 	private fun decorativeCubeAll(vararg blocks: Block) {
 		blockModelProvider.decorativeCubeAll(*blocks)
@@ -107,8 +112,9 @@ object DataGen {
 		blockModelProvider = BlockMatteryModelProvider(event)
 		blockStateProvider = MatteryBlockStateProvider(event)
 		itemModelProvider = MatteryItemModelProvider(event)
-		lootTableProvider = MatteryLootTableProvider(event.generator)
+		lootTableProvider = LootTables(event.generator)
 		recipeProvider = MatteryRecipeProvider(event.generator)
+		lootModifier = LootModifiers(event.generator)
 
 		event.generator.addProvider(blockModelProvider)
 		event.generator.addProvider(itemModelProvider)
@@ -118,6 +124,7 @@ object DataGen {
 		event.generator.addProvider(MatterBankProvider(event))
 		event.generator.addProvider(BatteryBankProvider(event))
 		event.generator.addProvider(lootTableProvider)
+		event.generator.addProvider(lootModifier)
 
 		decorativeCubeAll(*MBlocks.CRATES)
 		decorativeCubeAll(MBlocks.CARBON_FIBRE_BLOCK)
@@ -431,5 +438,42 @@ object DataGen {
 				}
 			}
 		}
+
+		lootModifier.lambda {
+			it.add("dungeon_pill", LootTableBasicAppender.Companion, LootTableBasicAppender(
+				arrayOf(LootTableIdCondition.Builder(ResourceLocation("chests/simple_dungeon")).build()),
+				ItemStack(MItems.PILL_ANDROID, 1) to 0.1,
+				ItemStack(MItems.PILL_HEAL, 2) to 0.5,
+			))
+
+			it.add("mineshaft_pill", LootTableBasicAppender.Companion, LootTableBasicAppender(
+				arrayOf(LootTableIdCondition.Builder(ResourceLocation("chests/abandoned_mineshaft")).build()),
+				ItemStack(MItems.PILL_ANDROID, 1) to 0.1,
+				ItemStack(MItems.PILL_HEAL, 1) to 0.25,
+			))
+
+			it.add("desert_pyramid_pill", LootTableBasicAppender.Companion, LootTableBasicAppender(
+				arrayOf(LootTableIdCondition.Builder(ResourceLocation("chests/abandoned_mineshaft")).build()),
+				ItemStack(MItems.PILL_ANDROID, 1) to 0.05,
+				ItemStack(MItems.PILL_HEAL, 1) to 0.3,
+			))
+
+			it.add("jungle_temple_pill", LootTableBasicAppender.Companion, LootTableBasicAppender(
+				arrayOf(LootTableIdCondition.Builder(ResourceLocation("chests/jungle_temple")).build()),
+				ItemStack(MItems.PILL_ANDROID, 1) to 0.1
+			))
+
+			it.add("end_city_pill", LootTableBasicAppender.Companion, LootTableBasicAppender(
+				arrayOf(LootTableIdCondition.Builder(ResourceLocation("chests/end_city_treasure")).build()),
+				ItemStack(MItems.PILL_ANDROID, 1) to 0.1,
+				ItemStack(MItems.PILL_HUMANE, 1) to 0.2,
+				ItemStack(MItems.PILL_OBLIVION, 1) to 0.4,
+			))
+
+			it.add("shipwreck_supply_pill", LootTableBasicAppender.Companion, LootTableBasicAppender(
+				arrayOf(LootTableIdCondition.Builder(ResourceLocation("chests/shipwreck_supply")).build()),
+				ItemStack(MItems.PILL_HUMANE, 1) to 0.3
+			))
+		}
 	}
 }

src/data/kotlin/ru/dbotthepony/mc/otm/datagen/loot/LootModifiers.kt

@@ -0,0 +1,19 @@
+package ru.dbotthepony.mc.otm.datagen.loot
+
+import net.minecraft.data.DataGenerator
+import net.minecraftforge.common.data.GlobalLootModifierProvider
+import ru.dbotthepony.mc.otm.datagen.DataGen
+
+class LootModifiers(generator: DataGenerator) : GlobalLootModifierProvider(generator, DataGen.MOD_ID) {
+	private val lambdas = ArrayList<(LootModifiers) -> Unit>()
+
+	fun lambda(lambda: (LootModifiers) -> Unit) {
+		lambdas.add(lambda)
+	}
+
+	override fun start() {
+		for (lambda in lambdas) {
+			lambda.invoke(this)
+		}
+	}
+}

src/data/kotlin/ru/dbotthepony/mc/otm/datagen/loot/LootTables.kt

@@ -10,7 +10,6 @@ import net.minecraft.world.level.storage.loot.LootTable
 import net.minecraft.world.level.storage.loot.ValidationContext
 import net.minecraft.world.level.storage.loot.entries.LootItem
 import net.minecraft.world.level.storage.loot.functions.CopyNbtFunction
-import net.minecraft.world.level.storage.loot.functions.LootItemFunctions
 import net.minecraft.world.level.storage.loot.parameters.LootContextParamSet
 import net.minecraft.world.level.storage.loot.parameters.LootContextParamSets
 import net.minecraft.world.level.storage.loot.providers.nbt.ContextNbtProvider
@@ -30,7 +29,7 @@ fun singleLootPool(f: (LootPool.Builder) -> Unit): LootTable.Builder {
 data class NbtCopy(val source: String, val destination: String, val strategy: CopyNbtFunction.MergeStrategy = CopyNbtFunction.MergeStrategy.REPLACE)
 data class TileNbtCopy(val source: String, val strategy: CopyNbtFunction.MergeStrategy = CopyNbtFunction.MergeStrategy.REPLACE)
 
-class MatteryLootTableProvider(generator: DataGenerator) : LootTableProvider(generator) {
+class LootTables(generator: DataGenerator) : LootTableProvider(generator) {
 	private val providers = ArrayList<LootTuple>()
 
 	override fun getTables(): List<LootTuple> {

src/main/java/ru/dbotthepony/mc/otm/OverdriveThatMatters.java

@@ -52,6 +52,10 @@ public final class OverdriveThatMatters {
 	};
 
 	public OverdriveThatMatters() {
+		if (INSTANCE != null) {
+			throw new IllegalStateException("yo what the fuck");
+		}
+
 		INSTANCE = this;
 
 		// Register the setup method for modloading
@@ -77,6 +81,7 @@ public final class OverdriveThatMatters {
 		FMLJavaModLoadingContext.get().getModEventBus().register(AndroidResearch.class);
 		FMLJavaModLoadingContext.get().getModEventBus().register(MStats.class);
 		FMLJavaModLoadingContext.get().getModEventBus().register(MRecipes.class);
+		FMLJavaModLoadingContext.get().getModEventBus().register(LootModifiers.INSTANCE);
 
 		MinecraftForge.EVENT_BUS.register(DrivePool.INSTANCE);
 		MinecraftForge.EVENT_BUS.register(PortableCondensationDriveItem.Companion);

src/main/kotlin/ru/dbotthepony/mc/otm/data/LootTableAppender.kt

@@ -0,0 +1,117 @@
+package ru.dbotthepony.mc.otm.data
+
+import com.google.common.collect.ImmutableList
+import com.google.gson.JsonArray
+import com.google.gson.JsonObject
+import com.google.gson.JsonParseException
+import net.minecraft.resources.ResourceLocation
+import net.minecraft.world.item.ItemStack
+import net.minecraft.world.level.storage.loot.Deserializers
+import net.minecraft.world.level.storage.loot.LootContext
+import net.minecraft.world.level.storage.loot.LootPool
+import net.minecraft.world.level.storage.loot.predicates.LootItemCondition
+import net.minecraftforge.common.loot.GlobalLootModifierSerializer
+import net.minecraftforge.common.loot.LootModifier
+import net.minecraftforge.registries.ForgeRegistries
+import ru.dbotthepony.mc.otm.OverdriveThatMatters
+
+class LootTableAppender(conditions: Array<out LootItemCondition>, private vararg val pools: LootPool) : LootModifier(conditions) {
+	override fun doApply(generatedLoot: MutableList<ItemStack>, context: LootContext): MutableList<ItemStack> {
+		pools.forEach { it.addRandomItems(generatedLoot::add, context) }
+		return generatedLoot
+	}
+
+	companion object : GlobalLootModifierSerializer<LootTableAppender>() {
+		private val lootPoolGson = Deserializers.createLootTableSerializer().create()
+
+		override fun read(
+			location: ResourceLocation,
+			`object`: JsonObject,
+			ailootcondition: Array<out LootItemCondition>
+		): LootTableAppender {
+			val pools = lootPoolGson.fromJson(`object`["pools"], Array<LootPool>::class.java) ?: throw JsonParseException("Invalid 'pools' entry")
+			pools.forEach {
+				@Suppress("UNNECESSARY_NOT_NULL_ASSERTION")
+				it!!
+			}
+			return LootTableAppender(ailootcondition, *pools)
+		}
+
+		override fun write(instance: LootTableAppender): JsonObject {
+			val build = makeConditions(instance.conditions)
+			build.add("pools", lootPoolGson.toJsonTree(instance.pools))
+			return build
+		}
+
+		init {
+			registryName = ResourceLocation(OverdriveThatMatters.MOD_ID, "loot_appender")
+		}
+	}
+}
+
+class LootTableBasicAppender(conditions: Array<out LootItemCondition>, items: List<Pair<ItemStack, Double>>) : LootModifier(conditions) {
+	private val items = ImmutableList.copyOf(items)
+
+	constructor(conditions: Array<out LootItemCondition>, vararg items: Pair<ItemStack, Double>) : this(conditions, items.toList())
+
+	override fun doApply(generatedLoot: MutableList<ItemStack>, context: LootContext): MutableList<ItemStack> {
+		for ((item, chance) in items) {
+			if (context.random.nextDouble() <= chance) {
+				generatedLoot.add(item.copy())
+			}
+		}
+
+		return generatedLoot
+	}
+
+	companion object : GlobalLootModifierSerializer<LootTableBasicAppender>() {
+		override fun read(
+			location: ResourceLocation,
+			obj: JsonObject,
+			ailootcondition: Array<out LootItemCondition>
+		): LootTableBasicAppender {
+			val entries = ArrayList<Pair<ItemStack, Double>>()
+
+			for (entry in obj["entries"].asJsonArray) {
+				entry as JsonObject
+				val item = entry["item"]?.asJsonObject ?: throw JsonParseException("Missing 'item'")
+				val name = item["name"]?.asString ?: throw JsonParseException("Missing 'item.name'")
+				val count = item["count"]?.asInt ?: throw JsonParseException("Missing 'item.count'")
+				val chance = entry["chance"]?.asDouble ?: throw JsonParseException("Missing 'chance'")
+
+				val itemObj = ForgeRegistries.ITEMS.getValue(ResourceLocation(name)) ?: throw JsonParseException("Unknown item $name")
+				val stack = ItemStack(itemObj, count)
+
+				entries.add(stack to chance)
+			}
+
+			return LootTableBasicAppender(ailootcondition, ImmutableList.copyOf(entries))
+		}
+
+		override fun write(instance: LootTableBasicAppender): JsonObject {
+			val obj = makeConditions(instance.conditions)
+			val listing = JsonArray()
+
+			for ((item, chance) in instance.items) {
+				val entry = JsonObject()
+
+				entry.add("item", JsonObject().also {
+					it.addProperty("name", item.item.registryName!!.toString())
+					it.addProperty("count", item.count)
+				})
+
+				entry.addProperty("chance", chance)
+
+				listing.add(entry)
+			}
+
+			obj.add("entries", listing)
+
+			return obj
+		}
+
+		init {
+			registryName = ResourceLocation(OverdriveThatMatters.MOD_ID, "loot_appender_basic")
+		}
+	}
+}

src/main/kotlin/ru/dbotthepony/mc/otm/registry/LootModifiers.kt

@@ -0,0 +1,15 @@
+package ru.dbotthepony.mc.otm.registry
+
+import net.minecraftforge.common.loot.GlobalLootModifierSerializer
+import net.minecraftforge.event.RegistryEvent
+import net.minecraftforge.eventbus.api.SubscribeEvent
+import ru.dbotthepony.mc.otm.data.LootTableAppender
+import ru.dbotthepony.mc.otm.data.LootTableBasicAppender
+
+object LootModifiers {
+	@SubscribeEvent
+	fun register(event: RegistryEvent.Register<GlobalLootModifierSerializer<*>>) {
+		event.registry.register(LootTableAppender.Companion)
+		event.registry.register(LootTableBasicAppender.Companion)
+	}
+}