More efficient ellipsoid placement

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

@@ -2,6 +2,7 @@ package ru.dbotthepony.mc.otm.worldgen.placement
 
 import com.mojang.serialization.MapCodec
 import com.mojang.serialization.codecs.RecordCodecBuilder
+import it.unimi.dsi.fastutil.objects.ObjectRBTreeSet
 import net.minecraft.core.BlockPos
 import net.minecraft.util.RandomSource
 import net.minecraft.util.valueproviders.FloatProvider
@@ -9,16 +10,10 @@ import net.minecraft.util.valueproviders.IntProvider
 import net.minecraft.world.level.levelgen.placement.PlacementContext
 import net.minecraft.world.level.levelgen.placement.PlacementModifier
 import net.minecraft.world.level.levelgen.placement.PlacementModifierType
-import ru.dbotthepony.mc.otm.core.math.component1
+import ru.dbotthepony.mc.otm.core.collect.Vec3iHashStrategy
-import ru.dbotthepony.mc.otm.core.math.component2
-import ru.dbotthepony.mc.otm.core.math.component3
-import ru.dbotthepony.mc.otm.core.math.minus
-import ru.dbotthepony.mc.otm.core.math.plus
 import ru.dbotthepony.mc.otm.registry.data.MPlacementModifiers
 import ru.dbotthepony.mc.otm.worldgen.EnhancedPlacementContext
 import java.util.stream.Stream
-import kotlin.math.PI
-import kotlin.math.roundToInt
 
 /**
  * Ellipsoid ("cloud") placement
@@ -65,11 +60,12 @@ data class EllipsoidPlacement(
 		require(yLength.minValue >= 1f) { "Bad ellipsoid y minimal size: $yLength" }
 	}
 
-	private fun evaluate(random: RandomSource, position: BlockPos, result: ArrayList<BlockPos>) {
+	private fun evaluate(random: RandomSource, position: BlockPos): ObjectRBTreeSet<BlockPos> {
 		val count = count.sample(random)
+		val results = ObjectRBTreeSet<BlockPos>(Vec3iHashStrategy)
 
 		if (count <= 0)
-			return
+			return results
 
 		val xLength = xLength.sample(random)
 		val zLength = zLength.sample(random)
@@ -79,26 +75,24 @@ data class EllipsoidPlacement(
 		val zPow = zLength * zLength
 		val yPow = yLength * yLength
 
-		result.ensureCapacity(result.size + count + 1)
+		var iterations = 0
+		val maxIterations = count * 10
 
-		Stream.generate {
+		while (results.size < count && ++iterations < maxIterations) {
 			val x = this.x.sample(random) * xLength
 			val y = this.y.sample(random) * yLength
 			val z = this.z.sample(random) * zLength
-			BlockPos(x.toInt(), y.toInt(), z.toInt())
-		}
-			.limit(count.toLong() * 10)
-			.filter {
-				val (ellipsoidX, ellipsoidY, ellipsoidZ) = it
 
-				(ellipsoidX * ellipsoidX) / xPow +
+			val isValidPoint = (x * x) / xPow +
-					(ellipsoidY * ellipsoidY) / yPow +
+				(y * y) / yPow +
-					(ellipsoidZ * ellipsoidZ) / zPow <= 1.0f
+				(z * z) / zPow <= 1.0f
+
+			if (isValidPoint) {
+				results.add(BlockPos(x.toInt() + position.x, y.toInt() + position.y, z.toInt() + position.z))
 			}
-			.distinct()
+		}
-			.limit(count.toLong())
+
-			.map { it + position }
+		return results
-			.forEach { result.add(it) }
 	}
 
 	override fun getPositions(
@@ -106,14 +100,14 @@ data class EllipsoidPlacement(
 		random: RandomSource,
 		position: BlockPos
 	): Stream<BlockPos> {
-		val result = ArrayList<BlockPos>()
+		return evaluate(random, position).stream()
-		evaluate(random, position, result)
-		return result.stream()
 	}
 
 	override fun evaluate(context: EnhancedPlacementContext, positions: List<BlockPos>): List<BlockPos> {
 		val result = ArrayList<BlockPos>()
-		positions.forEach { evaluate(context.random, it, result) }
+		val results = positions.map { evaluate(context.random, it) }
+		result.ensureCapacity(result.size + results.sumOf { it.size } + 1)
+		results.forEach { result.addAll(it) }
 		return result
 	}