Force directional multiblock and multiblock performance improvements

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

@@ -9,14 +9,17 @@ import it.unimi.dsi.fastutil.objects.ObjectArraySet
 import it.unimi.dsi.fastutil.objects.Reference2IntOpenHashMap
 import net.minecraft.core.BlockPos
 import net.minecraft.core.Direction
+import net.minecraft.core.SectionPos
 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.Blocks
 import net.minecraft.world.level.block.Rotation
 import net.minecraft.world.level.block.entity.BlockEntity
 import net.minecraft.world.level.block.state.BlockState
+import net.minecraft.world.level.chunk.ChunkStatus
 import ru.dbotthepony.mc.otm.core.collect.WeakHashSet
 import ru.dbotthepony.mc.otm.core.collect.collect
 import ru.dbotthepony.mc.otm.core.collect.map
@@ -119,9 +122,9 @@ class MultiblockBuilder {
 	}
 
 	@Suppress("unchecked_cast")
-	abstract class BlockPredicates<T : BlockPredicates<T>> {
+	abstract class Part<T : Part<T>> {
 		val predicates = ObjectArraySet<BlockPredicate>()
-		val children = ObjectArraySet<BlockPredicates<*>>()
+		val children = ObjectArraySet<Part<*>>()
 
 		val blockStateTags = ObjectArraySet<Any>()
 		val blockTags = ObjectArraySet<Any>()
@@ -153,12 +156,20 @@ class MultiblockBuilder {
 		}
 
 		fun block(block: Block): T {
-			predicates.add { pos, access -> access.getBlockState(pos).`is`(block) }
+			predicates.add { pos, access ->
+				// use getChunk directly because we don't want to chunkload
+				(access.getChunk(SectionPos.blockToSectionCoord(pos.x), SectionPos.blockToSectionCoord(pos.z), ChunkStatus.FULL, false)?.getBlockState(pos) ?: Blocks.AIR.defaultBlockState()).`is`(block)
+			}
+
 			return this as T
 		}
 
 		fun block(block: TagKey<Block>): T {
-			predicates.add { pos, access -> access.getBlockState(pos).`is`(block) }
+			predicates.add { pos, access ->
+				// use getChunk directly because we don't want to chunkload
+				(access.getChunk(SectionPos.blockToSectionCoord(pos.x), SectionPos.blockToSectionCoord(pos.z), ChunkStatus.FULL, false)?.getBlockState(pos) ?: Blocks.AIR.defaultBlockState()).`is`(block)
+			}
+
 			return this as T
 		}
 
@@ -168,7 +179,11 @@ class MultiblockBuilder {
 		}
 
 		fun block(state: BlockState): T {
-			predicates.add { pos, access -> access.getBlockState(pos) == state }
+			predicates.add { pos, access ->
+				// use getChunk directly because we don't want to chunkload
+				access.getChunk(SectionPos.blockToSectionCoord(pos.x), SectionPos.blockToSectionCoord(pos.z), ChunkStatus.FULL, false)?.getBlockState(pos) == state
+			}
+
 			return this as T
 		}
 
@@ -201,7 +216,7 @@ class MultiblockBuilder {
 		abstract val strategy: Strategy
 	}
 
-	class And<P : BlockPredicates<P>>(val parent: P) : BlockPredicates<And<P>>() {
+	class And<P : Part<P>>(val parent: P) : Part<And<P>>() {
 		init {
 			parent.children.add(this)
 		}
@@ -212,7 +227,7 @@ class MultiblockBuilder {
 			get() = Strategy.AND
 	}
 
-	class Or<P : BlockPredicates<P>>(val parent: P) : BlockPredicates<Or<P>>() {
+	class Or<P : Part<P>>(val parent: P) : Part<Or<P>>() {
 		init {
 			parent.children.add(this)
 		}
@@ -226,7 +241,7 @@ class MultiblockBuilder {
 	/**
 	 * Behaves as if being [Or] node
 	 */
-	inner class Node(val pos: BlockPos) : BlockPredicates<Node>() {
+	inner class Node(val pos: BlockPos) : Part<Node>() {
 		init {
 			check(nodes.put(pos, this) == null) { "Trying to create new node at $pos while already having one" }
 		}
@@ -433,8 +448,8 @@ class Multiblock(
 		}
 	}
 
-	private inner class Config(val pos: BlockPos, parts: Collection<MultiblockFactory.Part>) {
+	private inner class Config(val direction: Direction, val pos: BlockPos, parts: Collection<MultiblockFactory.Part>) {
-		private inner class Part(val pos: BlockPos, val prototype: MultiblockFactory.Part) {
+		private inner class Part(val pos: BlockPos, val prototype: MultiblockFactory.Part) : Comparable<Part> {
 			private var blockEntity: BlockEntity? = null
 			private var blockState: BlockState? = null
 			private val assignedBlockEntityLists = ArrayList<BEList<*>>(prototype.blockEntityTags.size)
@@ -442,6 +457,12 @@ class Multiblock(
 			private val assignedBlockLists = ArrayList<Object2IntOpenHashMap<Block>>()
 			private val children: ImmutableList<Part> = prototype.children.stream().map { Part(pos, it) }.collect(ImmutableList.toImmutableList())
 
+			override fun compareTo(other: Part): Int {
+				val cmp = SectionPos.blockToSectionCoord(pos.x).compareTo(SectionPos.blockToSectionCoord(other.pos.x))
+				if (cmp != 0) return cmp
+				return SectionPos.blockToSectionCoord(pos.z).compareTo(SectionPos.blockToSectionCoord(other.pos.z))
+			}
+
 			init {
 				prototype.blockEntityTags.forEach {
 					assignedBlockEntityLists.add(tag2BlockEntity.computeIfAbsent(it, Object2ObjectFunction { _ -> BEList(it as MultiblockBuilder.EntityTag<BlockEntity>).also { blockEntityLists.add(it) } }))
@@ -596,7 +617,10 @@ class Multiblock(
 		private val blockStateLists = ArrayList<Reference2IntOpenHashMap<BlockState>>()
 		private val blockLists = ArrayList<Object2IntOpenHashMap<Block>>()
 
-		private val parts: ImmutableSet<Part> = parts.stream().map { Part(it.pos + pos, it) }.collect(ImmutableSet.toImmutableSet())
+		private val parts: ImmutableList<Part> = parts.stream()
+			.map { Part(it.pos + pos, it) }
+			.sorted() // group/localize parts by in-world chunks to maximize getChunk() cache hit rate
+			.collect(ImmutableList.toImmutableList())
 
 		fun clear() {
 			blockEntityLists.forEach { it.clear() }
@@ -634,13 +658,16 @@ class Multiblock(
 		}
 	}
 
-	private val north = Config(pos, north).also { configurations.add(it) }
+	private val north = Config(Direction.NORTH, pos, north).also { configurations.add(it) }
-	private val south = Config(pos, south).also { configurations.add(it) }
+	private val south = Config(Direction.SOUTH, pos, south).also { configurations.add(it) }
-	private val west = Config(pos, west).also { configurations.add(it) }
+	private val west = Config(Direction.WEST, pos, west).also { configurations.add(it) }
-	private val east = Config(pos, east).also { configurations.add(it) }
+	private val east = Config(Direction.EAST, pos, east).also { configurations.add(it) }
 
 	private var activeConfig: Config? = null
 
+	val currentDirection: Direction?
+		get() = activeConfig?.direction
+
 	fun <T : BlockEntity> blockEntities(tag: MultiblockBuilder.EntityTag<T>): Set<T> {
 		if (!isValid) return emptySet()
 		return activeConfig?.blockEntities(tag) ?: setOf()
@@ -678,4 +705,26 @@ class Multiblock(
 			return false
 		}
 	}
+
+	fun update(levelAccessor: LevelAccessor, direction: Direction): Boolean {
+		if (activeConfig?.direction != direction) {
+			activeConfig?.clear()
+			activeConfig = null
+		}
+
+		val config = when (direction) {
+			Direction.NORTH -> north
+			Direction.SOUTH -> south
+			Direction.WEST -> west
+			Direction.EAST -> east
+			else -> throw IllegalArgumentException(direction.name)
+		}
+
+		isValid = config.update(levelAccessor)
+
+		if (isValid)
+			activeConfig = config
+
+		return isValid
+	}
 }