IndexMap improvements

Fixed issue converting from Range to Hash based Index maps
Fixed negative key lookup on Range based Index maps
Added value assertions on creation of Index maps

JAVA-1885

Author: rozza

src/main/com/mongodb/IndexMap.java

@@ -20,37 +20,56 @@
 import java.util.Map;
 
 /**
- * <p>Efficiently maps each integer in a set to another integer in a set, useful for merging bulk write errors when a bulk write must be 
- * split into multiple batches. Has the ability to switch from a range-based to a hash-based map depending on the mappings that have 
+ * <p>Efficiently maps each integer in a set to another integer in a set, useful for merging bulk write errors when a bulk write must be
+ * split into multiple batches. Has the ability to switch from a range-based to a hash-based map depending on the mappings that have
  * been added.</p>
  *
- * <p>This class is not part of the public API.</p>
+ * <p>This class should not be considered a part of the public API.</p>
  */
-abstract class IndexMap {
+public abstract class IndexMap {
 
     /**
      * Create an empty index map.
+     *
+     * @return a new index map
      */
     static IndexMap create() {
         return new RangeBased();
     }
 
     /**
      * Create an index map that maps the integers 0..count to startIndex..startIndex + count.
+     *
+     * @param startIndex the start index
+     * @param count      the count
+     * @return an index map
      */
     static IndexMap create(final int startIndex, final int count) {
         return new RangeBased(startIndex, count);
     }
 
+    /**
+     * Add a new index to the map
+     *
+     * @param index         the index
+     * @param originalIndex the original index
+     * @return an index map with this index added to it
+     */
     abstract IndexMap add(int index, int originalIndex);
 
+    /**
+     * Return the index that the specified index is mapped to.
+     *
+     * @param index the index
+     * @return the index it's mapped to
+     */
     abstract int map(int index);
 
     private static class HashBased extends IndexMap {
         private final Map<Integer, Integer> indexMap = new HashMap<Integer, Integer>();
 
-        public HashBased(int startIndex, int count) {
-            for (int i = startIndex; i <= count; i++) {
+        public HashBased(final int startIndex, final int count) {
+            for (int i = startIndex; i < startIndex + count; i++) {
                 indexMap.put(i - startIndex, i);
             }
         }
@@ -79,6 +98,12 @@ public RangeBased() {
         }
 
         public RangeBased(final int startIndex, final int count) {
+            if (startIndex < 0) {
+                throw new IllegalArgumentException("startIndex must be more than 0");
+            }
+            if (count < 0) {
+                throw new IllegalArgumentException("count must be more than 0");
+            }
             this.startIndex = startIndex;
             this.count = count;
         }
@@ -101,8 +126,10 @@ public IndexMap add(final int index, final int originalIndex) {
 
         @Override
         public int map(final int index) {
-            if (index >= count) {
-                throw new MongoInternalException("index should not be greater than count");
+            if (index < 0) {
+                throw new MongoInternalException("no mapping found for index " + index);
+            } else if (index >= count) {
+                throw new MongoInternalException("index should not be greater than or equal to count");
             }
             return startIndex + index;
         }

src/test/com/mongodb/BulkWriteOperationSpecification.groovy

@@ -480,6 +480,27 @@ class BulkWriteOperationSpecification extends FunctionalSpecification {
         ordered << [true, false]
     }
 
+    def 'should be able to merge upserts across batches'() {
+        given:
+        def operation = initializeBulkOperation(ordered)
+
+        (0..1002).each {
+            operation.addRequest(new UpdateRequest(new BasicDBObject('key', it), true,
+                    new BasicDBObject('$set', new BasicDBObject('key', it)), false));
+            operation.addRequest(new RemoveRequest(new BasicDBObject('key', it), false));
+        }
+
+        when:
+        def result = operation.execute(WriteConcern.ACKNOWLEDGED)
+
+        then:
+        result.removedCount == result.upserts.size()
+        collection.count() == 0
+
+        where:
+        ordered << [true, false]
+    }
+
     def 'error details should have correct index on ordered write failure'() {
         given:
         def operation = initializeBulkOperation(ordered)

src/test/com/mongodb/IndexMapSpecification.groovy

@@ -33,17 +33,6 @@ class IndexMapSpecification extends Specification {
         2 == indexMap.map(1)
     }
 
-    def 'should throw on unmapped contiguous index'() {
-        given:
-        def indexMap = IndexMap.create()
-        indexMap.add(0, 1)
-        when:
-        indexMap.map(3)
-
-        then:
-        thrown(MongoInternalException)
-    }
-
     def 'should map non-contiguous indexes'() {
         given:
         def indexMap = IndexMap.create()
@@ -59,17 +48,23 @@ class IndexMapSpecification extends Specification {
         5 == indexMap.map(2)
     }
 
-    def 'should throw on unmapped non-contiguous index'() {
-        given:
-        def indexMap = IndexMap.create()
-        indexMap = indexMap.add(0, 1)
-        indexMap = indexMap.add(3, 5)
+
+    def 'should throw on unmapped index'() {
+
+        when:
+        indexMap.map(-1)
+
+        then:
+        thrown(MongoInternalException)
 
         when:
-        indexMap.map(7)
+        indexMap.map(4)
 
         then:
         thrown(MongoInternalException)
+
+        where:
+        indexMap << [IndexMap.create().add(0, 1), IndexMap.create(1000, 3).add(5, 1005)]
     }
 
     def 'should map indexes when count is provided up front'() {
@@ -80,4 +75,32 @@ class IndexMapSpecification extends Specification {
         1 == indexMap.map(0)
         2 == indexMap.map(1)
     }
-}
+
+    def 'should include ranges when converting from range based to hash based indexMap'() {
+        given:
+        def indexMap = IndexMap.create(1000, 3)
+
+        when: 'converts from range based with a high startIndex to hash based'
+        indexMap = indexMap.add(5, 1005)
+
+        then:
+        1000 == indexMap.map(0)
+        1001 == indexMap.map(1)
+        1002 == indexMap.map(2)
+        1005 == indexMap.map(5)
+    }
+
+    def 'should not allow a negative startIndex or count'() {
+        when:
+        IndexMap.create(-1, 10)
+
+        then:
+        thrown(IllegalArgumentException)
+
+        when:
+        IndexMap.create(1, -10)
+
+        then:
+        thrown(IllegalArgumentException)
+    }
+}