add the functionality to pick multiple values from a collection or array to the FuzzedDataProvider

2 files changed, 47 insertions, 6 deletions

diff --git a/agent/src/main/java/com/code_intelligence/jazzer/api/FuzzedDataProvider.java b/agent/src/main/java/com/code_intelligence/jazzer/api/FuzzedDataProvider.java
index 3790a67b..30ae526a 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/api/FuzzedDataProvider.java
+++ b/agent/src/main/java/com/code_intelligence/jazzer/api/FuzzedDataProvider.java
@@ -14,6 +14,8 @@
 
 package com.code_intelligence.jazzer.api;
 
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Random;
@@ -284,6 +286,7 @@ public interface FuzzedDataProvider {
    * @param <T> the type of a collection element
    * @return an element from {@code collection} chosen based on the fuzzer input
    */
+  @SuppressWarnings("unchecked")
   default<T> T pickValue(Collection<T> collection) {
     int size = collection.size();
     if (size == 0) {
@@ -395,4 +398,47 @@ public interface FuzzedDataProvider {
   default char pickValue(char[] array) {
     return array[consumeInt(0, array.length - 1)];
   }
+
+  /**
+   * Picks {@code numOfElements} elements from {@code collection} based on the fuzzer input.
+   * <p><b>Note:</b> The distribution of picks is not perfectly uniform.
+   *
+   * @param collection the {@link Collection} to pick an element from.
+   * @param numOfElements the number of elements to pick.
+   * @param <T> the type of the collection element
+   * @return an array of size {@code numOfElements} from {@code collection} chosen based on the
+   *     fuzzer input
+   */
+  default<T> List<T> pickValues(Collection<T> collection, int numOfElements) {
+    int size = collection.size();
+    if (size == 0) {
+      throw new IllegalArgumentException("collection is empty");
+    }
+    if (numOfElements > collection.size()) {
+      throw new IllegalArgumentException("numOfElements exceeds collection.size()");
+    }
+
+    List<T> remainingElements = new ArrayList<>(collection);
+    List<T> pickedElements = new ArrayList<>();
+    for (int i = 0; i < numOfElements; i++) {
+      T element = pickValue(remainingElements);
+      pickedElements.add(element);
+      remainingElements.remove(element);
+    }
+    return pickedElements;
+  }
+
+  /**
+   * Picks {@code numOfElements} elements from {@code array} based on the fuzzer input.
+   * <p><b>Note:</b> The distribution of picks is not perfectly uniform.
+   *
+   * @param array the array to pick an element from.
+   * @param numOfElements the number of elements to pick.
+   * @param <T> the type of the array element
+   * @return an array of size {@code numOfElements} from {@code array} chosen based on the fuzzer
+   *     input
+   */
+  default<T> List<T> pickValues(T[] array, int numOfElements) {
+    return pickValues(Arrays.asList(array), numOfElements);
+  }
 }
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/autofuzz/Meta.java b/agent/src/main/java/com/code_intelligence/jazzer/autofuzz/Meta.java
index f544192c..0d0c3190 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/autofuzz/Meta.java
+++ b/agent/src/main/java/com/code_intelligence/jazzer/autofuzz/Meta.java
@@ -156,12 +156,7 @@ public class Meta {
                                                        .collect(Collectors.toList());
 
       int pickedMethodsNumber = data.consumeInt(0, cascadingBuilderMethods.size());
-      List<Method> pickedMethods = new ArrayList<>();
-      for (int i = 0; i < pickedMethodsNumber; i++) {
-        Method method = data.pickValue(cascadingBuilderMethods);
-        pickedMethods.add(method);
-        cascadingBuilderMethods.remove(method);
-      }
+      List<Method> pickedMethods = data.pickValues(cascadingBuilderMethods, pickedMethodsNumber);
 
       Method builderMethod = data.pickValue(originalObjectCreationMethods);