Per-block variable storage in enhanced placement

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

@@ -116,10 +116,10 @@ object EnhancedPlacedFeature : Feature<EnhancedPlacedFeature.Config>(
 ) {
 	class Node(val children: List<Node>, val contents: Either<EnhancedPlacementModifier, Holder<EnhancedFeature.Configured<*, *>>>){
 		fun evaluate(context: EnhancedPlacementContext) {
-			evaluate(context, listOf(BlockPos(context.origin.minBlockX, 0, context.origin.minBlockZ)))
+			evaluate(context, listOf(PlacementPos(BlockPos(context.origin.minBlockX, 0, context.origin.minBlockZ), emptyVariableMap)))
 		}
 
-		private fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>) {
+		private fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>) {
 			val actualPositions = if (contents.left().isPresent) {
 				contents.left().get().evaluate(context, positions)
 			} else {
@@ -136,18 +136,20 @@ object EnhancedPlacedFeature : Feature<EnhancedPlacedFeature.Config>(
 		}
 	}
 
+	private val emptyVariableMap = PlacementVariableMap()
+
 	class Config(
 		val chunkScanRange: Int,
 		val seedMix: Long,
 		val root: Node,
 	) : FeatureConfiguration {
 		private class GeneratedChunk : EnhancedPlacementContext.Placer {
-			private data class Placement(val context: EnhancedPlacementContext, val positions: BlockPosSet, val feature: EnhancedFeature.Configured<*, *>)
+			private data class Placement(val context: EnhancedPlacementContext, val positions: ObjectRBTreeSet<PlacementPos>, val feature: EnhancedFeature.Configured<*, *>)
 			private val placed = LinkedList<Placement>()
 
-			override fun place(context: EnhancedPlacementContext, positions: List<BlockPos>, feature: EnhancedFeature.Configured<*, *>) {
+			override fun place(context: EnhancedPlacementContext, positions: List<PlacementPos>, feature: EnhancedFeature.Configured<*, *>) {
 				if (positions.isNotEmpty()) {
-					placed.add(Placement(context, BlockPosSet().also { it.addAll(positions) }, feature))
+					placed.add(Placement(context, ObjectRBTreeSet<PlacementPos>().also { it.addAll(positions) }, feature))
 				}
 			}
 
@@ -156,10 +158,22 @@ object EnhancedPlacedFeature : Feature<EnhancedPlacedFeature.Config>(
 				val pos = ChunkPos(context.origin())
 
 				for ((eContext, positions, feature) in placed) {
-					val filtered = positions.subset(pos)
+					val itr = positions.iterator(PlacementPos(BlockPos(pos.minBlockX, Int.MIN_VALUE, pos.minBlockZ), emptyVariableMap))
 
-					if (filtered.isNotEmpty()) {
+					if (!itr.hasNext())
-						any = feature.place(eContext.push(context.level()), filtered, positions) || any
+						continue
+
+					val result = ObjectRBTreeSet<PlacementPos>()
+
+					for (placementPos in itr) {
+						if (SectionPos.blockToSectionCoord(placementPos.pos.x) != pos.x || SectionPos.blockToSectionCoord(placementPos.pos.z) != pos.z)
+							break
+
+						result.add(placementPos)
+					}
+
+					if (result.isNotEmpty()) {
+						any = feature.place(eContext.push(context.level()), result, positions) || any
 					}
 				}
 

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

@@ -16,7 +16,7 @@ import kotlin.collections.HashMap
 
 class EnhancedPlacementContext {
 	fun interface Placer {
-		fun place(context: EnhancedPlacementContext, positions: List<BlockPos>, feature: EnhancedFeature.Configured<*, *>)
+		fun place(context: EnhancedPlacementContext, positions: List<PlacementPos>, feature: EnhancedFeature.Configured<*, *>)
 	}
 
 	private data class SharedState(
@@ -43,29 +43,7 @@ class EnhancedPlacementContext {
 	val vanillaContext: PlacementContext
 		get() = state.vanillaContext
 
-	private val parent: EnhancedPlacementContext?
+	val variables: PlacementVariableMap
-	private val variables = Reference2ObjectOpenHashMap<PlacementVariable<*>, KOptional<*>>(0)
-
-	private fun <T> recursiveGet(index: PlacementVariable<T>): KOptional<T>? {
-		return variables[index] as KOptional<T>? ?: parent?.recursiveGet(index)
-	}
-
-	operator fun <T> get(index: PlacementVariable<T>): KOptional<T> {
-		return recursiveGet(index) ?: KOptional()
-	}
-
-	operator fun <T> set(index: PlacementVariable<T>, value: T) {
-		variables[index] = KOptional(value)
-	}
-
-	fun <T> remove(index: PlacementVariable<T>): KOptional<T> {
-		val old = variables.put(index, KOptional<T>())
-		return old as KOptional<T>? ?: KOptional()
-	}
-
-	fun remove(index: Collection<PlacementVariable<*>>) {
-		index.forEach { remove(it) }
-	}
 
 	constructor(
 		level: WorldGenLevel,
@@ -83,24 +61,24 @@ class EnhancedPlacementContext {
 			vanillaContext = PlacementContext(level, generator, Optional.empty())
 		)
 
-		parent = null
+		variables = PlacementVariableMap()
 	}
 
 	private constructor(parent: EnhancedPlacementContext) {
 		this.state = parent.state
-		this.parent = parent
+		this.variables = parent.variables.push()
 	}
 
 	private constructor(parent: EnhancedPlacementContext, context: WorldGenLevel) {
 		this.state = parent.state.copy(level = context)
-		this.parent = parent
+		this.variables = parent.variables.push()
 	}
 
 	fun push(): EnhancedPlacementContext {
 		return EnhancedPlacementContext(this)
 	}
 
-	fun place(positions: List<BlockPos>, feature: EnhancedFeature.Configured<*, *>) {
+	fun place(positions: List<PlacementPos>, feature: EnhancedFeature.Configured<*, *>) {
 		placer.place(this, positions, feature)
 	}
 

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

@@ -1,8 +1,16 @@
 package ru.dbotthepony.mc.otm.worldgen
 
+import net.minecraft.core.BlockPos
 import net.minecraft.world.level.levelgen.placement.PlacementModifier
+import ru.dbotthepony.mc.otm.core.collect.Vec3iHashStrategy
 import ru.dbotthepony.mc.otm.worldgen.placement.EnhancedPlacementModifier
 
 fun PlacementModifier.wrap(): EnhancedPlacementModifier {
 	return EnhancedPlacementModifier.Wrapper(this)
 }
+
+data class PlacementPos(val pos: BlockPos, val variables: PlacementVariableMap) : Comparable<PlacementPos> {
+	override fun compareTo(other: PlacementPos): Int {
+		return Vec3iHashStrategy.compare(pos, other.pos)
+	}
+}

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

@@ -2,8 +2,8 @@ package ru.dbotthepony.mc.otm.worldgen
 
 import net.minecraft.resources.ResourceLocation
 
-class PlacementVariable<T>(val name: ResourceLocation) {
+class PlacementVariable<T>(val name: ResourceLocation) : Comparable<PlacementVariable<T>> {
-	companion object {
+	override fun compareTo(other: PlacementVariable<T>): Int {
-
+		return name.compareTo(other.name)
 	}
 }

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

@@ -0,0 +1,56 @@
+package ru.dbotthepony.mc.otm.worldgen
+
+import it.unimi.dsi.fastutil.objects.Object2ObjectAVLTreeMap
+import ru.dbotthepony.kommons.util.KOptional
+
+/**
+ * This is a "map stack" of sorts, and when you want to push some data to it
+ * you need to push new map to stack (by calling [push]) and then write new data to
+ * the map returned by latter.
+ */
+class PlacementVariableMap {
+	private val parent: PlacementVariableMap?
+	private val variables = Object2ObjectAVLTreeMap<PlacementVariable<*>, KOptional<*>>()
+
+	constructor() {
+		parent = null
+	}
+
+	private constructor(parent: PlacementVariableMap) {
+		this.parent = parent
+	}
+
+	private fun <T> recursiveGet(index: PlacementVariable<T>): KOptional<T>? {
+		return variables[index] as KOptional<T>? ?: parent?.recursiveGet(index)
+	}
+
+	operator fun <T> get(index: PlacementVariable<T>): KOptional<T> {
+		return recursiveGet(index) ?: KOptional()
+	}
+
+	operator fun <T> set(index: PlacementVariable<T>, value: T): PlacementVariableMap {
+		variables[index] = KOptional(value)
+		return this
+	}
+
+	fun <T> remove(index: PlacementVariable<T>): KOptional<T> {
+		val old = variables.put(index, KOptional<T>())
+		return old as KOptional<T>? ?: KOptional()
+	}
+
+	fun remove(index: Collection<PlacementVariable<*>>) {
+		index.forEach { remove(it) }
+	}
+
+	fun push(): PlacementVariableMap {
+		return PlacementVariableMap(this)
+	}
+
+	override fun equals(other: Any?): Boolean {
+		return other === this || other is PlacementVariableMap && variables == other.variables && parent == other.parent
+	}
+
+	override fun hashCode(): Int {
+		return parent.hashCode() * 31 + variables.hashCode() + 5
+	}
+}

src/main/kotlin/ru/dbotthepony/mc/otm/worldgen/feature/EnhancedFeature.kt

@@ -12,12 +12,13 @@ import ru.dbotthepony.mc.otm.registry.MBuiltInRegistries
 import ru.dbotthepony.mc.otm.registry.MDeferredRegister
 import ru.dbotthepony.mc.otm.registry.MRegistries
 import ru.dbotthepony.mc.otm.worldgen.EnhancedPlacementContext
+import ru.dbotthepony.mc.otm.worldgen.PlacementPos
 
 abstract class EnhancedFeature<FC>(codec: Codec<FC>) {
-	abstract fun place(context: EnhancedPlacementContext, config: FC, positions: Set<BlockPos>, allPositions: Set<BlockPos>): Boolean
+	abstract fun place(context: EnhancedPlacementContext, config: FC, positions: Set<PlacementPos>, allPositions: Set<PlacementPos>): Boolean
 
 	data class Configured<F : EnhancedFeature<FC>, FC>(val feature: F, val config: FC) {
-		fun place(context: EnhancedPlacementContext, positions: Set<BlockPos>, allPositions: Set<BlockPos>): Boolean {
+		fun place(context: EnhancedPlacementContext, positions: Set<PlacementPos>, allPositions: Set<PlacementPos>): Boolean {
 			return feature.place(context, config, positions, allPositions)
 		}
 	}
@@ -29,11 +30,11 @@ abstract class EnhancedFeature<FC>(codec: Codec<FC>) {
 		override fun place(
 			context: EnhancedPlacementContext,
 			config: Holder<ConfiguredFeature<*, *>>,
-			positions: Set<BlockPos>,
+			positions: Set<PlacementPos>,
-			allPositions: Set<BlockPos>
+			allPositions: Set<PlacementPos>
 		): Boolean {
 			var any = false
-			positions.forEach { any = config.value().place(context.level, context.generator, context.random, it) || any }
+			positions.forEach { any = config.value().place(context.level, context.generator, context.random, it.pos) || any }
 			return any
 		}
 

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

@@ -13,6 +13,8 @@ import net.minecraft.world.level.levelgen.placement.PlacementModifierType
 import ru.dbotthepony.mc.otm.core.collect.Vec3iHashStrategy
 import ru.dbotthepony.mc.otm.registry.data.MPlacementModifiers
 import ru.dbotthepony.mc.otm.worldgen.EnhancedPlacementContext
+import ru.dbotthepony.mc.otm.worldgen.PlacementPos
+import ru.dbotthepony.mc.otm.worldgen.PlacementVariableMap
 import java.util.stream.Stream
 
 /**
@@ -103,9 +105,9 @@ data class EllipsoidPlacement(
 		return evaluate(random, position).stream()
 	}
 
-	override fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos> {
+	override fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>): List<PlacementPos> {
-		val result = ArrayList<BlockPos>()
+		val result = ArrayList<PlacementPos>()
-		val results = positions.map { evaluate(context.random, it) }
+		val results = positions.map { (it, c) -> evaluate(context.random, it).map { PlacementPos(it, c) } }
 		result.ensureCapacity(result.size + results.sumOf { it.size } + 1)
 		results.forEach { result.addAll(it) }
 		return result

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

@@ -6,6 +6,7 @@ import com.mojang.serialization.MapCodec
 import net.minecraft.core.BlockPos
 import net.minecraft.world.level.levelgen.placement.PlacementModifier
 import ru.dbotthepony.mc.otm.worldgen.EnhancedPlacementContext
+import ru.dbotthepony.mc.otm.worldgen.PlacementPos
 
 /**
  * Daisy-chaining placements. Required only when using placement modifiers which operate on children list
@@ -15,7 +16,7 @@ class EnhancedChainPlacement(
 ) : EnhancedPlacementModifier {
 	constructor(vararg children: EnhancedPlacementModifier) : this(ImmutableList.copyOf(children))
 
-	override fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos> {
+	override fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>): List<PlacementPos> {
 		var current = positions
 		children.forEach { current = it.evaluate(context, current) }
 		return current

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

@@ -1,12 +1,12 @@
 package ru.dbotthepony.mc.otm.worldgen.placement
 
 import com.mojang.serialization.MapCodec
-import net.minecraft.core.BlockPos
 import net.minecraft.util.valueproviders.IntProvider
 import ru.dbotthepony.mc.otm.worldgen.EnhancedPlacementContext
+import ru.dbotthepony.mc.otm.worldgen.PlacementPos
 
 class EnhancedCountPlacement(val provider: IntProvider) : EnhancedPlacementModifier {
-	override fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos> {
+	override fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>): List<PlacementPos> {
 		val count = provider.sample(context.random)
 
 		if (count <= 0) {
@@ -14,7 +14,7 @@ class EnhancedCountPlacement(val provider: IntProvider) : EnhancedPlacementModif
 		} else if (count == 1) {
 			return positions
 		} else {
-			val result = ArrayList<BlockPos>()
+			val result = ArrayList<PlacementPos>()
 			result.ensureCapacity(positions.size * count)
 			for (i in 0 until count) result.addAll(positions)
 			return result

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

@@ -3,12 +3,15 @@ package ru.dbotthepony.mc.otm.worldgen.placement
 import com.mojang.serialization.Codec
 import com.mojang.serialization.MapCodec
 import net.minecraft.core.BlockPos
+import net.minecraft.world.level.levelgen.placement.PlacementContext
 import net.minecraft.world.level.levelgen.placement.PlacementModifier
 import net.neoforged.bus.api.IEventBus
 import ru.dbotthepony.mc.otm.registry.MBuiltInRegistries
 import ru.dbotthepony.mc.otm.registry.MDeferredRegister
 import ru.dbotthepony.mc.otm.registry.MRegistries
 import ru.dbotthepony.mc.otm.worldgen.EnhancedPlacementContext
+import ru.dbotthepony.mc.otm.worldgen.PlacementPos
+import ru.dbotthepony.mc.otm.worldgen.PlacementVariableMap
 import java.util.stream.Collectors
 
 interface EnhancedPlacementModifier {
@@ -16,14 +19,13 @@ interface EnhancedPlacementModifier {
 		val codec: MapCodec<T>
 	}
 
-	fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos>
+	fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>): List<PlacementPos>
 	val type: Type<*>
 
 	class Wrapper(val parent: PlacementModifier) : EnhancedPlacementModifier {
-		override fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos> {
+		override fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>): List<PlacementPos> {
 			return positions.stream()
-				.flatMap { parent.getPositions(context.vanillaContext, context.random, it) }
+				.flatMap { (pos, vars) -> parent.getPositions(context.vanillaContext, context.random, pos).map { PlacementPos(it, vars) } }
-				// use Red-Black tree set instead of AVL tree set because we are write-intense
 				.collect(Collectors.toCollection(::ArrayList))
 		}
 
@@ -36,7 +38,7 @@ interface EnhancedPlacementModifier {
 	}
 
 	object Passthrough : EnhancedPlacementModifier, Type<Passthrough> {
-		override fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos> {
+		override fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>): List<PlacementPos> {
 			return positions
 		}
 

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

@@ -10,6 +10,7 @@ import net.minecraft.world.level.levelgen.placement.PlacementModifier
 import net.minecraft.world.level.levelgen.placement.PlacementModifierType
 import ru.dbotthepony.mc.otm.registry.data.MPlacementModifiers
 import ru.dbotthepony.mc.otm.worldgen.EnhancedPlacementContext
+import ru.dbotthepony.mc.otm.worldgen.PlacementPos
 import java.util.stream.Collectors
 import java.util.stream.Stream
 
@@ -23,8 +24,8 @@ class EnhancedSplitPlacement(
 ) : EnhancedPlacementModifier {
 	constructor(vararg children: EnhancedPlacementModifier) : this(ImmutableList.copyOf(children))
 
-	override fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos> {
+	override fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>): List<PlacementPos> {
-		val result = ArrayList<BlockPos>()
+		val result = ArrayList<PlacementPos>()
 		children.forEach { result.addAll(it.evaluate(context, positions)) }
 		return result
 	}

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

@@ -24,6 +24,7 @@ import ru.dbotthepony.mc.otm.data.codec.minRange
 import ru.dbotthepony.mc.otm.data.world.BooleanProvider
 import ru.dbotthepony.mc.otm.registry.data.MPlacementModifiers
 import ru.dbotthepony.mc.otm.worldgen.EnhancedPlacementContext
+import ru.dbotthepony.mc.otm.worldgen.PlacementPos
 import java.util.stream.Stream
 import kotlin.math.PI
 import kotlin.math.absoluteValue
@@ -135,9 +136,10 @@ class WormPlacement(
 		}
 	}
 
-	private fun evaluate(random: RandomSource, position: BlockPos, results: ArrayList<BlockPos>) {
+	private fun evaluate(random: RandomSource, position: BlockPos): List<BlockPos> {
 		val worms = ArrayList<Worm>()
 		worms.add(Worm(random, Vector.ZERO))
+		val results = ArrayList<BlockPos>()
 		results.add(position)
 
 		while (worms.isNotEmpty()) {
@@ -147,15 +149,19 @@ class WormPlacement(
 				it.hasFinished
 			}
 		}
+
+		return results
 	}
 
-	override fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos> {
+	override fun evaluate(context: EnhancedPlacementContext, positions: List<PlacementPos>): List<PlacementPos> {
 		if (positions.isEmpty())
 			return positions
 		else {
-			val results = ArrayList<BlockPos>()
+			val result = ArrayList<PlacementPos>()
-			positions.forEach { evaluate(context.random, it, results) }
+			val results = positions.map { (it, c) -> evaluate(context.random, it).map { PlacementPos(it, c) } }
-			return results
+			result.ensureCapacity(result.size + results.sumOf { it.size } + 1)
+			results.forEach { result.addAll(it) }
+			return result
 		}
 	}
 
@@ -163,9 +169,7 @@ class WormPlacement(
 		get() = Companion
 
 	override fun getPositions(context: PlacementContext, random: RandomSource, center: BlockPos): Stream<BlockPos> {
-		val results = ArrayList<BlockPos>()
+		return evaluate(random, center).stream()
-		evaluate(random, center, results)
-		return results.stream()
 	}
 
 	override fun type(): PlacementModifierType<*> {