Enable random NotNestedSet creation to accept a custom element genera…

…tor.

PiperOrigin-RevId: 632239739
Change-Id: I0079746eb826684ecc3209e8bc02e3445dd8ea0d

src/test/java/com/google/devtools/build/lib/skyframe/serialization/NotNestedSet.java

@@ -23,6 +23,7 @@
 import java.util.Random;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
+import java.util.function.Function;
 
 /** A nested-set like class for codec testing. */
 final class NotNestedSet {
@@ -43,11 +44,11 @@ Object[] getContents() {
   private static final int MAX_RANDOM_ELEMENTS = 5;
 
   /** Helper for constructing contents or nested contents. */
-  static Object[] createRandomLeafArray(Random rng) {
+  static Object[] createRandomLeafArray(Random rng, Function<Random, Object> elementFactory) {
     int count = rng.nextInt(MAX_RANDOM_ELEMENTS - 1) + 1;
     Object[] array = new Object[count];
     for (int i = 0; i < count; i++) {
-      array[i] = rng.nextInt();
+      array[i] = elementFactory.apply(rng);
     }
     return array;
   }
@@ -61,7 +62,8 @@ static Object[] createRandomLeafArray(Random rng) {
    */
   private static final double EXTRA_PARENT_PROBABILITY = 0.01;
 
-  static NotNestedSet createRandom(Random rng, int maxLayers, int maxNodesPerLayer) {
+  static NotNestedSet createRandom(
+      Random rng, int maxLayers, int maxNodesPerLayer, Function<Random, Object> elementFactory) {
     // Creates a random DAG layer by layer.
     int layerCount = rng.nextInt(maxLayers - 1) + 1;
 
@@ -103,17 +105,17 @@ static NotNestedSet createRandom(Random rng, int maxLayers, int maxNodesPerLayer
       ArrayList<NodeBuilder> layer = layers.get(i);
       for (NodeBuilder builder : layer) {
         if (i == layerCount - 1) {
-          builder.value = createRandomLeafArray(rng);
+          builder.value = createRandomLeafArray(rng, elementFactory);
           continue;
         }
-        // Inserts additional random integer elements into each non-leaf node.
+        // Inserts additional random elements into each non-leaf node.
         int randomElementCount = rng.nextInt(MAX_RANDOM_ELEMENTS);
         ArrayList<Object> values = new ArrayList<>(builder.children.size() + randomElementCount);
         for (Coordinate child : builder.children) {
           values.add(layers.get(child.layer()).get(child.index()).value);
         }
         for (int j = 0; j < randomElementCount; j++) {
-          values.add(rng.nextInt());
+          values.add(elementFactory.apply(rng));
         }
         Collections.shuffle(values, rng);
         builder.value = values.toArray(new Object[0]);

src/test/java/com/google/devtools/build/lib/skyframe/serialization/SharedValueDeserializationContextTest.java

@@ -60,7 +60,7 @@ public final class SharedValueDeserializationContextTest {
   @TestParameters("{size: 64, useDeferredCodec: false}")
   @TestParameters("{size: 128, useDeferredCodec: false}")
   public void codec_roundTrips(int size, boolean useDeferredCodec) throws Exception {
-    new SerializationTester(NotNestedSet.createRandom(rng, size, size))
+    new SerializationTester(NotNestedSet.createRandom(rng, size, size, Random::nextInt))
         .addCodec(
             useDeferredCodec
                 ? new NotNestedSetDeferredCodec(new NestedArrayCodec())
@@ -125,7 +125,9 @@ public void getsShouldBeConcurrent() throws Exception {
     NotNestedSet subject =
         new NotNestedSet(
             new Object[] {
-              createRandomLeafArray(rng), createRandomLeafArray(rng), createRandomLeafArray(rng)
+              createRandomLeafArray(rng, Random::nextInt),
+              createRandomLeafArray(rng, Random::nextInt),
+              createRandomLeafArray(rng, Random::nextInt)
             });
 
     SerializationResult<ByteString> serialized =
@@ -284,9 +286,10 @@ public void valueDependsOnFuture(
     if (doesSecondAliasFirst) {
       subject =
           new NotNestedSetContainer(
-              NotNestedSet.createRandom(rng, 4, 4), NotNestedSet.createRandom(rng, 4, 4));
+              NotNestedSet.createRandom(rng, 4, 4, Random::nextInt),
+              NotNestedSet.createRandom(rng, 4, 4, Random::nextInt));
     } else {
-      NotNestedSet contained = NotNestedSet.createRandom(rng, 5, 5);
+      NotNestedSet contained = NotNestedSet.createRandom(rng, 5, 5, Random::nextInt);
       subject = new NotNestedSetContainer(contained, contained);
     }
     new SerializationTester(subject)

src/test/java/com/google/devtools/build/lib/skyframe/serialization/SharedValueSerializationContextTest.java

@@ -308,7 +308,7 @@ private ListenableFuture<SerializationResult<ByteString>> serializeWithExecutor(
   }
 
   private Object[] createRandomLeafArray() {
-    return NotNestedSet.createRandomLeafArray(rng);
+    return NotNestedSet.createRandomLeafArray(rng, Random::nextInt);
   }
 
   private static final long POLL_MS = 100;