1 files changed, 31 insertions, 4 deletions

diff --git a/lib/Fuzzer/test/FuzzerUnittest.cpp b/lib/Fuzzer/test/FuzzerUnittest.cpp
index 8c00127..b33e0c9 100644
--- a/lib/Fuzzer/test/FuzzerUnittest.cpp
+++ b/lib/Fuzzer/test/FuzzerUnittest.cpp
@@ -247,8 +247,8 @@ void TestAddWordFromDictionary(Mutator M, int NumIter) {
   MutationDispatcher MD(Rand);
   uint8_t Word1[4] = {0xAA, 0xBB, 0xCC, 0xDD};
   uint8_t Word2[3] = {0xFF, 0xEE, 0xEF};
-  MD.AddWordToDictionary(Word1, sizeof(Word1));
-  MD.AddWordToDictionary(Word2, sizeof(Word2));
+  MD.AddWordToManualDictionary(Unit(Word1, Word1 + sizeof(Word1)));
+  MD.AddWordToManualDictionary(Unit(Word2, Word2 + sizeof(Word2)));
   int FoundMask = 0;
   uint8_t CH0[7] = {0x00, 0x11, 0x22, 0xAA, 0xBB, 0xCC, 0xDD};
   uint8_t CH1[7] = {0x00, 0x11, 0xAA, 0xBB, 0xCC, 0xDD, 0x22};
@@ -274,14 +274,41 @@ void TestAddWordFromDictionary(Mutator M, int NumIter) {
 }
 
 TEST(FuzzerMutate, AddWordFromDictionary1) {
-  TestAddWordFromDictionary(&MutationDispatcher::Mutate_AddWordFromDictionary,
-                            1 << 15);
+  TestAddWordFromDictionary(
+      &MutationDispatcher::Mutate_AddWordFromManualDictionary, 1 << 15);
 }
 
 TEST(FuzzerMutate, AddWordFromDictionary2) {
   TestAddWordFromDictionary(&MutationDispatcher::Mutate, 1 << 15);
 }
 
+void TestAddWordFromDictionaryWithHint(Mutator M, int NumIter) {
+  FuzzerRandomLibc Rand(0);
+  MutationDispatcher MD(Rand);
+  uint8_t Word[] = {0xAA, 0xBB, 0xCC, 0xDD, 0xFF, 0xEE, 0xEF};
+  size_t PosHint = 7777;
+  MD.AddWordToAutoDictionary(Unit(Word, Word + sizeof(Word)), PosHint);
+  int FoundMask = 0;
+  for (int i = 0; i < NumIter; i++) {
+    uint8_t T[10000];
+    memset(T, 0, sizeof(T));
+    size_t NewSize = (MD.*M)(T, 9000, 10000);
+    if (NewSize >= PosHint + sizeof(Word) &&
+        !memcmp(Word, T + PosHint, sizeof(Word)))
+      FoundMask = 1;
+  }
+  EXPECT_EQ(FoundMask, 1);
+}
+
+TEST(FuzzerMutate, AddWordFromDictionaryWithHint1) {
+  TestAddWordFromDictionaryWithHint(
+      &MutationDispatcher::Mutate_AddWordFromAutoDictionary, 1 << 5);
+}
+
+TEST(FuzzerMutate, AddWordFromDictionaryWithHint2) {
+  TestAddWordFromDictionaryWithHint(&MutationDispatcher::Mutate, 1 << 10);
+}
+
 void TestChangeASCIIInteger(Mutator M, int NumIter) {
   FuzzerRandomLibc Rand(0);
   MutationDispatcher MD(Rand);