From 8551e2fa64833ae9729fc48a176aff333b1ff222 Mon Sep 17 00:00:00 2001
From: DBotThePony <dbotthepony@yandex.ru>
Date: Fri, 30 Aug 2024 17:25:09 +0700
Subject: [PATCH] Fix machines not getting capabilities

---
 .../mc/otm/block/entity/MatteryBlockEntity.kt   |  3 +++
 .../ru/dbotthepony/mc/otm/registry/MBlocks.kt   | 17 +++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/src/main/kotlin/ru/dbotthepony/mc/otm/block/entity/MatteryBlockEntity.kt b/src/main/kotlin/ru/dbotthepony/mc/otm/block/entity/MatteryBlockEntity.kt
index 4b72a5937..bfb44d88f 100644
--- a/src/main/kotlin/ru/dbotthepony/mc/otm/block/entity/MatteryBlockEntity.kt
+++ b/src/main/kotlin/ru/dbotthepony/mc/otm/block/entity/MatteryBlockEntity.kt
@@ -52,6 +52,7 @@ import ru.dbotthepony.mc.otm.core.util.TickList
 import ru.dbotthepony.mc.otm.core.util.countingLazy
 import ru.dbotthepony.mc.otm.network.BlockEntitySyncPacket
 import ru.dbotthepony.mc.otm.onceServer
+import ru.dbotthepony.mc.otm.registry.MBlocks
 import ru.dbotthepony.mc.otm.sometimeServer
 import java.lang.ref.WeakReference
 import java.util.*
@@ -124,6 +125,7 @@ abstract class MatteryBlockEntity(p_155228_: BlockEntityType<*>, p_155229_: Bloc
 	 */
 	protected fun <T : Any> exposeSideless(capability: BlockCapability<T, *>, value: T) {
 		check(!sidelessCaps.containsKey(capability)) { "Already has globally exposed $capability!" }
+		MBlocks.ensureCapabilityIsKnown(capability)
 		sidelessCaps[capability] = value
 		setChanged()
 		level?.invalidateCapabilities(blockPos)
@@ -132,6 +134,7 @@ abstract class MatteryBlockEntity(p_155228_: BlockEntityType<*>, p_155229_: Bloc
 	protected fun <T : Any> exposeSided(side: RelativeSide, capability: BlockCapability<T, *>, value: T) {
 		val map = sidedCaps[side.ordinal]
 		check(!map.containsKey(capability)) { "Already has exposed $capability on $side!" }
+		MBlocks.ensureCapabilityIsKnown(capability)
 		map[capability] = value
 		setChanged()
 		level?.invalidateCapabilities(blockPos)
diff --git a/src/main/kotlin/ru/dbotthepony/mc/otm/registry/MBlocks.kt b/src/main/kotlin/ru/dbotthepony/mc/otm/registry/MBlocks.kt
index 0cb52c9d9..a82e451df 100644
--- a/src/main/kotlin/ru/dbotthepony/mc/otm/registry/MBlocks.kt
+++ b/src/main/kotlin/ru/dbotthepony/mc/otm/registry/MBlocks.kt
@@ -1,5 +1,6 @@
 package ru.dbotthepony.mc.otm.registry
 
+import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet
 import net.minecraft.core.Direction
 import net.minecraft.core.registries.BuiltInRegistries
 import net.minecraft.util.valueproviders.UniformInt
@@ -19,6 +20,7 @@ import net.minecraft.world.level.material.MapColor
 import net.minecraft.world.level.material.PushReaction
 import net.neoforged.bus.api.IEventBus
 import net.neoforged.neoforge.capabilities.BlockCapability
+import net.neoforged.neoforge.capabilities.Capabilities
 import net.neoforged.neoforge.capabilities.IBlockCapabilityProvider
 import net.neoforged.neoforge.capabilities.RegisterCapabilitiesEvent
 import ru.dbotthepony.mc.otm.block.BlackHoleBlock
@@ -77,6 +79,7 @@ import ru.dbotthepony.mc.otm.block.tech.PlatePressBlock
 import ru.dbotthepony.mc.otm.block.tech.PoweredBlastFurnaceBlock
 import ru.dbotthepony.mc.otm.block.tech.PoweredFurnaceBlock
 import ru.dbotthepony.mc.otm.block.tech.PoweredSmokerBlock
+import ru.dbotthepony.mc.otm.capability.MatteryCapability
 import ru.dbotthepony.mc.otm.config.CablesConfig
 import ru.dbotthepony.mc.otm.core.collect.SupplierList
 import ru.dbotthepony.mc.otm.core.collect.SupplierMap
@@ -85,9 +88,23 @@ import java.util.function.Supplier
 object MBlocks {
 	private val registry = MDeferredRegister(BuiltInRegistries.BLOCK)
 
+	private val knownCaps = ObjectOpenHashSet<BlockCapability<*, *>>()
+
+	fun ensureCapabilityIsKnown(cap: BlockCapability<*, *>) {
+		check(cap in knownCaps) { "Tried to expose capability which wasn't known on mod startup. This is unfortunately limitation with new capability system and can not be fixed, unless NeoForge team decides to rethink their decision to remove per block(entity) dynamic capabilities." }
+	}
+
 	private fun registerCapabilities(event: RegisterCapabilitiesEvent) {
+		// static initializers
+		MatteryCapability.BLOCK_ENERGY
+		Capabilities.FluidHandler.BLOCK
+		Capabilities.EnergyStorage.BLOCK
+		Capabilities.ItemHandler.BLOCK
+
 		// ugly
 		for (cap in BlockCapability.getAll()) {
+			knownCaps.add(cap)
+
 			val provider = IBlockCapabilityProvider<Any?, Direction?> { level, pos, state, be, context: Direction? ->
 				if (be is MatteryBlockEntity) {
 					return@IBlockCapabilityProvider be.getCapability(cap as BlockCapability<Any, Direction?>, context)