Allow repeated keys in RandomCollection (#2624)

Author: SirYwell

worldedit-core/src/main/java/com/fastasyncworldedit/core/extent/transform/RandomTransform.java

@@ -7,10 +7,10 @@
 import com.sk89q.worldedit.extent.AbstractDelegateExtent;
 import com.sk89q.worldedit.extent.Extent;
 
-import java.util.HashMap;
-import java.util.LinkedHashSet;
-import java.util.Map;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
@@ -20,10 +20,9 @@
 public class RandomTransform extends SelectTransform {
 
     private final SimpleRandom random;
-    private final Map<ResettableExtent, Double> weights = new HashMap<>();
+    private final List<RandomCollection.Weighted<ResettableExtent>> weights;
 
     private transient RandomCollection<ResettableExtent> collection;
-    private transient LinkedHashSet<ResettableExtent> extents = new LinkedHashSet<>();
 
     public RandomTransform() {
         this(new TrueRandom());
@@ -36,27 +35,27 @@ public RandomTransform() {
      */
     public RandomTransform(SimpleRandom random) {
         this.random = random;
+        this.weights = new ArrayList<>();
     }
 
     @Override
     public AbstractDelegateExtent getExtent(int x, int y, int z) {
-        return collection.next(x, y, z);
+        return collection.next(this.random, x, y, z);
     }
 
     @Override
     public AbstractDelegateExtent getExtent(int x, int z) {
-        return collection.next(x, 0, z);
+        return collection.next(this.random, x, 0, z);
     }
 
     @Override
     public ResettableExtent setExtent(Extent extent) {
         if (collection == null) {
-            collection = RandomCollection.of(weights, random);
-            extents = new LinkedHashSet<>(weights.keySet());
+            collection = RandomCollection.of(weights);
         }
         super.setExtent(extent);
-        for (ResettableExtent current : extents) {
-            current.setExtent(extent);
+        for (RandomCollection.Weighted<ResettableExtent> current : this.weights) {
+            current.value().setExtent(extent);
         }
         return this;
     }
@@ -72,13 +71,12 @@ public ResettableExtent setExtent(Extent extent) {
      */
     public void add(ResettableExtent extent, double chance) {
         checkNotNull(extent);
-        weights.put(extent, chance);
-        collection = RandomCollection.of(weights, random);
-        this.extents.add(extent);
+        weights.add(new RandomCollection.Weighted<>(extent, chance));
+        collection = RandomCollection.of(weights);
     }
 
     public Set<ResettableExtent> getExtents() {
-        return extents;
+        return this.weights.stream().map(RandomCollection.Weighted::value).collect(Collectors.toSet());
     }
 
     public RandomCollection<ResettableExtent> getCollection() {

worldedit-core/src/main/java/com/fastasyncworldedit/core/util/collection/FastRandomCollection.java

@@ -4,15 +4,14 @@
 import com.fastasyncworldedit.core.util.MathMan;
 
 import java.util.ArrayList;
-import java.util.Map;
+import java.util.List;
 import java.util.Optional;
 
-public class FastRandomCollection<T> extends RandomCollection<T> {
+public final class FastRandomCollection<T> implements RandomCollection<T> {
 
     private final T[] values;
 
-    private FastRandomCollection(T[] values, SimpleRandom random) {
-        super(random);
+    private FastRandomCollection(T[] values) {
         this.values = values;
     }
 
@@ -22,16 +21,15 @@ private FastRandomCollection(T[] values, SimpleRandom random) {
      * {@code Optional} in any case.
      *
      * @param weights the weight of the values.
-     * @param random  the random generator to use for this collection.
      * @param <T>     the value type.
      * @return an {@link Optional} containing the new collection if it could be created, {@link
      *         Optional#empty()} otherwise.
      * @see RandomCollection for API usage.
      */
-    public static <T> Optional<RandomCollection<T>> create(Map<T, Double> weights, SimpleRandom random) {
+    public static <T> Optional<RandomCollection<T>> create(List<Weighted<T>> weights) {
         int max = 0;
         int[] counts = new int[weights.size()];
-        Double[] weightDoubles = weights.values().toArray(new Double[0]);
+        double[] weightDoubles = weights.stream().mapToDouble(Weighted::weight).toArray();
         for (int i = 0; i < weightDoubles.length; i++) {
             int weight = (int) (weightDoubles[i] * 100);
             counts[i] = weight;
@@ -47,21 +45,21 @@ public static <T> Optional<RandomCollection<T>> create(Map<T, Double> weights, S
             return Optional.empty();
         }
         ArrayList<T> parsed = new ArrayList<>();
-        for (Map.Entry<T, Double> entry : weights.entrySet()) {
-            int num = (int) (100 * entry.getValue());
+        for (Weighted<T> entry : weights) {
+            int num = (int) (100 * entry.weight());
             for (int j = 0; j < num / gcd; j++) {
-                parsed.add(entry.getKey());
+                parsed.add(entry.value());
             }
         }
         @SuppressWarnings("unchecked")
         T[] values = (T[]) parsed.toArray();
-        FastRandomCollection<T> fastRandomCollection = new FastRandomCollection<>(values, random);
+        FastRandomCollection<T> fastRandomCollection = new FastRandomCollection<>(values);
         return Optional.of(fastRandomCollection);
     }
 
     @Override
-    public T next(int x, int y, int z) {
-        return values[getRandom().nextInt(x, y, z, values.length)];
+    public T next(final SimpleRandom random, int x, int y, int z) {
+        return values[random.nextInt(x, y, z, values.length)];
     }
 
 }

worldedit-core/src/main/java/com/fastasyncworldedit/core/util/collection/RandomCollection.java

@@ -2,49 +2,35 @@
 
 import com.fastasyncworldedit.core.math.random.SimpleRandom;
 
-import java.util.Map;
+import java.util.List;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
 /**
  * A RandomCollection holds multiple values that can be accessed by using
- * {@link RandomCollection#next(int, int, int)}. The returned value is
+ * {@link RandomCollection#next(SimpleRandom, int, int, int)}. The returned value is
  * determined by a given {@link SimpleRandom} implementation.
  *
  * @param <T> the type of values the collection holds.
  */
-public abstract class RandomCollection<T> {
-
-    private SimpleRandom random;
-
-    protected RandomCollection(SimpleRandom random) {
-        this.random = random;
-    }
+public sealed interface RandomCollection<T> permits FastRandomCollection, SimpleRandomCollection {
 
     /**
      * Return a new RandomCollection. The implementation may differ depending on the
      * given arguments but there is no need to differ.
      *
      * @param weights the weighted map.
-     * @param random  the random number generator.
      * @param <T>     the type the collection holds.
      * @return a RandomCollection using the given weights and the RNG.
      */
-    public static <T> RandomCollection<T> of(Map<T, Double> weights, SimpleRandom random) {
-        checkNotNull(random);
-        return FastRandomCollection.create(weights, random)
-                .orElseGet(() -> new SimpleRandomCollection<>(weights, random));
+    static <T> RandomCollection<T> of(List<Weighted<T>> weights) {
+        return FastRandomCollection.create(weights)
+                .orElseGet(() -> new SimpleRandomCollection<>(weights));
     }
 
-    public void setRandom(SimpleRandom random) {
-        checkNotNull(random);
-        this.random = random;
-    }
+    T next(SimpleRandom random, int x, int y, int z);
 
-    public SimpleRandom getRandom() {
-        return random;
-    }
-
-    public abstract T next(int x, int y, int z);
+    record Weighted<T>(T value, double weight) {
 
+    }
 }

worldedit-core/src/main/java/com/fastasyncworldedit/core/util/collection/SimpleRandomCollection.java

@@ -2,41 +2,39 @@
 
 import com.fastasyncworldedit.core.math.random.SimpleRandom;
 
-import java.util.Map;
+import java.util.List;
 import java.util.NavigableMap;
 import java.util.TreeMap;
 
-public class SimpleRandomCollection<E> extends RandomCollection<E> {
+public final class SimpleRandomCollection<T> implements RandomCollection<T> {
 
-    private final NavigableMap<Double, E> map = new TreeMap<>();
-    private double total = 0;
+    private final NavigableMap<Double, T> map;
+    private final double total;
 
     /**
      * Create a {@link RandomCollection} from a weighted map and a RNG.
-     * It is recommended to use {@link RandomCollection#of(Map, SimpleRandom)}
+     * It is recommended to use {@link RandomCollection#of(List)}
      * instead of this constructor.
      *
      * @param weights the weighted map.
-     * @param random  the random number generator.
      */
-    public SimpleRandomCollection(Map<E, Double> weights, SimpleRandom random) {
-        super(random);
-        for (Map.Entry<E, Double> entry : weights.entrySet()) {
-            add(entry.getValue(), entry.getKey());
+    public SimpleRandomCollection(List<Weighted<T>> weights) {
+        this.map = new TreeMap<>();
+        double total = 0;
+        for (Weighted<T> entry : weights) {
+            final double weight = entry.weight();
+            if (weight <= 0) {
+                throw new IllegalArgumentException("Weights must be positive");
+            }
+            total += weight;
+            this.map.put(total, entry.value());
         }
-    }
-
-    public void add(double weight, E result) {
-        if (weight <= 0) {
-            return;
-        }
-        total += weight;
-        map.put(total, result);
+        this.total = total;
     }
 
     @Override
-    public E next(int x, int y, int z) {
-        return map.ceilingEntry(getRandom().nextDouble(x, y, z, this.total)).getValue();
+    public T next(final SimpleRandom random, int x, int y, int z) {
+        return map.ceilingEntry(random.nextDouble(x, y, z, this.total)).getValue();
     }
 
 }

worldedit-core/src/main/java/com/sk89q/worldedit/function/pattern/RandomPattern.java

@@ -27,10 +27,10 @@
 import com.sk89q.worldedit.math.BlockVector3;
 import com.sk89q.worldedit.world.block.BaseBlock;
 
-import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
-import java.util.Map;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
@@ -39,11 +39,10 @@
  */
 public class RandomPattern extends AbstractPattern {
 
-    //FAWE start - SimpleRandom > Random, LHS<P> > List
+    //FAWE start - SimpleRandom > Random, RandomCollection
     private final SimpleRandom random;
-    private Map<Pattern, Double> weights = new LinkedHashMap<>();
+    private final List<RandomCollection.Weighted<Pattern>> weights;
     private RandomCollection<Pattern> collection;
-    private LinkedHashSet<Pattern> patterns = new LinkedHashSet<>();
     //FAWE end
 
     //FAWE start
@@ -53,6 +52,7 @@ public RandomPattern() {
 
     public RandomPattern(SimpleRandom random) {
         this.random = random;
+        this.weights = new ArrayList<>();
     }
 
     /**
@@ -63,16 +63,10 @@ public RandomPattern(SimpleRandom random) {
      */
     public RandomPattern(SimpleRandom random, RandomPattern parent) {
         this.random = random;
-        this.weights = parent.weights;
-        this.collection = RandomCollection.of(weights, random);
-        this.patterns = parent.patterns;
-    }
-
-    private RandomPattern(SimpleRandom random, Map<Pattern, Double> weights) {
-        this.random = random;
-        this.weights = weights;
-        this.collection = RandomCollection.of(weights, random);
-        this.patterns = new LinkedHashSet<>(weights.keySet());
+        this.weights = parent.weights.stream()
+                .map(weighted -> new RandomCollection.Weighted<>(weighted.value().fork(), weighted.weight()))
+                .collect(Collectors.toCollection(ArrayList::new));
+        this.collection = RandomCollection.of(weights);
     }
     //FAWE end
 
@@ -87,18 +81,15 @@ private RandomPattern(SimpleRandom random, Map<Pattern, Double> weights) {
      */
     public void add(Pattern pattern, double chance) {
         checkNotNull(pattern);
-        //FAWE start - Double, weights, patterns and collection
-        Double existingWeight = weights.get(pattern);
-        if (existingWeight != null) {
-            chance += existingWeight;
-        }
-        weights.put(pattern, chance);
-        collection = RandomCollection.of(weights, random);
-        this.patterns.add(pattern);
+        //FAWE start - Double, weights, repeating patterns, and collection
+        this.weights.add(new RandomCollection.Weighted<>(pattern, chance));
+        this.collection = RandomCollection.of(weights);
     }
 
     public Set<Pattern> getPatterns() {
-        return patterns;
+        return this.weights.stream()
+                .map(RandomCollection.Weighted::value)
+                .collect(Collectors.toSet());
     }
 
     public RandomCollection<Pattern> getCollection() {
@@ -107,19 +98,17 @@ public RandomCollection<Pattern> getCollection() {
 
     @Override
     public BaseBlock applyBlock(BlockVector3 position) {
-        return collection.next(position.x(), position.y(), position.z()).applyBlock(position);
+        return collection.next(this.random, position.x(), position.y(), position.z()).applyBlock(position);
     }
 
     @Override
     public boolean apply(Extent extent, BlockVector3 get, BlockVector3 set) throws WorldEditException {
-        return collection.next(get.x(), get.y(), get.z()).apply(extent, get, set);
+        return collection.next(this.random, get.x(), get.y(), get.z()).apply(extent, get, set);
     }
 
     @Override
     public Pattern fork() {
-        final LinkedHashMap<Pattern, Double> newWeights = new LinkedHashMap<>();
-        this.weights.forEach((p, w) -> newWeights.put(p.fork(), w));
-        return new RandomPattern(this.random, newWeights);
+        return new RandomPattern(this.random, this);
     }
 
     //FAWE end