1 files changed, 77 insertions, 0 deletions

diff --git a/src/processor/range_map_unittest.cc b/src/processor/range_map_unittest.cc
index cff5b0a9..b9657f64 100644
--- a/src/processor/range_map_unittest.cc
+++ b/src/processor/range_map_unittest.cc
@@ -243,6 +243,71 @@ static bool RetrieveTest(TestMap *range_map, const RangeTest *range_test) {
 }
 
 
+// Test RetrieveRangeAtIndex, which is supposed to return objects in order
+// according to their addresses.  This test is performed by looping through
+// the map, calling RetrieveRangeAtIndex for all possible indices in sequence,
+// and verifying that each call returns a different object than the previous
+// call, and that ranges are returned with increasing base addresses.  Returns
+// false if the test fails.
+static bool RetrieveIndexTest(TestMap *range_map, int set) {
+  linked_ptr<CountedObject> object;
+  CountedObject *last_object = NULL;
+  AddressType last_base = 0;
+
+  int object_count = range_map->GetCount();
+  for (int object_index = 0; object_index < object_count; ++object_index) {
+    AddressType base;
+    if (!range_map->RetrieveRangeAtIndex(object_index, &object, &base, NULL)) {
+      fprintf(stderr, "FAILED: RetrieveRangeAtIndex set %d index %d, "
+              "expected success, observed failure\n",
+              set, object_index);
+      return false;
+    }
+
+    if (!object.get()) {
+      fprintf(stderr, "FAILED: RetrieveRangeAtIndex set %d index %d, "
+              "expected object, observed NULL\n",
+              set, object_index);
+      return false;
+    }
+
+    // It's impossible to do these comparisons unless there's a previous
+    // object to compare against.
+    if (last_object) {
+      // The object must be different from the last one.
+      if (object->id() == last_object->id()) {
+        fprintf(stderr, "FAILED: RetrieveRangeAtIndex set %d index %d, "
+                "expected different objects, observed same objects (%d)\n",
+                set, object_index, object->id());
+        return false;
+      }
+
+      // Each object must have a base greater than the previous object's base.
+      if (base <= last_base) {
+        fprintf(stderr, "FAILED: RetrieveRangeAtIndex set %d index %d, "
+                "expected different bases, observed same bases (%d)\n",
+                set, object_index, base);
+        return false;
+      }
+    }
+
+    last_object = object.get();
+    last_base = base;
+  }
+
+  // Make sure that RetrieveRangeAtIndex doesn't allow lookups at indices that
+  // are too high.
+  if (range_map->RetrieveRangeAtIndex(object_count, &object, NULL, NULL)) {
+    fprintf(stderr, "FAILED: RetrieveRangeAtIndex set %d index %d (too large), "
+            "expected failure, observed success\n",
+            set, object_count);
+    return false;
+  }
+
+  return true;
+}
+
+
 // RunTests runs a series of test sets.
 static bool RunTests() {
   // These tests will be run sequentially.  The first set of tests exercises
@@ -373,6 +438,15 @@ static bool RunTests() {
       return false;
     }
 
+    // The RangeMap's own count of objects should also match.
+    if (range_map->GetCount() != stored_count) {
+      fprintf(stderr, "FAILED: stored object count doesn't match GetCount, "
+              "expected %d, observed %d\n",
+              stored_count, range_map->GetCount());
+
+      return false;
+    }
+
     // Run the RetrieveRange test
     for (unsigned int range_test_index = 0;
          range_test_index < range_test_count;
@@ -382,6 +456,9 @@ static bool RunTests() {
         return false;
     }
 
+    if (!RetrieveIndexTest(range_map.get(), range_test_set_index))
+      return false;
+
     // Clear the map between test sets.  If this is the final test set,
     // delete the map instead to test destruction.
     if (range_test_set_index < range_test_set_count - 1)