Global block/state tag, additional predicates for multiblock tests, more multiblock performance improvements

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

@@ -2,10 +2,16 @@ package ru.dbotthepony.mc.otm.core
 
 import com.google.common.collect.ImmutableList
 import com.google.common.collect.ImmutableSet
+import it.unimi.dsi.fastutil.objects.Object2IntMap
+import it.unimi.dsi.fastutil.objects.Object2IntMaps
 import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap
 import it.unimi.dsi.fastutil.objects.Object2ObjectFunction
 import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap
 import it.unimi.dsi.fastutil.objects.ObjectArraySet
+import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet
+import it.unimi.dsi.fastutil.objects.ObjectSets
+import it.unimi.dsi.fastutil.objects.Reference2IntMap
+import it.unimi.dsi.fastutil.objects.Reference2IntMaps
 import it.unimi.dsi.fastutil.objects.Reference2IntOpenHashMap
 import net.minecraft.core.BlockPos
 import net.minecraft.core.Direction
@@ -20,6 +26,7 @@ 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.SupplierList
 import ru.dbotthepony.mc.otm.core.collect.WeakHashSet
 import ru.dbotthepony.mc.otm.core.collect.collect
 import ru.dbotthepony.mc.otm.core.collect.map
@@ -40,6 +47,10 @@ fun onBlockEntityInvalidated(blockEntity: BlockEntity) {
 fun interface BlockPredicate {
 	fun test(pos: BlockPos, access: LevelAccessor): Boolean
 
+	fun rotate(rotation: Rotation): BlockPredicate {
+		return this
+	}
+
 	fun and(other: BlockPredicate): BlockPredicate {
 		return BlockPredicate { pos, access -> test(pos, access) && other.test(pos, access) }
 	}
@@ -107,13 +118,15 @@ inline fun <reified T : BlockEntity> multiblockEntity(): MultiblockBuilder.Entit
 	return MultiblockBuilder.EntityTag(T::class)
 }
 
+private val GLOBAL_BLOCK_TAG = Any()
+
 class MultiblockBuilder {
 	private val nodes = Object2ObjectOpenHashMap<BlockPos, Node>()
+	private val customChecks = ArrayList<Predicate<Multiblock>>()
 
-	// not data classes to allow having multiple tags with same target
+	class EntityTag<T : BlockEntity>(val clazz: KClass<T>, val predicate: Predicate<in T> = Predicate { true }) : Predicate<BlockEntity> {
-	class EntityTag<T : BlockEntity>(val clazz: KClass<T>) : Predicate<BlockEntity> {
 		override fun test(t: BlockEntity): Boolean {
-			return clazz.isInstance(t)
+			return clazz.isInstance(t) && predicate.test(t as T)
 		}
 	}
 
@@ -121,11 +134,19 @@ class MultiblockBuilder {
 		OR, AND
 	}
 
+	fun customCheck(predicate: Predicate<Multiblock>): MultiblockBuilder {
+		customChecks.add(predicate)
+		return this
+	}
+
 	@Suppress("unchecked_cast")
-	abstract class Part<T : Part<T>> {
+	abstract inner class Part<T : Part<T>> {
 		val predicates = ObjectArraySet<BlockPredicate>()
 		val children = ObjectArraySet<Part<*>>()
 
+		val builder: MultiblockBuilder
+			get() = this@MultiblockBuilder
+
 		val blockStateTags = ObjectArraySet<Any>()
 		val blockTags = ObjectArraySet<Any>()
 		val blockEntityTags = ObjectArraySet<EntityTag<*>>()
@@ -135,11 +156,21 @@ class MultiblockBuilder {
 			return this as T
 		}
 
+		fun tagBlockState(): T {
+			blockStateTags.add(GLOBAL_BLOCK_TAG)
+			return this as T
+		}
+
 		fun tagBlock(value: Any): T {
 			blockTags.add(value)
 			return this as T
 		}
 
+		fun tagBlock(): T {
+			blockTags.add(GLOBAL_BLOCK_TAG)
+			return this as T
+		}
+
 		fun tag(value: EntityTag<*>): T {
 			blockEntityTags.add(value)
 			return this as T
@@ -216,7 +247,7 @@ class MultiblockBuilder {
 		abstract val strategy: Strategy
 	}
 
-	class And<P : Part<P>>(val parent: P) : Part<And<P>>() {
+	inner class And<P : Part<P>>(val parent: P) : Part<And<P>>() {
 		init {
 			parent.children.add(this)
 		}
@@ -227,7 +258,7 @@ class MultiblockBuilder {
 			get() = Strategy.AND
 	}
 
-	class Or<P : Part<P>>(val parent: P) : Part<Or<P>>() {
+	inner class Or<P : Part<P>>(val parent: P) : Part<Or<P>>() {
 		init {
 			parent.children.add(this)
 		}
@@ -258,6 +289,18 @@ class MultiblockBuilder {
 			return relative(dir.default)
 		}
 
+		fun relative(diff: Vec3i, configurator: Node.() -> Unit): Node {
+			return node(pos + diff).also(configurator)
+		}
+
+		fun relative(dir: Direction, configurator: Node.() -> Unit): Node {
+			return relative(dir.normal).also(configurator)
+		}
+
+		fun relative(dir: RelativeSide, configurator: Node.() -> Unit): Node {
+			return relative(dir.default).also(configurator)
+		}
+
 		fun front() = relative(RelativeSide.FRONT)
 		fun back() = relative(RelativeSide.BACK)
 		fun left() = relative(RelativeSide.LEFT)
@@ -341,11 +384,11 @@ class MultiblockBuilder {
 	}
 
 	fun build(): MultiblockFactory {
-		return MultiblockFactory(nodes.values.iterator().map { it.build() }.collect(ImmutableSet.toImmutableSet()))
+		return MultiblockFactory(nodes.values.iterator().map { it.build() }.collect(ImmutableSet.toImmutableSet()), ImmutableList.copyOf(customChecks))
 	}
 }
 
-class MultiblockFactory(val north: ImmutableSet<Part>) {
+class MultiblockFactory(val north: ImmutableSet<Part>, val customChecks: ImmutableList<Predicate<Multiblock>>) {
 	data class Part(
 		val pos: BlockPos,
 		val strategy: MultiblockBuilder.Strategy,
@@ -357,13 +400,7 @@ class MultiblockFactory(val north: ImmutableSet<Part>) {
 	)
 
 	fun create(pos: BlockPos): Multiblock {
-		return Multiblock(
+		return Multiblock(pos, this)
-			pos,
-			north,
-			south,
-			west,
-			east,
-		)
 	}
 
 	val south: ImmutableSet<Part> = north.iterator().map { it.copy(pos = it.pos.rotate(Rotation.CLOCKWISE_180)) }.collect(ImmutableSet.toImmutableSet())
@@ -371,22 +408,22 @@ class MultiblockFactory(val north: ImmutableSet<Part>) {
 	val east: ImmutableSet<Part> = north.iterator().map { it.copy(pos = it.pos.rotate(Rotation.CLOCKWISE_90)) }.collect(ImmutableSet.toImmutableSet())
 }
 
-class Multiblock(
+class Multiblock(pos: BlockPos, factory: MultiblockFactory) {
-	pos: BlockPos,
-	north: Collection<MultiblockFactory.Part>,
-	south: Collection<MultiblockFactory.Part>,
-	west: Collection<MultiblockFactory.Part>,
-	east: Collection<MultiblockFactory.Part>,
-) {
 	var isValid = false
 		private set
 
-	private val configurations = ArrayList<Config>()
+	private val customChecks = factory.customChecks
+	private var collectFailedParts = false
+	private val north by lazy(LazyThreadSafetyMode.NONE) { Config(Direction.NORTH, pos, factory.north) }
+	private val south by lazy(LazyThreadSafetyMode.NONE) { Config(Direction.SOUTH, pos, factory.south) }
+	private val west by lazy(LazyThreadSafetyMode.NONE) { Config(Direction.WEST, pos, factory.west) }
+	private val east by lazy(LazyThreadSafetyMode.NONE) { Config(Direction.EAST, pos, factory.east) }
+	private val configurations = SupplierList(::north, ::south, ::west, ::east)
 
 	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)
+		val setView: Set<T> = ObjectSets.unmodifiable(set)
 
 		fun add(blockEntity: BlockEntity) {
 			if (tag.test(blockEntity)) {
@@ -453,9 +490,9 @@ class Multiblock(
 			private var blockEntity: BlockEntity? = null
 			private var blockState: BlockState? = null
 			private val assignedBlockEntityLists = ArrayList<BEList<*>>(prototype.blockEntityTags.size)
-			private val assignedBlockStateLists = ArrayList<Reference2IntOpenHashMap<BlockState>>()
+			private val assignedBlockStateLists = ArrayList<Reference2IntMap<BlockState>>()
-			private val assignedBlockLists = ArrayList<Object2IntOpenHashMap<Block>>()
+			private val assignedBlockLists = ArrayList<Reference2IntMap<Block>>()
-			private val children: ImmutableList<Part> = prototype.children.stream().map { Part(pos, it) }.collect(ImmutableList.toImmutableList())
+			private val children: ImmutableList<Part> by lazy(LazyThreadSafetyMode.NONE) { 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))
@@ -465,15 +502,15 @@ class Multiblock(
 
 			init {
 				prototype.blockEntityTags.forEach {
-					assignedBlockEntityLists.add(tag2BlockEntity.computeIfAbsent(it, Object2ObjectFunction { _ -> BEList(it as MultiblockBuilder.EntityTag<BlockEntity>).also { blockEntityLists.add(it) } }))
+					assignedBlockEntityLists.add(getBlockEntityList(it))
 				}
 
 				prototype.blockStateTags.forEach {
-					assignedBlockStateLists.add(tag2BlockState.computeIfAbsent(it, Object2ObjectFunction { Reference2IntOpenHashMap<BlockState>().also { blockStateLists.add(it) } }))
+					assignedBlockStateLists.add(getBlockStateList(it))
 				}
 
 				prototype.blockTags.forEach {
-					assignedBlockLists.add(tag2Block.computeIfAbsent(it, Object2ObjectFunction { Object2IntOpenHashMap<Block>().also { blockLists.add(it) } }))
+					assignedBlockLists.add(getBlockList(it))
 				}
 			}
 
@@ -609,13 +646,59 @@ class Multiblock(
 			}
 		}
 
+		private fun <T : BlockEntity> getBlockEntityList(tag: MultiblockBuilder.EntityTag<T>): BEList<T> {
+			val existing = tag2BlockEntity[tag]
+
+			if (existing != null)
+				return existing as BEList<T>
+
+			val new = BEList(tag)
+			blockEntityLists.add(new)
+			tag2BlockEntity[tag] = new
+			return new
+		}
+
+		private fun getBlockList(tag: Any): Reference2IntMap<Block> {
+			val existing = tag2Block[tag]
+
+			if (existing != null)
+				return existing
+
+			val new = Reference2IntOpenHashMap<Block>()
+			tag2Block[tag] = new
+			tag2BlockViews[tag] = Reference2IntMaps.unmodifiable(new)
+			blockLists.add(new)
+			return new
+		}
+
+		private fun getBlockStateList(tag: Any): Reference2IntMap<BlockState> {
+			val existing = tag2BlockState[tag]
+
+			if (existing != null)
+				return existing
+
+			val new = Reference2IntOpenHashMap<BlockState>()
+			tag2BlockState[tag] = new
+			tag2BlockStateViews[tag] = Reference2IntMaps.unmodifiable(new)
+			blockStateLists.add(new)
+			return new
+		}
+
 		private val tag2BlockEntity = Object2ObjectOpenHashMap<MultiblockBuilder.EntityTag<*>, BEList<*>>()
-		private val tag2BlockState = Object2ObjectOpenHashMap<Any, Reference2IntOpenHashMap<BlockState>>()
+		private val tag2BlockState = Object2ObjectOpenHashMap<Any, Reference2IntMap<BlockState>>()
-		private val tag2Block = Object2ObjectOpenHashMap<Any, Object2IntOpenHashMap<Block>>()
+		private val tag2Block = Object2ObjectOpenHashMap<Any, Reference2IntMap<Block>>()
+
+		private val tag2BlockStateViews = Object2ObjectOpenHashMap<Any, Reference2IntMap<BlockState>>()
+		private val tag2BlockViews = Object2ObjectOpenHashMap<Any, Reference2IntMap<Block>>()
 
 		private val blockEntityLists = ArrayList<BEList<*>>()
-		private val blockStateLists = ArrayList<Reference2IntOpenHashMap<BlockState>>()
+		private val blockStateLists = ArrayList<Reference2IntMap<BlockState>>()
-		private val blockLists = ArrayList<Object2IntOpenHashMap<Block>>()
+		private val blockLists = ArrayList<Reference2IntMap<Block>>()
+
+		private val globalBlocks = getBlockList(GLOBAL_BLOCK_TAG)
+		private val globalBlockStates = getBlockStateList(GLOBAL_BLOCK_TAG)
+		private val globalBlocksView = tag2BlockViews[GLOBAL_BLOCK_TAG]!!
+		private val globalBlockStatesView = tag2BlockStateViews[GLOBAL_BLOCK_TAG]!!
 
 		private val parts: ImmutableList<Part> = parts.stream()
 			.map { Part(it.pos + pos, it) }
@@ -647,6 +730,22 @@ class Multiblock(
 			return (tag2BlockEntity[tag]?.setView ?: setOf()) as Set<T>
 		}
 
+		fun blocks(tag: Any): Reference2IntMap<Block> {
+			return tag2BlockViews[tag] ?: Reference2IntMaps.emptyMap()
+		}
+
+		fun blockStates(tag: Any): Reference2IntMap<BlockState> {
+			return tag2BlockStateViews[tag] ?: Reference2IntMaps.emptyMap()
+		}
+
+		fun blocks(): Reference2IntMap<Block> {
+			return globalBlocksView
+		}
+
+		fun blockStates(): Reference2IntMap<BlockState> {
+			return globalBlockStatesView
+		}
+
 		fun blockEntityRemoved(blockEntity: BlockEntity): Boolean {
 			var any = false
 
@@ -658,11 +757,6 @@ class Multiblock(
 		}
 	}
 
-	private val north = Config(Direction.NORTH, pos, north).also { configurations.add(it) }
-	private val south = Config(Direction.SOUTH, pos, south).also { configurations.add(it) }
-	private val west = Config(Direction.WEST, pos, west).also { configurations.add(it) }
-	private val east = Config(Direction.EAST, pos, east).also { configurations.add(it) }
-
 	private var activeConfig: Config? = null
 
 	val currentDirection: Direction?
@@ -673,6 +767,26 @@ class Multiblock(
 		return activeConfig?.blockEntities(tag) ?: setOf()
 	}
 
+	fun blocks(tag: Any): Reference2IntMap<Block> {
+		if (!isValid) return Reference2IntMaps.emptyMap()
+		return activeConfig?.blocks(tag) ?: Reference2IntMaps.emptyMap()
+	}
+
+	fun blockStates(tag: Any): Reference2IntMap<BlockState> {
+		if (!isValid) return Reference2IntMaps.emptyMap()
+		return activeConfig?.blockStates(tag) ?: Reference2IntMaps.emptyMap()
+	}
+
+	fun blocks(): Reference2IntMap<Block> {
+		if (!isValid) return Reference2IntMaps.emptyMap()
+		return activeConfig?.blocks() ?: Reference2IntMaps.emptyMap()
+	}
+
+	fun blockStates(): Reference2IntMap<BlockState> {
+		if (!isValid) return Reference2IntMaps.emptyMap()
+		return activeConfig?.blockStates() ?: Reference2IntMaps.emptyMap()
+	}
+
 	fun blockEntityRemoved(blockEntity: BlockEntity) {
 		activeConfig?.blockEntityRemoved(blockEntity)
 	}
@@ -680,11 +794,11 @@ class Multiblock(
 	fun update(levelAccessor: LevelAccessor): Boolean {
 		val activeConfig = activeConfig
 
-		if (activeConfig != null && activeConfig.update(levelAccessor)) {
+		if (activeConfig != null && activeConfig.update(levelAccessor) && customChecks.all { it.test(this) }) {
 			return true
 		} else if (activeConfig != null) {
 			for (config in configurations) {
-				if (config !== activeConfig && config.update(levelAccessor)) {
+				if (config !== activeConfig && config.update(levelAccessor) && customChecks.all { it.test(this) }) {
 					this.activeConfig = config
 					return true
 				}
@@ -695,7 +809,7 @@ class Multiblock(
 			return false
 		} else {
 			for (config in configurations) {
-				if (config.update(levelAccessor)) {
+				if (config.update(levelAccessor) && customChecks.all { it.test(this) }) {
 					this.activeConfig = config
 					this.isValid = true
 					return true
@@ -720,7 +834,7 @@ class Multiblock(
 			else -> throw IllegalArgumentException(direction.name)
 		}
 
-		isValid = config.update(levelAccessor)
+		isValid = config.update(levelAccessor) && customChecks.all { it.test(this) }
 
 		if (isValid)
 			activeConfig = config