diff --git a/roaring/roaring.go b/roaring/roaring.go index 743513562..3eca20357 100644 --- a/roaring/roaring.go +++ b/roaring/roaring.go @@ -1156,8 +1156,11 @@ func (c *container) runAdd(v uint32) bool { func (c *container) contains(v uint32) bool { if c.isArray() { return c.arrayContains(v) + } else if c.isBitmap() { + return c.bitmapContains(v) + } else { + return c.runContains(v) } - return c.bitmapContains(v) } func (c *container) arrayContains(v uint32) bool { @@ -1183,11 +1186,18 @@ func (c *container) runContains(v uint32) bool { } // remove adds a value to the container. -func (c *container) remove(v uint32) bool { +func (c *container) remove(v uint32) (removed bool) { if c.isArray() { - return c.arrayRemove(v) + removed = c.arrayRemove(v) + } else if c.isBitmap() { + removed = c.bitmapRemove(v) + } else { + removed = c.runRemove(v) } - return c.bitmapRemove(v) + if removed { + c.n-- + } + return removed } func (c *container) arrayRemove(v uint32) bool { @@ -1197,7 +1207,6 @@ func (c *container) arrayRemove(v uint32) bool { } c.unmap() - c.n-- c.array = append(c.array[:i], c.array[i+1:]...) return true } @@ -1219,6 +1228,30 @@ func (c *container) bitmapRemove(v uint32) bool { return true } +func (c *container) runRemove(v uint32) bool { + v16 := uint16(v) + for i, iv := range c.runs { + if v16 <= iv.last { + if v16 < iv.start { + return false + } + c.unmap() + if v16 == iv.last && v16 == iv.start { + c.runs = append(c.runs[:i], c.runs[i+1:]...) + } else if v16 == iv.last { + c.runs[i].last -= 1 + } else if v16 == iv.start { + c.runs[i].start += 1 + } else if v16 > iv.start { + c.runs[i].last = v16 - 1 + c.runs = append(c.runs[:i+1], append([]interval16{{start: v16 + 1, last: iv.last}}, c.runs[i+1:]...)...) + } + return true + } + } + return false +} + // max returns the maximum value in the container. func (c *container) max() uint32 { if c.isArray() { diff --git a/roaring/roaring_internal_test.go b/roaring/roaring_internal_test.go index c9a7160cd..c5736ea11 100644 --- a/roaring/roaring_internal_test.go +++ b/roaring/roaring_internal_test.go @@ -47,6 +47,7 @@ func TestContainerRunAdd(t *testing.T) { {8, []interval16{{start: 0, last: 4}, {start: 6, last: 8}, {start: 10, last: 10}}}, } for _, test := range tests { + c.mapped = true ret := c.add(test.op) if !ret { t.Fatalf("result of adding new bit should be true: %v", c.runs) @@ -54,6 +55,9 @@ func TestContainerRunAdd(t *testing.T) { if !reflect.DeepEqual(c.runs, test.exp) { t.Fatalf("Should have %v, but got %v after adding %v", test.exp, c.runs, test.op) } + if c.mapped { + t.Fatalf("container should not be mapped after adding bit %v", test.op) + } } } @@ -221,3 +225,37 @@ func TestIntersectionCountArrayBitmap2(t *testing.T) { } } } + +func TestRunRemove(t *testing.T) { + c := container{runs: []interval16{{start: 2, last: 10}, {start: 12, last: 13}, {start: 15, last: 16}}} + tests := []struct { + op uint32 + exp []interval16 + expRet bool + }{ + {2, []interval16{{start: 3, last: 10}, {start: 12, last: 13}, {start: 15, last: 16}}, true}, + {10, []interval16{{start: 3, last: 9}, {start: 12, last: 13}, {start: 15, last: 16}}, true}, + {12, []interval16{{start: 3, last: 9}, {start: 13, last: 13}, {start: 15, last: 16}}, true}, + {13, []interval16{{start: 3, last: 9}, {start: 15, last: 16}}, true}, + {16, []interval16{{start: 3, last: 9}, {start: 15, last: 15}}, true}, + {6, []interval16{{start: 3, last: 5}, {start: 7, last: 9}, {start: 15, last: 15}}, true}, + {8, []interval16{{start: 3, last: 5}, {start: 7, last: 7}, {start: 9, last: 9}, {start: 15, last: 15}}, true}, + {8, []interval16{{start: 3, last: 5}, {start: 7, last: 7}, {start: 9, last: 9}, {start: 15, last: 15}}, false}, + {1, []interval16{{start: 3, last: 5}, {start: 7, last: 7}, {start: 9, last: 9}, {start: 15, last: 15}}, false}, + {44, []interval16{{start: 3, last: 5}, {start: 7, last: 7}, {start: 9, last: 9}, {start: 15, last: 15}}, false}, + } + + for _, test := range tests { + c.mapped = true + ret := c.remove(test.op) + if ret != test.expRet || !reflect.DeepEqual(c.runs, test.exp) { + t.Fatalf("Unexpected result removing %v from runs. Expected %v, got %v. Expected %v, got %v", test.op, test.expRet, ret, test.exp, c.runs) + } + if ret && c.mapped { + t.Fatalf("container was not unmapped although bit %v was removed", test.op) + } + if !ret && !c.mapped { + t.Fatalf("container was unmapped although bit %v was not removed", test.op) + } + } +}