feat: Integral mutators now use @DictionaryProvider

Author: oetr

src/main/java/com/code_intelligence/jazzer/mutation/mutator/lang/IntegralMutatorFactory.java

@@ -24,8 +24,11 @@
 import com.code_intelligence.jazzer.mutation.api.MutatorFactory;
 import com.code_intelligence.jazzer.mutation.api.PseudoRandom;
 import com.code_intelligence.jazzer.mutation.api.SerializingMutator;
+import com.code_intelligence.jazzer.mutation.combinator.SamplingUtils;
+import com.code_intelligence.jazzer.mutation.combinator.SamplingUtils.WeightedValue;
 import com.code_intelligence.jazzer.mutation.mutator.libfuzzer.LibFuzzerMutate;
 import com.code_intelligence.jazzer.mutation.runtime.MutatorRuntime;
+import com.code_intelligence.jazzer.mutation.support.DictionaryProviderSupport;
 import com.code_intelligence.jazzer.mutation.support.RangeSupport;
 import com.code_intelligence.jazzer.mutation.support.RangeSupport.LongRange;
 import com.google.errorprone.annotations.ForOverride;
@@ -34,7 +37,10 @@
 import java.io.IOException;
 import java.lang.reflect.AnnotatedType;
 import java.lang.reflect.ParameterizedType;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.LongStream;
 
@@ -49,7 +55,7 @@ public Optional<SerializingMutator<?>> tryCreate(
 
     if (clazz == byte.class || clazz == Byte.class) {
       return Optional.of(
-          new AbstractIntegralMutator<Byte>(type, Byte.MIN_VALUE, Byte.MAX_VALUE) {
+          new AbstractIntegralMutator<Byte>(runtime, type, Byte.MIN_VALUE, Byte.MAX_VALUE) {
             @Override
             protected long mutateWithLibFuzzer(long value) {
               return LibFuzzerMutate.mutateDefault((byte) value, this, 0);
@@ -82,7 +88,7 @@ public void write(Byte value, DataOutputStream out) throws IOException {
           });
     } else if (clazz == short.class || clazz == Short.class) {
       return Optional.of(
-          new AbstractIntegralMutator<Short>(type, Short.MIN_VALUE, Short.MAX_VALUE) {
+          new AbstractIntegralMutator<Short>(runtime, type, Short.MIN_VALUE, Short.MAX_VALUE) {
             @Override
             protected long mutateWithLibFuzzer(long value) {
               return LibFuzzerMutate.mutateDefault((short) value, this, 0);
@@ -115,7 +121,8 @@ public void write(Short value, DataOutputStream out) throws IOException {
           });
     } else if (clazz == int.class || clazz == Integer.class) {
       return Optional.of(
-          new AbstractIntegralMutator<Integer>(type, Integer.MIN_VALUE, Integer.MAX_VALUE) {
+          new AbstractIntegralMutator<Integer>(
+              runtime, type, Integer.MIN_VALUE, Integer.MAX_VALUE) {
             @Override
             protected long mutateWithLibFuzzer(long value) {
               return LibFuzzerMutate.mutateDefault((int) value, this, 0);
@@ -148,7 +155,7 @@ public void write(Integer value, DataOutputStream out) throws IOException {
           });
     } else if (clazz == long.class || clazz == Long.class) {
       return Optional.of(
-          new AbstractIntegralMutator<Long>(type, Long.MIN_VALUE, Long.MAX_VALUE) {
+          new AbstractIntegralMutator<Long>(runtime, type, Long.MIN_VALUE, Long.MAX_VALUE) {
             @Override
             protected long mutateWithLibFuzzer(long value) {
               return LibFuzzerMutate.mutateDefault(value, this, 0);
@@ -197,9 +204,14 @@ abstract static class AbstractIntegralMutator<T extends Number> extends Serializ
     private final int largestMutableBitNegative;
     private final int largestMutableBitPositive;
     private final long[] specialValues;
+    private final long[] dictionaryValues;
+    private final Function<PseudoRandom, MutationFunction> mutationFunctionSampler;
 
     AbstractIntegralMutator(
-        AnnotatedType type, long defaultMinValueForType, long defaultMaxValueForType) {
+        MutatorRuntime runtime,
+        AnnotatedType type,
+        long defaultMinValueForType,
+        long defaultMaxValueForType) {
       LongRange resolved =
           RangeSupport.resolveIntegralRange(type, defaultMinValueForType, defaultMaxValueForType);
       long minValue = resolved.min;
@@ -232,6 +244,50 @@ abstract static class AbstractIntegralMutator<T extends Number> extends Serializ
         largestMutableBitPositive = bitWidth(maxValue);
       }
       this.specialValues = collectSpecialValues(minValue, maxValue);
+
+      this.dictionaryValues =
+          DictionaryProviderSupport.extractRawValues(runtime, type)
+              .map(
+                  stream ->
+                      stream
+                          .filter(v -> v instanceof Number)
+                          .map(v -> ((Number) v).longValue())
+                          .filter(v -> v >= minValue)
+                          .filter(v -> v <= maxValue)
+                          .sorted()
+                          .mapToLong(Long::longValue)
+                          .toArray())
+              .orElse(null);
+      List<WeightedValue<MutationFunction>> f = new ArrayList<>();
+      f.add(new WeightedValue<>(1.0, MutationFunction.BIT_FLIP));
+      f.add(new WeightedValue<>(1.0, MutationFunction.RANDOM_WALK));
+      f.add(new WeightedValue<>(1.0, MutationFunction.RANDOM_VALUE));
+      f.add(new WeightedValue<>(1.0, MutationFunction.LIBFUZZER));
+      if (dictionaryValues != null && dictionaryValues.length > 0) {
+        // Since weights here are relative, we need to adjust the weight of user dictionary mutator
+        // so that it is taken proportionate the inverse probability specified in the annotation.
+        // Basically, we need to scale up the weight for pInv:
+        // 1/p    --- x?
+        // 1- 1/p --- totalFuncWeights
+        // x = (1/p * totalFuncWeights) / (1 - 1/p)
+        //   =  totalFuncWeights / (p - 1)
+        double totalFuncWeights = 0.0;
+        for (WeightedValue<MutationFunction> wf : f) {
+          totalFuncWeights += wf.weight;
+        }
+        int invProbability = DictionaryProviderSupport.extractFirstInvProbability(type);
+        double perValueWeight = totalFuncWeights / (invProbability - 1);
+        f.add(new WeightedValue<>(perValueWeight, MutationFunction.DICTIONARY_VALUE));
+      }
+      this.mutationFunctionSampler = SamplingUtils.weightedSampler(f);
+    }
+
+    private enum MutationFunction {
+      BIT_FLIP,
+      DICTIONARY_VALUE,
+      LIBFUZZER,
+      RANDOM_VALUE,
+      RANDOM_WALK,
     }
 
     private static long[] collectSpecialValues(long minValue, long maxValue) {
@@ -263,20 +319,25 @@ protected final long mutateImpl(long value, PseudoRandom prng) {
       final long previousValue = value;
       // Mutate in a loop to verify that we really mutated.
       do {
-        switch (prng.indexIn(4)) {
-          case 0:
+        switch (mutationFunctionSampler.apply(prng)) {
+          case BIT_FLIP:
             value = bitFlip(value, prng);
             break;
-          case 1:
+          case RANDOM_WALK:
             value = randomWalk(value, prng);
             break;
-          case 2:
+          case RANDOM_VALUE:
             value = prng.closedRange(minValue, maxValue);
             break;
-          case 3:
+          case LIBFUZZER:
             // TODO: Replace this with a structure-aware dictionary/TORC search similar to fuzztest.
             value = forceInRange(mutateWithLibFuzzer(value));
             break;
+          case DICTIONARY_VALUE:
+            value = dictionaryValues[prng.indexIn(dictionaryValues.length)];
+            break;
+          default:
+            throw new AssertionError("Invalid mutation function.");
         }
       } while (value == previousValue);
       return value;