Fix lootable ranges for any

Signed-off-by: Pablo Herrera <[email protected]>

core/src/main/java/tc/oc/pgm/util/compose/Any.java

@@ -2,8 +2,10 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Range;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Random;
+import java.util.function.Function;
 import java.util.stream.Stream;
 import tc.oc.pgm.api.filter.Filter;
 import tc.oc.pgm.api.filter.query.Query;
@@ -12,6 +14,8 @@
 
 public class Any<T> implements Composition<T> {
 
+  private static final Random RANDOM = new Random();
+
   public static class Option<T> {
     private final double weight;
     private final Filter filter;
@@ -24,19 +28,13 @@ public Option(double weight, Filter filter, Composition<T> element) {
     }
   }
 
-  private int countMaxAmount;
-  private final int countOffset;
+  private final Range<Integer> range;
   private final boolean unique;
   private final List<Option<T>> options;
   private final double totalWeight;
 
   public Any(Range<Integer> count, boolean unique, Iterable<Option<T>> choices) {
-    // Range is enforced to be bounded in constructor
-    count = Ranges.toClosed(count);
-    this.countOffset = count.lowerEndpoint();
-    // If the range is a singleton we need to manually set this to 1 since
-    // lower and upper endpoint will be identical
-    this.countMaxAmount = Math.max(count.upperEndpoint() - count.lowerEndpoint(), 1);
+    this.range = Ranges.toClosed(count);
     this.unique = unique;
     this.options = ImmutableList.copyOf(choices);
     this.totalWeight = this.options.stream().mapToDouble(c -> c.weight).sum();
@@ -53,16 +51,18 @@ public Stream<T> elements(Query query) {
       }
     }
 
-    final Random random = new Random();
-    Stream<T> result = Stream.empty();
+    int picks = random(range.lowerEndpoint(), range.upperEndpoint());
+    List<Stream<T>> result = new ArrayList<>(picks);
 
-    for (int count = random.nextInt(this.countMaxAmount + 1) + this.countOffset;
-        count > 0 && !chooser.isEmpty();
-        count--) {
-      final Option<T> option = chooser.choose(random);
-      result = Stream.concat(result, option.element.elements(query));
+    for (int i = picks; i > 0 && !chooser.isEmpty(); i--) {
+      final Option<T> option = chooser.choose(RANDOM);
+      result.add(option.element.elements(query));
       if (unique) chooser.remove(option);
     }
-    return result;
+    return result.stream().flatMap(Function.identity());
+  }
+
+  private int random(int min, int max) {
+    return min >= max ? min : min + RANDOM.nextInt(max + 1 - min);
   }
 }