Provide BlockPosSet for vastly improved memory efficiency

src/main/kotlin/ru/dbotthepony/mc/otm/core/collect/BlockPosSet.kt

@@ -0,0 +1,193 @@
+package ru.dbotthepony.mc.otm.core.collect
+
+import it.unimi.dsi.fastutil.objects.Object2ObjectFunction
+import it.unimi.dsi.fastutil.objects.Object2ObjectRBTreeMap
+import it.unimi.dsi.fastutil.objects.ObjectIterators
+import net.minecraft.core.BlockPos
+import net.minecraft.core.SectionPos
+import net.minecraft.world.level.ChunkPos
+import ru.dbotthepony.kommons.collect.iterateSetBits
+import java.util.BitSet
+import java.util.Collections
+
+class BlockPosSet : MutableSet<BlockPos> {
+	private class Segment {
+		val bits = BitSet(16 * 16 * 16)
+
+		private fun pos2index(pos: BlockPos): Int {
+			val x = pos.x.and(15)
+			val z = pos.z.and(15) shl 4
+			val y = pos.y.and(15) shl 8
+			return x or y or z
+		}
+
+		private fun index2pos(index: Int, gX: Int, gY: Int, gZ: Int): BlockPos {
+			val x = index.and(15) + gX
+			val z = index.ushr(4).and(15) + gZ
+			val y = index.ushr(8).and(15) + gY
+			return BlockPos(x, y, z)
+		}
+
+		operator fun contains(pos: BlockPos): Boolean {
+			return bits[pos2index(pos)]
+		}
+
+		fun add(pos: BlockPos): Boolean {
+			val index = pos2index(pos)
+
+			if (bits[index])
+				return false
+
+			bits[index] = true
+			return true
+		}
+
+		fun remove(pos: BlockPos): Boolean {
+			val index = pos2index(pos)
+
+			if (!bits[index])
+				return false
+
+			bits[index] = false
+			return true
+		}
+
+		fun iterator(gX: Int, gY: Int, gZ: Int): MutableIterator<BlockPos> {
+			return object : MutableIterator<BlockPos> {
+				private val iterator = bits.iterateSetBits(4096)
+				private var last = -1
+
+				override fun hasNext(): Boolean {
+					return iterator.hasNext()
+				}
+
+				override fun next(): BlockPos {
+					last = iterator.nextInt()
+					return index2pos(last, gX, gY, gZ)
+				}
+
+				override fun remove() {
+					if (last == -1)
+						throw NoSuchElementException()
+
+					bits[last] = false
+				}
+			}
+		}
+
+		val isEmpty: Boolean
+			get() = bits.isEmpty
+
+		val size: Int
+			get() = bits.cardinality()
+	}
+
+	private val segments = Object2ObjectRBTreeMap<SectionPos, Segment>(Vec3iHashStrategy)
+
+	fun subset(pos: ChunkPos): Set<BlockPos> {
+		val result = BlockPosSet()
+
+		for (key in segments.keys.iterator(SectionPos.of(pos.x, 0, pos.z))) {
+			if (key.x != pos.x || key.z != pos.z) break
+			result.segments[key] = segments[key]
+		}
+
+		return Collections.unmodifiableSet(result)
+	}
+
+	override fun add(element: BlockPos): Boolean {
+		return segments.computeIfAbsent(SectionPos.of(element), Object2ObjectFunction { Segment() }).add(element)
+	}
+
+	override val size: Int
+		get() = segments.values.sumOf { it.size }
+
+	override fun addAll(elements: Collection<BlockPos>): Boolean {
+		var any = false
+		elements.forEach { any = add(it) || any }
+		return any
+	}
+
+	override fun clear() {
+		segments.clear()
+	}
+
+	override fun contains(element: BlockPos): Boolean {
+		return segments[SectionPos.of(element)]?.contains(element) == true
+	}
+
+	override fun containsAll(elements: Collection<BlockPos>): Boolean {
+		return elements.all { contains(it) }
+	}
+
+	override fun isEmpty(): Boolean {
+		var any = false
+
+		segments.entries.removeIf { (_, s) ->
+			any = any || s.isEmpty
+			s.isEmpty
+		}
+
+		return any
+	}
+
+	override fun iterator(): MutableIterator<BlockPos> {
+		return object : MutableIterator<BlockPos> {
+			private val iterator = segments.entries.iterator()
+			private var current: MutableIterator<BlockPos> = ObjectIterators.emptyIterator()
+			private var last = current
+			private var foundNext = false
+
+			private fun findNext() {
+				if (!foundNext) {
+					foundNext = true
+
+					while (!current.hasNext() && iterator.hasNext()) {
+						val (pos, itr) = iterator.next()
+						current = itr.iterator(pos.x shl 4, pos.y shl 4, pos.z shl 4)
+					}
+				}
+			}
+
+			override fun hasNext(): Boolean {
+				findNext()
+				return current.hasNext()
+			}
+
+			override fun next(): BlockPos {
+				findNext()
+				last = current
+				val blockPos = current.next()
+				return blockPos
+			}
+
+			override fun remove() {
+				last.remove()
+			}
+		}
+	}
+
+	override fun remove(element: BlockPos): Boolean {
+		val index = SectionPos.of(element)
+		val segment = segments[index] ?: return false
+
+		if (segment.remove(element)) {
+			if (segment.isEmpty)
+				segments.remove(index)
+
+			return true
+		}
+
+		return false
+	}
+
+	override fun removeAll(elements: Collection<BlockPos>): Boolean {
+		var any = false
+		elements.forEach { any = remove(it) || any }
+		return any
+	}
+
+	override fun retainAll(elements: Collection<BlockPos>): Boolean {
+		return removeIf { it !in elements }
+	}
+}

src/main/kotlin/ru/dbotthepony/mc/otm/core/collect/Vec3iHashStrategy.kt

@@ -3,19 +3,20 @@ package ru.dbotthepony.mc.otm.core.collect
 import it.unimi.dsi.fastutil.Hash
 import it.unimi.dsi.fastutil.HashCommon
 import net.minecraft.core.BlockPos
+import net.minecraft.core.Vec3i
 
-object BlockPosHashStrategy : Hash.Strategy<BlockPos>, Comparator<BlockPos> {
+object Vec3iHashStrategy : Hash.Strategy<Vec3i>, Comparator<Vec3i> {
-	override fun equals(a: BlockPos?, b: BlockPos?): Boolean {
+	override fun equals(a: Vec3i?, b: Vec3i?): Boolean {
 		return a == b
 	}
 
 	// while this avoids collisions, it is rather slow when compared to just using a tree set here
-	override fun hashCode(o: BlockPos?): Int {
+	override fun hashCode(o: Vec3i?): Int {
 		o ?: return 0
 		return HashCommon.murmurHash3(o.x.toLong().and(1 shl 26 - 1) or o.z.toLong().and(1 shl 26 - 1).shl(26) or o.y.toLong().and(1 shl 12 - 1).shl(52)).toInt()
 	}
 
-	override fun compare(o1: BlockPos?, o2: BlockPos?): Int {
+	override fun compare(o1: Vec3i?, o2: Vec3i?): Int {
 		if (o1 == null && o2 == null)
 			return 0
 		else if (o1 == null)

src/main/kotlin/ru/dbotthepony/mc/otm/worldgen/EnhancedPlacedFeature.kt

@@ -25,14 +25,14 @@ import net.minecraft.world.level.levelgen.feature.configurations.FeatureConfigur
 import net.minecraft.world.level.levelgen.placement.PlacedFeature
 import net.minecraft.world.level.levelgen.placement.PlacementModifier
 import ru.dbotthepony.kommons.util.XXHash64
-import ru.dbotthepony.mc.otm.core.collect.BlockPosHashStrategy
+import ru.dbotthepony.mc.otm.core.collect.BlockPosSet
+import ru.dbotthepony.mc.otm.core.collect.Vec3iHashStrategy
 import ru.dbotthepony.mc.otm.core.util.GJRAND64RandomSource
 import ru.dbotthepony.mc.otm.data.codec.minRange
 import ru.dbotthepony.mc.otm.worldgen.feature.EnhancedFeature
 import ru.dbotthepony.mc.otm.worldgen.placement.EnhancedPlacementModifier
 import java.io.DataOutputStream
 import java.time.Duration
-import java.util.IdentityHashMap
 import kotlin.math.sqrt
 
 private object NodeCodec : Codec<EnhancedPlacedFeature.Node> {
@@ -140,13 +140,13 @@ object EnhancedPlacedFeature : Feature<EnhancedPlacedFeature.Config>(
 		val root: Node,
 	) : FeatureConfiguration {
 		private class GeneratedChunk : EnhancedPlacementContext.Placer {
-			private data class Placement(val context: EnhancedPlacementContext, val positions: Set<BlockPos>, val feature: EnhancedFeature.Configured<*, *>)
+			private data class Placement(val context: EnhancedPlacementContext, val positions: BlockPosSet, val feature: EnhancedFeature.Configured<*, *>)
 
-			// TODO: extremely inefficient
+			// TODO: inefficient
 			private val placed = ArrayList<Placement>()
 
 			override fun place(context: EnhancedPlacementContext, positions: List<BlockPos>, feature: EnhancedFeature.Configured<*, *>) {
-				placed.add(Placement(context, ObjectRBTreeSet(BlockPosHashStrategy).also { it.addAll(positions) }, feature))
+				placed.add(Placement(context, BlockPosSet().also { it.addAll(positions) }, feature))
 			}
 
 			fun place(context: FeaturePlaceContext<*>): Boolean {
@@ -154,10 +154,10 @@ object EnhancedPlacedFeature : Feature<EnhancedPlacedFeature.Config>(
 				val pos = ChunkPos(context.origin())
 
 				for ((eContext, positions, feature) in placed) {
-					val filtered = positions.filter { SectionPos.blockToSectionCoord(it.x) == pos.x && SectionPos.blockToSectionCoord(it.z) == pos.z }.toTypedArray()
+					val filtered = positions.subset(pos)
 
 					if (filtered.isNotEmpty()) {
-						any = feature.place(eContext, ObjectArraySet.ofUnchecked(*filtered), positions) || any
+						any = feature.place(eContext, filtered, positions) || any
 					}
 				}
 

src/main/kotlin/ru/dbotthepony/mc/otm/worldgen/placement/EnhancedPlacementModifier.kt

@@ -2,13 +2,9 @@ package ru.dbotthepony.mc.otm.worldgen.placement
 
 import com.mojang.serialization.Codec
 import com.mojang.serialization.MapCodec
-import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenCustomHashSet
-import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet
-import it.unimi.dsi.fastutil.objects.ObjectRBTreeSet
 import net.minecraft.core.BlockPos
 import net.minecraft.world.level.levelgen.placement.PlacementModifier
 import net.neoforged.bus.api.IEventBus
-import ru.dbotthepony.mc.otm.core.collect.BlockPosHashStrategy
 import ru.dbotthepony.mc.otm.registry.MBuiltInRegistries
 import ru.dbotthepony.mc.otm.registry.MDeferredRegister
 import ru.dbotthepony.mc.otm.registry.MRegistries