BlockLootTableHolder#fill

src/main/kotlin/ru/dbotthepony/mc/otm/block/entity/decorative/CargoCrateBlockEntity.kt

@@ -1,18 +1,11 @@
 package ru.dbotthepony.mc.otm.block.entity.decorative
 
-import net.minecraft.advancements.CriteriaTriggers
 import net.minecraft.core.BlockPos
 import net.minecraft.core.HolderLookup
-import net.minecraft.core.registries.Registries
 import net.minecraft.nbt.CompoundTag
-import net.minecraft.nbt.Tag
 import net.minecraft.network.chat.Component
-import net.minecraft.resources.ResourceKey
-import net.minecraft.resources.ResourceLocation
 import net.minecraft.server.level.ServerLevel
-import net.minecraft.server.level.ServerPlayer
 import net.minecraft.sounds.SoundSource
-import net.minecraft.world.Containers
 import net.minecraft.world.entity.player.Inventory
 import net.minecraft.world.entity.player.Player
 import net.minecraft.world.inventory.AbstractContainerMenu
@@ -21,18 +14,14 @@ import net.minecraft.world.level.Level
 import net.minecraft.world.level.block.Block
 import net.minecraft.world.level.block.state.BlockState
 import net.minecraft.world.level.gameevent.GameEvent
-import net.minecraft.world.level.storage.loot.LootParams
-import net.minecraft.world.level.storage.loot.LootTable
-import net.minecraft.world.level.storage.loot.parameters.LootContextParamSets
-import net.minecraft.world.level.storage.loot.parameters.LootContextParams
-import net.minecraft.world.phys.Vec3
 import net.neoforged.neoforge.capabilities.Capabilities
 import ru.dbotthepony.mc.otm.block.decorative.CargoCrateBlock
 import ru.dbotthepony.mc.otm.block.entity.MatteryDeviceBlockEntity
-import ru.dbotthepony.mc.otm.container.MatteryContainer
 import ru.dbotthepony.mc.otm.container.HandlerFilter
+import ru.dbotthepony.mc.otm.container.MatteryContainer
 import ru.dbotthepony.mc.otm.core.TranslatableComponent
 import ru.dbotthepony.mc.otm.core.otmRandom
+import ru.dbotthepony.mc.otm.core.util.BlockLootTableHolder
 import ru.dbotthepony.mc.otm.menu.decorative.CargoCrateMenu
 import ru.dbotthepony.mc.otm.registry.game.MBlockEntities
 import ru.dbotthepony.mc.otm.registry.game.MSoundEvents
@@ -47,11 +36,11 @@ class CargoCrateBlockEntity(
 
 	val handler = container.handler(object : HandlerFilter {
 		override fun canInsert(slot: Int, stack: ItemStack): Boolean {
-			return lootTable == null
+			return loot.lootTable == null
 		}
 
 		override fun canExtract(slot: Int, amount: Int, stack: ItemStack): Boolean {
-			return lootTable == null
+			return loot.lootTable == null
 		}
 	})
 
@@ -70,8 +59,7 @@ class CargoCrateBlockEntity(
 		unpackLootTable(player)
 	}
 
-	var lootTable: ResourceKey<LootTable>? = null
+	val loot = BlockLootTableHolder(::markDirtyFast)
-	var lootTableSeed: Long? = null
 
 	fun onPlayerOpen() {
 		val level = level
@@ -99,52 +87,24 @@ class CargoCrateBlockEntity(
 
 	override fun saveLevel(nbt: CompoundTag, registry: HolderLookup.Provider) {
 		super.saveLevel(nbt, registry)
-
+		loot.save(nbt, registry)
-		if (lootTable != null) {
-			nbt.putString(LOOT_TABLE_KEY, lootTable!!.location().toString())
-			nbt.putLong(LOOT_TABLE_SEED_KEY, lootTableSeed ?: 0L)
-		}
 	}
 
 	override fun loadAdditional(nbt: CompoundTag, registry: HolderLookup.Provider) {
 		super.loadAdditional(nbt, registry)
-
+		loot.load(nbt, registry)
-		if (nbt.contains(LOOT_TABLE_KEY, Tag.TAG_STRING.toInt())) {
-			lootTable = ResourceLocation.tryParse(nbt.getString(LOOT_TABLE_KEY))?.let { ResourceKey.create(Registries.LOOT_TABLE, it) }
-			lootTableSeed = if (nbt.contains(LOOT_TABLE_SEED_KEY, Tag.TAG_LONG.toInt())) nbt.getLong(LOOT_TABLE_SEED_KEY) else 0L
-		}
 	}
 
-	fun unpackLootTable(ply: Player? = null) {
+	private fun unpackLootTable(ply: Player? = null) {
-		val lootTable = lootTable ?: return
+		loot.fill(level as? ServerLevel ?: return, blockPos, container, ply = ply, blockEntity = this)
-		val lootTableSeed = lootTableSeed ?: 0L
+		loot.clear()
-		val server = level?.server ?: return
-
-		val loot = server.reloadableRegistries().getLootTable(lootTable)
-
-		if (ply is ServerPlayer) {
-			CriteriaTriggers.GENERATE_LOOT.trigger(ply, lootTable)
-		}
-
-		this.lootTable = null
-		this.lootTableSeed = null
-
-		val params = LootParams.Builder(level as ServerLevel)
-			.withParameter(LootContextParams.ORIGIN, Vec3.atCenterOf(this.worldPosition))
-
-		if (ply != null) {
-			params.withLuck(ply.luck).withParameter(LootContextParams.THIS_ENTITY, ply)
-		}
-
-		Containers.dropContents(level as ServerLevel, blockPos, container)
-		loot.fill(container, params.create(LootContextParamSets.CHEST), lootTableSeed)
 	}
 
 	override val defaultDisplayName: Component
 		get() = MACHINE_NAME
 
 	override fun createMenu(containerID: Int, inventory: Inventory, ply: Player): AbstractContainerMenu? {
-		if (lootTable != null && ply.isSpectator)
+		if (loot.lootTable != null && ply.isSpectator)
 			return null
 
 		unpackLootTable(ply)

src/main/kotlin/ru/dbotthepony/mc/otm/core/Ext.kt

@@ -10,6 +10,7 @@ import com.google.common.collect.ImmutableSet
 import com.google.gson.JsonElement
 import com.google.gson.JsonObject
 import com.google.gson.JsonPrimitive
+import it.unimi.dsi.fastutil.ints.IntList
 import it.unimi.dsi.fastutil.objects.ObjectComparators
 import net.minecraft.Util
 import net.minecraft.core.BlockPos
@@ -199,6 +200,15 @@ fun IntArray.shuffle(random: RandomSource): IntArray {
 	return this
 }
 
+fun <L : IntList> L.shuffle(random: RandomSource): L {
+	for (i in lastIndex downTo 1) {
+		val j = random.nextInt(i + 1)
+		set(j, set(i, getInt(j)))
+	}
+
+	return this
+}
+
 fun <T, L : MutableList<T>> L.shuffle(random: RandomSource): L {
 	Util.shuffle(this, random)
 	return this

src/main/kotlin/ru/dbotthepony/mc/otm/core/util/BlockLootTableHolder.kt

@@ -1,5 +1,7 @@
 package ru.dbotthepony.mc.otm.core.util
 
+import net.minecraft.advancements.CriteriaTriggers
+import net.minecraft.core.BlockPos
 import net.minecraft.core.HolderLookup
 import net.minecraft.core.registries.Registries
 import net.minecraft.nbt.CompoundTag
@@ -7,15 +9,22 @@ import net.minecraft.nbt.Tag
 import net.minecraft.resources.ResourceKey
 import net.minecraft.resources.ResourceLocation
 import net.minecraft.server.level.ServerLevel
+import net.minecraft.server.level.ServerPlayer
 import net.minecraft.util.RandomSource
+import net.minecraft.world.Container
+import net.minecraft.world.Containers
+import net.minecraft.world.entity.player.Player
 import net.minecraft.world.item.ItemStack
+import net.minecraft.world.level.block.entity.BlockEntity
 import net.minecraft.world.level.block.state.BlockState
 import net.minecraft.world.level.storage.loot.LootParams
 import net.minecraft.world.level.storage.loot.LootTable
 import net.minecraft.world.level.storage.loot.parameters.LootContextParamSets
 import net.minecraft.world.level.storage.loot.parameters.LootContextParams
+import net.minecraft.world.phys.Vec3
 import ru.dbotthepony.mc.otm.block.entity.MatteryBlockEntity.Companion.LOOT_TABLE_KEY
 import ru.dbotthepony.mc.otm.block.entity.MatteryBlockEntity.Companion.LOOT_TABLE_SEED_KEY
+import ru.dbotthepony.mc.otm.core.getBlockEntityNow
 import ru.dbotthepony.mc.otm.core.otmRandom
 
 class BlockLootTableHolder(private val listener: Runnable = Runnable {  }) {
@@ -75,13 +84,49 @@ class BlockLootTableHolder(private val listener: Runnable = Runnable {  }) {
 		return if (lootTableSeed == null) fallback else GJRAND64RandomSource(lootTableSeed)
 	}
 
-	fun getItems(params: LootParams, level: ServerLevel, overrideSeed: Long? = null): MutableList<ItemStack> {
+	fun getItems(params: LootParams, level: ServerLevel, overrideSeed: Long? = null, rounds: Int = 1): MutableList<ItemStack> {
-		return lookup(level).getRandomItems(params, selectRandom(overrideSeed, level.otmRandom))
+		return lookup(level).getRandomItems(params, selectRandom(overrideSeed, level.otmRandom), rounds)
 	}
 
-	fun getItems(params: LootParams.Builder, blockState: BlockState, overrideSeed: Long? = null): MutableList<ItemStack> {
+	fun getItems(params: LootParams.Builder, blockState: BlockState, overrideSeed: Long? = null, rounds: Int = 1): MutableList<ItemStack> {
 		val level = params.level
 		val create = params.withParameter(LootContextParams.BLOCK_STATE, blockState).create(LootContextParamSets.BLOCK)
-		return getItems(create, level, overrideSeed)
+		return getItems(create, level, overrideSeed, rounds)
+	}
+
+	fun fill(level: ServerLevel, blockPos: BlockPos, container: Container, ply: Player? = null, blockEntity: BlockEntity? = level.getBlockEntityNow(blockPos), overrideSeed: Long? = null, dropContents: Boolean = true, rounds: Int = 1) {
+		if (rounds == 0)
+			return
+		else if (rounds < 0)
+			throw IllegalArgumentException("Negative amount of rounds: $rounds")
+
+		val lootTable = lootTable ?: return
+		val server = level.server
+		val loot = server.reloadableRegistries().getLootTable(lootTable)
+
+		if (ply is ServerPlayer)
+			CriteriaTriggers.GENERATE_LOOT.trigger(ply, lootTable)
+
+		val params = LootParams.Builder(level)
+			.withParameter(LootContextParams.ORIGIN, Vec3.atCenterOf(blockPos))
+
+		if (ply != null) {
+			params
+				.withLuck(ply.luck)
+				.withParameter(LootContextParams.THIS_ENTITY, ply)
+		}
+
+		if (blockEntity != null)
+			params.withParameter(LootContextParams.BLOCK_ENTITY, blockEntity)
+
+		if (dropContents)
+			Containers.dropContents(level, blockPos, container)
+
+		loot.fill(params.create(LootContextParamSets.CHEST), selectRandom(overrideSeed, level.otmRandom), container, rounds)
+	}
+
+	fun clear() {
+		lootTable = null
+		lootTableSeed = null
 	}
 }

src/main/kotlin/ru/dbotthepony/mc/otm/core/util/LootTableUtils.kt

@@ -0,0 +1,132 @@
+package ru.dbotthepony.mc.otm.core.util
+
+import it.unimi.dsi.fastutil.ints.IntArrayList
+import net.minecraft.util.Mth
+import net.minecraft.util.RandomSource
+import net.minecraft.world.Container
+import net.minecraft.world.item.Item
+import net.minecraft.world.item.ItemStack
+import net.minecraft.world.level.storage.loot.LootParams
+import net.minecraft.world.level.storage.loot.LootTable
+import org.apache.logging.log4j.LogManager
+import ru.dbotthepony.mc.otm.container.get
+import ru.dbotthepony.mc.otm.container.set
+import ru.dbotthepony.mc.otm.core.isNotEmpty
+import ru.dbotthepony.mc.otm.core.shuffle
+import java.util.stream.Collectors
+import java.util.stream.IntStream
+import kotlin.math.min
+
+private val LOGGER = LogManager.getLogger()
+
+private class Bucket {
+	val items = ArrayList<ItemStack>()
+	var last = 0
+
+	fun add(item: ItemStack) {
+		if (item.isEmpty) return
+
+		for (i in last until items.size) {
+			val existing = items[i]
+			var available = existing.maxStackSize - existing.count
+
+			if (available > 0 && ItemStack.isSameItemSameComponents(existing, item)) {
+				val diff = min(available, item.count)
+				existing.grow(diff)
+				item.shrink(diff)
+				available -= diff
+
+				if (available == 0 && i == last)
+					last++
+
+				if (item.isEmpty)
+					return
+			}
+		}
+
+		if (item.isNotEmpty) {
+			items.add(item)
+		}
+	}
+}
+
+private fun recombine(lists: Collection<Collection<ItemStack>>): MutableList<ItemStack> {
+	val result = HashMap<Item, Bucket>()
+
+	for (list in lists)
+		for (item in list)
+			result.computeIfAbsent(item.item) { Bucket() }.add(item)
+
+	return result.values
+		.stream()
+		.flatMap { it.items.stream() }
+		.collect(Collectors.toCollection(::ArrayList))
+}
+
+fun LootTable.getRandomItems(params: LootParams, random: RandomSource, rounds: Int = 1): MutableList<ItemStack> {
+	if (rounds == 0) {
+		return ArrayList()
+	} else if (rounds == 1) {
+		return getRandomItems(params, random)
+	} else {
+		return recombine(IntStream.range(0, rounds).mapToObj { getRandomItems(params, random) }.toList())
+	}
+}
+
+private fun shuffle(items: MutableList<ItemStack>, emptySlotCount: Int, random: RandomSource) {
+	val pool = ArrayList<ItemStack>(items)
+	items.clear()
+
+	while (items.size + pool.size < emptySlotCount && pool.isNotEmpty()) {
+		val select = pool.removeAt(random.nextInt(pool.size))
+		val maxStackSize = select.maxStackSize
+
+		if (maxStackSize == 1 || select.count == 1) {
+			items.add(select)
+		} else {
+			val split = select.split(Mth.nextInt(random, 1, select.count / 2))
+
+			if (split.count > 1 && random.nextBoolean())
+				pool.add(split)
+			else
+				items.add(split)
+
+			if (select.count > 1 && random.nextBoolean())
+				pool.add(select)
+			else
+				items.add(select)
+		}
+	}
+
+	items.addAll(pool)
+	items.shuffle(random)
+}
+
+fun LootTable.fill(params: LootParams, random: RandomSource, container: Container, rounds: Int = 1) {
+	val emptySlots = IntArrayList()
+
+	for (i in 0 until container.containerSize)
+		if (container[i].isEmpty)
+			emptySlots.add(i)
+
+	emptySlots.shuffle(random)
+
+	if (emptySlots.isEmpty) {
+		LOGGER.warn("Tried to fill container with no empty slots")
+		return
+	}
+
+	val items = getRandomItems(params, random, rounds)
+	shuffle(items, emptySlots.size, random)
+
+	val slotItr = emptySlots.iterator()
+	val itemItr = items.iterator()
+
+	while (slotItr.hasNext() && itemItr.hasNext()) {
+		container[slotItr.nextInt()] = itemItr.next()
+	}
+
+	if (itemItr.hasNext()) {
+		LOGGER.warn("Tried to overfill a container")
+	}
+}