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

driver-core/src/main/com/mongodb/internal/connection/IndexMap.java

@@ -21,6 +21,8 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import static com.mongodb.assertions.Assertions.isTrueArgument;
+
 /**
  * <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
@@ -71,7 +73,7 @@ private static class HashBased extends IndexMap {
         private final Map<Integer, Integer> indexMap = new HashMap<Integer, Integer>();
 
         public HashBased(final int startIndex, final int count) {
-            for (int i = startIndex; i <= count; i++) {
+            for (int i = startIndex; i < startIndex + count; i++) {
                 indexMap.put(i - startIndex, i);
             }
         }
@@ -100,6 +102,8 @@ public RangeBased() {
         }
 
         public RangeBased(final int startIndex, final int count) {
+            isTrueArgument("startIndex", startIndex >= 0);
+            isTrueArgument("count", count > 0);
             this.startIndex = startIndex;
             this.count = count;
         }
@@ -122,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;
         }

driver-core/src/test/functional/com/mongodb/operation/MixedBulkWriteOperationAsyncSpecification.groovy

@@ -526,6 +526,28 @@ class MixedBulkWriteOperationAsyncSpecification extends OperationFunctionalSpeci
         ordered << [true, false]
     }
 
+    def 'should be able to merge upserts across batches'() {
+        given:
+        def writeOperations = [];
+        (0..1002).each {
+            def upsert = new UpdateRequest(new BsonDocument('key', new BsonInt32(it)),
+                    new BsonDocument('$set', new BsonDocument('key', new BsonInt32(it))),
+                    UPDATE).upsert(true)
+            writeOperations.add(upsert);
+            writeOperations.add(new DeleteRequest(new BsonDocument('key', new BsonInt32(it))));
+        }
+
+        when:
+        def result = executeAsync(new MixedBulkWriteOperation(getNamespace(), writeOperations, ordered, ACKNOWLEDGED))
+
+        then:
+        result.deletedCount == result.upserts.size()
+        getCollectionHelper().count() == 0
+
+        where:
+        ordered << [true, false]
+    }
+
     def 'error details should have correct index on ordered write failure'() {
         given:
         def op = new MixedBulkWriteOperation(getNamespace(),

driver-core/src/test/functional/com/mongodb/operation/MixedBulkWriteOperationSpecification.groovy

@@ -530,6 +530,28 @@ class MixedBulkWriteOperationSpecification extends OperationFunctionalSpecificat
         ordered << [true, false]
     }
 
+    def 'should be able to merge upserts across batches'() {
+        given:
+        def writeOperations = [];
+        (0..1002).each {
+            def upsert = new UpdateRequest(new BsonDocument('key', new BsonInt32(it)),
+                                           new BsonDocument('$set', new BsonDocument('key', new BsonInt32(it))),
+                        UPDATE).upsert(true)
+            writeOperations.add(upsert);
+            writeOperations.add(new DeleteRequest(new BsonDocument('key', new BsonInt32(it))));
+        }
+
+        when:
+        def result = new MixedBulkWriteOperation(getNamespace(), writeOperations, ordered, ACKNOWLEDGED).execute(getBinding())
+
+        then:
+        result.deletedCount == result.upserts.size()
+        getCollectionHelper().count() == 0
+
+        where:
+        ordered << [true, false]
+    }
+
     def 'error details should have correct index on ordered write failure'() {
         given:
         def op = new MixedBulkWriteOperation(getNamespace(),

driver-core/src/test/unit/com/mongodb/internal/connection/IndexMapSpecification.groovy

@@ -34,17 +34,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()
@@ -60,17 +49,22 @@ 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'() {
@@ -81,4 +75,33 @@ 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)
+    }
+
+}