implements bitmap batch Direct* operations which are optimized for sorted data

also reset the data on AddN and RemoveN ops if the log write fails

roaring/roaring.go

@@ -156,7 +156,8 @@ func (b *Bitmap) Clone() *Bitmap {
 	return other
 }
 
-// Add adds values to the bitmap.
+// Add adds values to the bitmap. TODO(2.0) deprecate - use the more general
+// AddN (though be aware that it modifies 'a' in place).
 func (b *Bitmap) Add(a ...uint64) (changed bool, err error) {
 	changed = false
 	for _, v := range a {
@@ -177,19 +178,14 @@ func (b *Bitmap) Add(a ...uint64) (changed bool, err error) {
 	return changed, nil
 }
 
-// AddN adds values to the bitmap, appending them all to the op log in a batched write. It returns the number of changed bits.
+// AddN adds values to the bitmap, appending them all to the op log in a batched
+// write. It returns the number of changed bits.
 func (b *Bitmap) AddN(a ...uint64) (changed int, err error) {
 	if len(a) == 0 {
 		return 0, nil
 	}
 
-	for _, v := range a {
-		// Apply to the in-memory bitmap.
-		if b.DirectAdd(v) {
-			a[changed] = v
-			changed++
-		}
-	}
+	changed = b.DirectAddN(a...) // modifies a in-place
 
 	if b.OpWriter != nil {
 		op := &op{
@@ -197,6 +193,7 @@ func (b *Bitmap) AddN(a ...uint64) (changed int, err error) {
 			values: a[:changed],
 		}
 		if err := b.writeOp(op); err != nil {
+			b.DirectRemoveN(op.values...) // reset data since we're returning an error
 			return 0, errors.Wrap(err, "writing to op log")
 		}
 	}
@@ -204,7 +201,46 @@ func (b *Bitmap) AddN(a ...uint64) (changed int, err error) {
 	return changed, nil
 }
 
-// DirectAdd adds a value to the bitmap by bypassing the op log.
+// DirectAddN sets multiple bits in the bitmap, returning how many changed. It
+// modifies the slice 'a' in place such that once it's complete a[:changed] will
+// be list of changed bits. It is more efficient than repeated calls to
+// DirectAdd for semi-dense sorted data because it reuses the container from the
+// previous value if the new value has the same highbits instead of looking it
+// up each time. TODO: if Containers implementations cached the last few
+// Container objects returned from calls like Get and GetOrCreate, this
+// optimization would be less useful.
+func (b *Bitmap) DirectAddN(a ...uint64) (changed int) {
+	return b.directOpN((*Container).add, a...)
+}
+
+// DirectRemoveN behaves analgously to DirectAddN.
+func (b *Bitmap) DirectRemoveN(a ...uint64) (changed int) {
+	return b.directOpN((*Container).remove, a...)
+}
+
+// directOpN contains the logic for DirectAddN and DirectRemoveN. Theoretically,
+// it could be used by anything that wanted to apply a boolean-returning
+// container level operation across a list of values and return the number of
+// trues while modifying the list of values in place to contain the
+// true-returning values in order.
+func (b *Bitmap) directOpN(op func(c *Container, v uint16) bool, a ...uint64) (changed int) {
+	hb := uint64(0xFFFFFFFFFFFFFFFF) // impossible sentinel value
+	var cont *Container
+	for _, v := range a {
+		if newhb := highbits(v); newhb != hb {
+			hb = newhb
+			cont = b.Containers.GetOrCreate(hb)
+		}
+		if op(cont, lowbits(v)) {
+			a[changed] = v
+			changed++
+		}
+	}
+	return changed
+}
+
+// DirectAdd adds a value to the bitmap by bypassing the op log. TODO(2.0)
+// deprecate in favor of DirectAddN.
 func (b *Bitmap) DirectAdd(v uint64) bool {
 	cont := b.Containers.GetOrCreate(highbits(v))
 	return cont.add(lowbits(v))
@@ -219,7 +255,9 @@ func (b *Bitmap) Contains(v uint64) bool {
 	return c.Contains(lowbits(v))
 }
 
-// Remove removes values from the bitmap.
+// Remove removes values from the bitmap (writing to the op log if available).
+// TODO(2.0) deprecate - use the more general RemoveN (though be aware that it
+// modifies 'a' in place).
 func (b *Bitmap) Remove(a ...uint64) (changed bool, err error) {
 	changed = false
 	for _, v := range a {
@@ -239,18 +277,13 @@ func (b *Bitmap) Remove(a ...uint64) (changed bool, err error) {
 	return changed, nil
 }
 
+// RemoveN behaves analagously to AddN.
 func (b *Bitmap) RemoveN(a ...uint64) (changed int, err error) {
 	if len(a) == 0 {
 		return 0, nil
 	}
 
-	for _, v := range a {
-		// Apply to the in-memory bitmap.
-		if b.remove(v) {
-			a[changed] = v
-			changed++
-		}
-	}
+	changed = b.DirectRemoveN(a...) // modifies a in-place
 
 	if b.OpWriter != nil {
 		op := &op{
@@ -258,6 +291,7 @@ func (b *Bitmap) RemoveN(a ...uint64) (changed int, err error) {
 			values: a[:changed],
 		}
 		if err := b.writeOp(op); err != nil {
+			b.DirectAddN(op.values...) // reset data since we're returning an error
 			return 0, errors.Wrap(err, "writing to op log")
 		}
 	}

roaring/roaring_internal_test.go

@@ -3549,3 +3549,107 @@ func compareOps(op1, op2 *op) error {
 	}
 	return nil
 }
+
+func TestDirectAddN(t *testing.T) {
+	tests := []struct {
+		call1     []uint64
+		expn1     int
+		call2     []uint64
+		expn2     int
+		exp       []uint64
+		expcall2n []uint64
+	}{
+		{
+			call1:     []uint64{0},
+			expn1:     1,
+			call2:     []uint64{0, 1},
+			expn2:     1,
+			exp:       []uint64{0, 1},
+			expcall2n: []uint64{1},
+		},
+		{
+			call1:     []uint64{0, 22, 55},
+			expn1:     3,
+			call2:     []uint64{0, 14, 22, 99, 55},
+			expn2:     2,
+			exp:       []uint64{0, 14, 22, 55, 99},
+			expcall2n: []uint64{14, 99},
+		},
+	}
+
+	for i, test := range tests {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			b := NewBitmap()
+			n1 := b.DirectAddN(test.call1...)
+			if n1 != test.expn1 {
+				t.Errorf("mismatched n1 exp:%d got:%d", test.expn1, n1)
+			}
+			n2 := b.DirectAddN(test.call2...)
+			if n2 != test.expn2 {
+				t.Errorf("mismatched n2 exp:%d got:%d", test.expn2, n2)
+			}
+			if !reflect.DeepEqual(test.exp, b.Slice()) {
+				t.Errorf("misatched results \n%v\n%v", test.exp, b.Slice())
+			}
+			if !reflect.DeepEqual(test.expcall2n, test.call2[:n2]) {
+				t.Errorf("unexpected arg change \n%v\n%v", test.expcall2n, test.call2[:n2])
+			}
+		})
+	}
+}
+
+func TestDirectAddNVsAdd(t *testing.T) {
+	tests := [][]uint64{
+		{},
+		{0},
+		{0, 1, 2, 3},
+		{0, 1, 2, 101000, 9384932},
+		{9384932, 101000, 2, 1, 0},
+		{3489, 19230, 394, 0, 893982, 890283, 14, 7},
+	}
+	// TODO generate more tests and fuzz
+	testsCopy := make([][]uint64, len(tests))
+	copy(testsCopy, tests)
+	for i, test := range testsCopy {
+		t.Run(fmt.Sprintf("Fresh%d", i), func(t *testing.T) {
+			ba := NewBitmap()
+			bd := NewBitmap()
+			na, err := ba.Add(test...)
+			if err != nil {
+				t.Fatalf("adding bits: %v", err)
+			}
+			nd := bd.DirectAddN(test...)
+			if na != (nd > 0) {
+				t.Errorf("differing changed numbers %v, %d", na, nd)
+			}
+			if ba.Count() != bd.Count() {
+				t.Errorf("different counts")
+			}
+			if !reflect.DeepEqual(ba.Slice(), bd.Slice()) {
+				t.Errorf("unequal values\n%v\n%v", ba.Slice(), bd.Slice())
+			}
+		})
+	}
+
+	ba := NewBitmap()
+	bd := NewBitmap()
+	for i, test := range tests {
+		t.Run(fmt.Sprintf("ContinuousAdd%d", i), func(t *testing.T) {
+			na, err := ba.Add(test...)
+			if err != nil {
+				t.Fatalf("adding bits: %v", err)
+			}
+			nd := bd.DirectAddN(test...)
+			if na != (nd > 0) {
+				t.Errorf("differing changed numbers %v, %d", na, nd)
+			}
+			if ba.Count() != bd.Count() {
+				t.Errorf("different counts")
+			}
+			if !reflect.DeepEqual(ba.Slice(), bd.Slice()) {
+				t.Errorf("unequal values\n%v\n%v", ba.Slice(), bd.Slice())
+			}
+		})
+	}
+
+}