unionBitmapRun, bitmapSet/ZeroRange, special case optimization for RLE

roaring/roaring.go

@@ -1915,6 +1915,9 @@ func unionArrayArray(a, b *container) *container {
 // unionArrayRun optimistically assumes that the result will be a run container,
 // and converts to a bitmap or array container afterwards if necessary.
 func unionArrayRun(a, b *container) *container {
+	if b.n == 65536 {
+		return b.clone()
+	}
 	output := &container{}
 	na, nb := len(a.array), len(b.runs)
 	var vb interval32
@@ -1966,6 +1969,12 @@ func (c *container) runAppendInterval(v interval32) int {
 }
 
 func unionRunRun(a, b *container) *container {
+	if a.n == 65536 {
+		return a.clone()
+	}
+	if b.n == 65536 {
+		return b.clone()
+	}
 	na, nb := len(a.runs), len(b.runs)
 	output := &container{
 		runs: make([]interval32, 0, na+nb),
@@ -1993,8 +2002,62 @@ func unionRunRun(a, b *container) *container {
 }
 
 func unionBitmapRun(a, b *container) *container {
-	// TODO
-	return nil
+	if b.n == 65536 {
+		return b.clone()
+	}
+	output := a.clone()
+	for j := 0; j < len(b.runs); j++ {
+		output.bitmapSetRange(uint64(b.runs[j].start), uint64(b.runs[j].last))
+	}
+	return output
+}
+
+const Z = 0xFFFFFFFFFFFFFFFF
+
+// sets all bits in [i, j] (inclusive) (c must be a bitmap container)
+func (c *container) bitmapSetRange(i, j uint64) {
+	j += 1
+	x := i / 64
+	y := (j - 1) / 64
+	var X uint64 = Z << (i % 64)
+	var Y uint64 = Z >> (64 - (j % 64))
+	xcnt := popcnt(X)
+	ycnt := popcnt(Y)
+	if x == y {
+		c.n += int((j - i) - popcnt(c.bitmap[x]&(X&Y)))
+		c.bitmap[x] |= (X & Y)
+	} else {
+		c.n += int(xcnt - popcnt(c.bitmap[x]&X))
+		c.bitmap[x] |= X
+		for i := x + 1; i < y; i++ {
+			c.n += int(64 - popcnt(c.bitmap[i]))
+			c.bitmap[i] = Z
+		}
+		c.n += int(ycnt - popcnt(c.bitmap[y]&Y))
+		c.bitmap[y] |= Y
+	}
+}
+
+// zeroes all bits in [i, j] (inclusive) (c must be a bitmap container)
+func (c *container) bitmapZeroRange(i, j uint64) {
+	j += 1
+	x := i / 64
+	y := (j - 1) / 64
+	var X uint64 = Z << (i % 64)
+	var Y uint64 = Z >> (64 - (j % 64))
+	if x == y {
+		c.n -= int(popcnt(c.bitmap[x] & (X & Y)))
+		c.bitmap[x] &= ^(X & Y)
+	} else {
+		c.n -= int(popcnt(c.bitmap[x] & X))
+		c.bitmap[x] &= ^X
+		for i := x + 1; i < y; i++ {
+			c.n -= int(popcnt(c.bitmap[i]))
+			c.bitmap[i] = 0
+		}
+		c.n -= int(popcnt(c.bitmap[y] & Y))
+		c.bitmap[y] &= ^Y
+	}
 }
 
 func unionArrayBitmap(a, b *container) *container {

roaring/roaring_internal_test.go

@@ -690,3 +690,126 @@ func TestUnionArrayRun(t *testing.T) {
 		}
 	}
 }
+
+func TestBitmapSetRange(t *testing.T) {
+	c := &container{bitmap: make([]uint64, bitmapN)}
+	tests := []struct {
+		bitmap []uint64
+		start  uint64
+		last   uint64
+		exp    []uint64
+		expN   int
+	}{
+		{
+			bitmap: []uint64{0x0000000000FFF900},
+			start:  9,
+			last:   10,
+			exp:    []uint64{0x0000000000FFFF00},
+			expN:   16,
+		},
+		{
+			bitmap: []uint64{0xFF0, 0xFF, 0xFF},
+			start:  60,
+			last:   130,
+			exp:    []uint64{0xF000000000000FF0, 0xFFFFFFFFFFFFFFFF, 0xFF},
+			expN:   84,
+		},
+	}
+
+	for i, test := range tests {
+		for i, v := range test.bitmap {
+			c.bitmap[i] = v
+		}
+		c.n = c.countRange(0, 65535)
+		c.bitmapSetRange(test.start, test.last)
+		if !reflect.DeepEqual(c.bitmap[:len(test.exp)], test.exp) {
+			t.Fatalf("test %#v expected %x, got %x", i, test.exp, c.bitmap[:len(test.bitmap)])
+		}
+		if test.expN != c.n {
+			t.Fatalf("test #%v expected n to be %v, but got %v", i, test.expN, c.n)
+		}
+	}
+}
+
+func TestBitmapZeroRange(t *testing.T) {
+	c := &container{bitmap: make([]uint64, bitmapN)}
+	tests := []struct {
+		bitmap []uint64
+		start  uint64
+		last   uint64
+		exp    []uint64
+		expN   int
+	}{
+		{
+			bitmap: []uint64{0x0000000000FFFF00},
+			start:  9,
+			last:   10,
+			exp:    []uint64{0x0000000000FFF900},
+			expN:   14,
+		},
+		{
+			bitmap: []uint64{0xFF0, 0xFF, 0xFF},
+			start:  60,
+			last:   130,
+			exp:    []uint64{0xFF0, 0, 0xF8},
+			expN:   13,
+		},
+	}
+
+	for i, test := range tests {
+		for i, v := range test.bitmap {
+			c.bitmap[i] = v
+		}
+		c.n = c.countRange(0, 65535)
+		c.bitmapZeroRange(test.start, test.last)
+		if !reflect.DeepEqual(c.bitmap[:len(test.exp)], test.exp) {
+			t.Fatalf("test %#v expected %x, got %x", i, test.exp, c.bitmap[:len(test.bitmap)])
+		}
+		if test.expN != c.n {
+			t.Fatalf("test #%v expected n to be %v, but got %v", i, test.expN, c.n)
+		}
+		for i, _ := range test.bitmap {
+			c.bitmap[i] = 0
+		}
+	}
+
+}
+
+func TestUnionBitmapRun(t *testing.T) {
+	a := &container{bitmap: make([]uint64, bitmapN)}
+	b := &container{}
+	tests := []struct {
+		bitmap []uint64
+		runs   []interval32
+		exp    []uint64
+		expN   int
+	}{
+		{
+			bitmap: []uint64{2},
+			runs:   []interval32{{start: 0, last: 0}, {start: 2, last: 5}, {start: 62, last: 71}, {start: 77, last: 78}},
+			exp:    []uint64{0xC00000000000003F, 0x60FF},
+			expN:   18,
+		},
+	}
+	for i, test := range tests {
+		for i, v := range test.bitmap {
+			a.bitmap[i] = v
+		}
+		a.n = a.bitmapCountRange(0, 65535)
+		b.runs = test.runs
+		ret := unionBitmapRun(a, b)
+		if ret.isArray() {
+			ret.arrayToBitmap()
+		}
+		if !reflect.DeepEqual(ret.bitmap[:len(test.exp)], test.exp) {
+			t.Fatalf("test #%v expected %x, but got %x", i, test.exp, ret.bitmap[:len(test.exp)])
+		}
+		if ret.n != test.expN {
+			t.Fatalf("test #%v expected n to be %v, but got %v", i, test.expN, ret.n)
+		}
+		for i, _ := range test.bitmap {
+			a.bitmap[i] = 0
+		}
+	}
+
+}