Merge pull request pholser#248 from pholser/issues/241

For pholser#241, optimize SourceOfRandomness.choose()
vlsi · Oct 8, 2019 · e60611d · e60611d
2 parents 9f4b690 + 8b6e436
commit e60611d
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 12 deletions.
diff --git a/core/src/main/java/com/pholser/junit/quickcheck/internal/Items.java b/core/src/main/java/com/pholser/junit/quickcheck/internal/Items.java
@@ -26,6 +26,8 @@ a copy of this software and associated documentation files (the
 package com.pholser.junit.quickcheck.internal;
 
 import java.util.Collection;
+import java.util.List;
+import java.util.RandomAccess;
 
 import com.pholser.junit.quickcheck.random.SourceOfRandomness;
 
@@ -36,11 +38,31 @@ private Items() {
 
     @SuppressWarnings("unchecked")
     public static <T> T choose(Collection<T> items, SourceOfRandomness random) {
-        Object[] asArray = items.toArray(new Object[items.size()]);
-        return (T) asArray[random.nextInt(items.size())];
+        int size = items.size();
+        if (size == 0) {
+            throw new IllegalArgumentException(
+                "Collection is empty, can't pick an element from it");
+        }
+
+        if (items instanceof RandomAccess && items instanceof List<?>) {
+            List<T> list = (List<T>) items;
+            return size == 1
+                ? list.get(0)
+                : list.get(random.nextInt(size));
+        }
+
+        if (size == 1) {
+            return items.iterator().next();
+        }
+
+        Object[] array = items.toArray(new Object[0]);
+        return (T) array[random.nextInt(array.length)];
     }
 
-    public static <T> T chooseWeighted(Collection<Weighted<T>> items, SourceOfRandomness random) {
+    public static <T> T chooseWeighted(
+        Collection<Weighted<T>> items,
+        SourceOfRandomness random) {
+
         if (items.size() == 1)
             return items.iterator().next().item;
 
@@ -54,6 +76,7 @@ public static <T> T chooseWeighted(Collection<Weighted<T>> items, SourceOfRandom
                 return each.item;
         }
 
-        throw new AssertionError(String.format("sample = %d, range = %d", sample, range));
+        throw new AssertionError(
+            String.format("sample = %d, range = %d", sample, range));
     }
 }
diff --git a/core/src/main/java/com/pholser/junit/quickcheck/random/SourceOfRandomness.java b/core/src/main/java/com/pholser/junit/quickcheck/random/SourceOfRandomness.java
@@ -36,6 +36,7 @@ a copy of this software and associated documentation files (the
 import java.util.Collection;
 import java.util.Random;
 
+import com.pholser.junit.quickcheck.internal.Items;
 import com.pholser.junit.quickcheck.internal.Ranges;
 
 import static java.util.concurrent.TimeUnit.*;
@@ -352,10 +353,8 @@ public Duration nextDuration(Duration min, Duration max) {
      * @param items a collection
      * @return a randomly chosen element from the collection
      */
-    @SuppressWarnings("unchecked")
     public <T> T choose(Collection<T> items) {
-        Object[] array = items.toArray(new Object[items.size()]);
-        return (T) array[nextInt(array.length)];
+        return Items.choose(items, this);
     }
 
     /**

diff --git a/core/src/test/java/com/pholser/junit/quickcheck/internal/ItemsTest.java b/core/src/test/java/com/pholser/junit/quickcheck/internal/ItemsTest.java
@@ -32,6 +32,7 @@ a copy of this software and associated documentation files (the
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
@@ -43,6 +44,7 @@ a copy of this software and associated documentation files (the
 
 public class ItemsTest {
     @Rule public final MockitoRule mockito = MockitoJUnit.rule();
+    @Rule public final ExpectedException thrown = ExpectedException.none();
 
     @Mock private SourceOfRandomness random;
     private Weighted<String> first;
@@ -93,10 +95,15 @@ public class ItemsTest {
     }
 
     @Test public void choosingFromEmptyCollection() {
-        try {
-            Items.chooseWeighted(emptyList(), random);
-        } catch (AssertionError expected) {
-            assertEquals("sample = 0, range = 0", expected.getMessage());
-        }
+        thrown.expect(IllegalArgumentException.class);
+
+        Items.choose(emptyList(), random);
+    }
+
+    @Test public void choosingWeightedFromEmptyCollection() {
+        thrown.expect(AssertionError.class);
+        thrown.expectMessage("sample = 0, range = 0");
+
+        Items.chooseWeighted(emptyList(), random);
     }
 }