:mind_blown: More multiblock code

src/main/java/ru/dbotthepony/mc/otm/mixin/BlockEntityMixin.java

@@ -0,0 +1,23 @@
+package ru.dbotthepony.mc.otm.mixin;
+
+import net.minecraft.world.level.block.entity.BlockEntity;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import ru.dbotthepony.mc.otm.core.MultiblockKt;
+
+// because i know
+// someone
+// somewhere
+// will try to break the system
+// either on purpose or accidentally because of other mod
+@Mixin(BlockEntity.class)
+public abstract class BlockEntityMixin {
+	@Inject(
+		at = @At("TAIL"),
+		method = "setRemoved()V"
+	)
+	public void setRemovedListener() {
+		MultiblockKt.onBlockEntityInvalidated((BlockEntity) (Object) this);
+	}
+}

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

@@ -11,19 +11,29 @@ import net.minecraft.core.BlockPos
 import net.minecraft.core.Direction
 import net.minecraft.core.Vec3i
 import net.minecraft.tags.TagKey
+import net.minecraft.world.level.Level
 import net.minecraft.world.level.LevelAccessor
 import net.minecraft.world.level.block.Block
 import net.minecraft.world.level.block.Rotation
 import net.minecraft.world.level.block.entity.BlockEntity
 import net.minecraft.world.level.block.state.BlockState
+import ru.dbotthepony.mc.otm.core.collect.WeakHashSet
 import ru.dbotthepony.mc.otm.core.collect.collect
 import ru.dbotthepony.mc.otm.core.collect.map
 import ru.dbotthepony.mc.otm.core.math.RelativeSide
 import ru.dbotthepony.mc.otm.core.math.plus
 import java.util.Collections
+import java.util.WeakHashMap
 import java.util.function.Predicate
 import kotlin.reflect.KClass
 
+private val blockEntityListeners = WeakHashMap<Level, WeakHashMap<BlockEntity, WeakHashSet<Multiblock>>>()
+
+fun onBlockEntityInvalidated(blockEntity: BlockEntity) {
+	blockEntityListeners[blockEntity.level]?.get(blockEntity)?.forEach { it.blockEntityRemoved(blockEntity) }
+	blockEntityListeners[blockEntity.level]?.remove(blockEntity)
+}
+
 fun interface BlockPredicate {
 	fun test(pos: BlockPos, access: LevelAccessor): Boolean
 
@@ -358,14 +368,20 @@ class Multiblock(
 
 	private val configurations = ArrayList<Config>()
 
-	private class BEList<T : BlockEntity>(val tag: MultiblockBuilder.EntityTag<T>) {
+	private inner class BEList<T : BlockEntity>(val tag: MultiblockBuilder.EntityTag<T>) {
 		val list = ArrayList<T>()
 		val set = ObjectArraySet<T>()
 		val setView: Set<T> = Collections.unmodifiableSet(set)
 
 		fun add(blockEntity: BlockEntity) {
 			if (tag.test(blockEntity)) {
-				set.add(blockEntity as T)
+				if (set.add(blockEntity as T)) {
+					blockEntityListeners
+						.computeIfAbsent(blockEntity.level) { WeakHashMap() }
+						.computeIfAbsent(blockEntity) { WeakHashSet() }
+						.add(this@Multiblock)
+				}
+
 				list.add(blockEntity)
 			}
 		}
@@ -376,11 +392,42 @@ class Multiblock(
 
 				if (blockEntity !in list) {
 					set.remove(blockEntity)
+
+					val getSet = blockEntityListeners[blockEntity.level]?.get(blockEntity)
+
+					if (getSet != null) {
+						getSet.remove(this@Multiblock)
+
+						if (getSet.isEmpty()) {
+							blockEntityListeners[blockEntity.level]?.remove(blockEntity)
+						}
+					}
 				}
 			}
 		}
 
+		fun blockEntityRemoved(blockEntity: BlockEntity): Boolean {
+			if (blockEntity in list) {
+				while (list.remove(blockEntity)) {}
+				set.remove(blockEntity)
+
+				return true
+			}
+
+			return false
+		}
+
 		fun clear() {
+			set.forEach {
+				val getSet = checkNotNull(blockEntityListeners[it.level]) { "Consistency check failed: No subscriber lists for level ${it.level}" }.get(it)
+				checkNotNull(getSet) { "Consistency check failed: No subscriber list for $it" }
+				check(getSet.remove(this@Multiblock)) { "Consistency check failed: Can't remove ${this@Multiblock} from $it subscriber list" }
+
+				if (getSet.isEmpty()) {
+					blockEntityListeners[it.level]?.remove(it)
+				}
+			}
+
 			list.clear()
 			set.clear()
 		}
@@ -471,12 +518,12 @@ class Multiblock(
 							if (old != null) {
 								assignedBlockStateLists.forEach {
 									it[old] = it.getInt(old) - 1
-									check(it.getInt(old) >= 0) { "Counter for block state $old turned negative" }
+									check(it.getInt(old) >= 0) { "Consistency check failed: Counter for block state $old turned negative" }
 								}
 
 								assignedBlockLists.forEach {
 									it[old.block] = it.getInt(old.block) - 1
-									check(it.getInt(old.block) >= 0) { "Counter for block ${old.block.registryName} turned negative" }
+									check(it.getInt(old.block) >= 0) { "Consistency check failed: Counter for block ${old.block.registryName} turned negative" }
 								}
 							}
 
@@ -510,12 +557,12 @@ class Multiblock(
 				if (blockState != null) {
 					assignedBlockStateLists.forEach {
 						it[blockState] = it.getInt(blockState) - 1
-						check(it.getInt(blockState) >= 0) { "Counter for block state $blockState turned negative" }
+						check(it.getInt(blockState) >= 0) { "Consistency check failed: Counter for block state $blockState turned negative" }
 					}
 
 					assignedBlockLists.forEach {
 						it[blockState.block] = it.getInt(blockState.block) - 1
-						check(it.getInt(blockState.block) >= 0) { "Counter for block ${blockState.block.registryName} turned negative" }
+						check(it.getInt(blockState.block) >= 0) { "Consistency check failed: Counter for block ${blockState.block.registryName} turned negative" }
 					}
 				}
 
@@ -575,6 +622,16 @@ class Multiblock(
 		fun <T : BlockEntity> blockEntities(tag: MultiblockBuilder.EntityTag<T>): Set<T> {
 			return (tag2BlockEntity[tag]?.setView ?: setOf()) as Set<T>
 		}
+
+		fun blockEntityRemoved(blockEntity: BlockEntity): Boolean {
+			var any = false
+
+			blockEntityLists.forEach {
+				any = it.blockEntityRemoved(blockEntity) || any
+			}
+
+			return any
+		}
 	}
 
 	private val north = Config(pos, north).also { configurations.add(it) }
@@ -585,9 +642,14 @@ class Multiblock(
 	private var activeConfig: Config? = null
 
 	fun <T : BlockEntity> blockEntities(tag: MultiblockBuilder.EntityTag<T>): Set<T> {
+		if (!isValid) return emptySet()
 		return activeConfig?.blockEntities(tag) ?: setOf()
 	}
 
+	fun blockEntityRemoved(blockEntity: BlockEntity) {
+		activeConfig?.blockEntityRemoved(blockEntity)
+	}
+
 	fun update(levelAccessor: LevelAccessor): Boolean {
 		val activeConfig = activeConfig
 

src/main/resources/overdrive_that_matters.mixins.json

@@ -6,6 +6,7 @@
   "minVersion": "0.8",
   "refmap": "overdrive_that_matters.refmap.json",
   "mixins": [
+    "BlockEntityMixin",
     "EnchantmentHelperMixin",
     "FoodDataMixin",
     "MixinPatchProjectileFinder",