create BitmapRewriter/ApplyRewriter, parallel to BitmapFilter

This in a parallel to ApplyFilter/BitmapFilter which allows writebacks
while it's running. It's a write operation, so it needs a write lock
on the Tx, and needs to create bitmaps if they don't already exist.
The semantics are a bit messy and need better documentation still.

catcher.go

@@ -234,6 +234,10 @@ func (c *catcherTx) ApplyFilter(index, field, view string, shard uint64, ckey ui
 	return GenericApplyFilter(c, index, field, view, shard, ckey, filter)
 }
 
+func (c *catcherTx) ApplyRewriter(index, field, view string, shard uint64, ckey uint64, filter roaring.BitmapRewriter) (err error) {
+	return c.b.ApplyRewriter(index, field, view, shard, ckey, filter)
+}
+
 func (c *catcherTx) GetSortedFieldViewList(idx *Index, shard uint64) (fvs []txkey.FieldView, err error) {
 	return c.b.GetSortedFieldViewList(idx, shard)
 }

rbf.go

@@ -411,6 +411,10 @@ func (tx *RBFTx) ApplyFilter(index, field, view string, shard uint64, ckey uint6
 	return tx.tx.ApplyFilter(rbfName(index, field, view, shard), ckey, filter)
 }
 
+func (tx *RBFTx) ApplyRewriter(index, field, view string, shard uint64, ckey uint64, filter roaring.BitmapRewriter) (err error) {
+	return tx.tx.ApplyRewriter(rbfName(index, field, view, shard), ckey, filter)
+}
+
 func (tx *RBFTx) GetSortedFieldViewList(idx *Index, shard uint64) (fvs []txkey.FieldView, err error) {
 	return tx.tx.GetSortedFieldViewList()
 }

rbf/cursorx.go

@@ -191,8 +191,37 @@ func intoContainer(l leafCell, tx *Tx, replacing *roaring.Container, target []by
 	return c
 }
 
-func toContainer(l leafCell, tx *Tx) (c *roaring.Container) {
+// intoWritableContainer always uses the provided target for a copy of
+// the container's contents, so the container can be modified safely.
+func intoWritableContainer(l leafCell, tx *Tx, replacing *roaring.Container, target []byte) (c *roaring.Container) {
+	if len(l.Data) == 0 {
+		return nil
+	}
+	orig := l.Data
+	target = target[:len(orig)]
+	copy(target, orig)
+	switch l.Type {
+	case ContainerTypeArray:
+		c = roaring.RemakeContainerArray(replacing, toArray16(target))
+	case ContainerTypeBitmapPtr:
+		pgno := toPgno(target)
+		target = target[:PageSize] // reslice back to full size
+		_, bm, _ := tx.leafCellBitmapInto(pgno, target)
+		c = roaring.RemakeContainerBitmapN(replacing, bm, int32(l.BitN))
+	case ContainerTypeBitmap:
+		c = roaring.RemakeContainerBitmapN(replacing, toArray64(target), int32(l.BitN))
+	case ContainerTypeRLE:
+		c = roaring.RemakeContainerRunN(replacing, toInterval16(target), int32(l.BitN))
+	}
+	// Note: If the "roaringparanoia" build tag isn't set, this
+	// should be optimized away entirely. Otherwise it's moderately
+	// expensive.
+	c.CheckN()
+	c.SetMapped(false)
+	return c
+}
 
+func toContainer(l leafCell, tx *Tx) (c *roaring.Container) {
 	if len(l.Data) == 0 {
 		return nil
 	}

rbf/tx.go

@@ -3,6 +3,7 @@ package rbf
 
 import (
 	"bufio"
+	"errors"
 	"fmt"
 	"io"
 	"math"
@@ -1301,6 +1302,16 @@ func (tx *Tx) leafCellBitmap(pgno uint32) (uint32, []uint64, error) {
 	return pgno, toArray64(page), err
 }
 
+// leafCellBitmapInto copies the bitmap into provided space
+func (tx *Tx) leafCellBitmapInto(pgno uint32, into []byte) (uint32, []uint64, error) {
+	page, _, err := tx.readPage(pgno)
+	if err != nil {
+		return 0, nil, err
+	}
+	copy(into, page)
+	return pgno, toArray64(into), err
+}
+
 func (tx *Tx) ContainerIterator(name string, key uint64) (citer roaring.ContainerIterator, found bool, err error) {
 	tx.mu.RLock()
 	defer tx.mu.RUnlock()
@@ -1319,22 +1330,52 @@ func (tx *Tx) ContainerIterator(name string, key uint64) (citer roaring.Containe
 	return &containerIterator{cursor: c}, exact, nil
 }
 
-// Shared pool for in-memory database pages.
-// These are used before being flushed to disk.
+// Shared pool for container filters, used because they contain cursors
+// which are large.
 var containerFilterPool = &sync.Pool{}
 
-func getContainerFilter(c *Cursor, filter roaring.BitmapFilter, tx *Tx) *containerFilter {
+// getContainerFilter generates a containerFilter, which may actually secretly
+// be used for rewriting; the data structures are similar enough that sharing
+// a pool for both types seems advantageous.
+func getContainerFilter(c *Cursor, name string, filter roaring.BitmapFilter, rewriter roaring.BitmapRewriter, tx *Tx) *containerFilter {
 	existing := containerFilterPool.Get()
 	if existing == nil {
-		return &containerFilter{cursor: c, filter: filter, tx: tx}
+		return &containerFilter{cursor: c, name: name, filter: filter, rewriter: rewriter, tx: tx}
 	}
 	f := existing.(*containerFilter)
 	f.cursor = c
+	f.name = name
 	f.filter = filter
+	f.rewriter = rewriter
 	f.tx = tx
 	return f
 }
 
+func (tx *Tx) ApplyRewriter(name string, key uint64, rewriter roaring.BitmapRewriter) (err error) {
+	tx.mu.Lock()
+	defer tx.mu.Unlock()
+
+	// Unlike a Filter, Rewriter makes sense to apply to an empty bitmap.
+	if err = tx.createBitmapIfNotExists(name); err != nil {
+		return err
+	}
+
+	c, err := tx.cursor(name)
+	if err == ErrBitmapNotFound {
+		return nil // nothing available.
+	} else if err != nil {
+		return err
+	}
+
+	_, err = c.Seek(key)
+	if err != nil {
+		return err
+	}
+	f := getContainerFilter(c, name, nil, rewriter, tx)
+	defer f.Close()
+	return f.ApplyRewriter()
+}
+
 func (tx *Tx) ApplyFilter(name string, key uint64, filter roaring.BitmapFilter) (err error) {
 	tx.mu.RLock()
 	defer tx.mu.RUnlock()
@@ -1350,9 +1391,9 @@ func (tx *Tx) ApplyFilter(name string, key uint64, filter roaring.BitmapFilter)
 	if err != nil {
 		return err
 	}
-	f := getContainerFilter(c, filter, tx)
+	f := getContainerFilter(c, name, filter, nil, tx)
 	defer f.Close()
-	return f.Apply()
+	return f.ApplyFilter()
 }
 
 func (tx *Tx) ForEach(name string, fn func(i uint64) error) error {
@@ -1681,13 +1722,16 @@ func (tx *Tx) OffsetRange(name string, offset, start, endx uint64) (*roaring.Bit
 	return other, nil
 }
 
-// containerFilter is like ContainerIterator, but implements ApplyFilter
+// containerFilter is like ContainerIterator, but implements ApplyFilter and
+// also ApplyRewriter, depending on which is provided to it.
 type containerFilter struct {
-	cursor *Cursor
-	filter roaring.BitmapFilter
-	tx     *Tx
-	header roaring.Container
-	body   [8192]byte
+	cursor   *Cursor
+	name     string
+	filter   roaring.BitmapFilter
+	rewriter roaring.BitmapRewriter
+	tx       *Tx
+	header   roaring.Container
+	body     [8192]byte
 }
 
 func (s *containerFilter) Close() {
@@ -1696,9 +1740,12 @@ func (s *containerFilter) Close() {
 	containerFilterPool.Put(s)
 }
 
-func (s *containerFilter) Apply() (err error) {
+func (s *containerFilter) ApplyFilter() (err error) {
 	var minKey roaring.FilterKey
 	var cell leafCell
+	if s.filter == nil {
+		return errors.New("can't apply filter without a filter")
+	}
 	for err := s.cursor.Next(); err == nil; err = s.cursor.Next() {
 		elem := &s.cursor.stack.elems[s.cursor.stack.top]
 		leafPage, _, _ := s.cursor.tx.readPage(elem.pgno)
@@ -1733,6 +1780,81 @@ func (s *containerFilter) Apply() (err error) {
 	return nil
 }
 
+func (s *containerFilter) ApplyRewriter() (err error) {
+	var minKey roaring.FilterKey
+	var cell leafCell
+	if s.rewriter == nil {
+		return errors.New("can't apply rewriter without a rewriter")
+	}
+	var dirty bool
+	var key roaring.FilterKey
+	var writeback roaring.ContainerWriteback = func(updateKey roaring.FilterKey, data *roaring.Container) (err error) {
+		var exact bool
+		if updateKey != key {
+			exact, err = s.cursor.Seek(uint64(updateKey))
+			if err != nil {
+				return err
+			}
+		} else {
+			exact = true
+		}
+		if data.N() == 0 {
+			if exact {
+				err = s.cursor.deleteLeafCell(uint64(updateKey))
+				key = ^roaring.FilterKey(0)
+			}
+			// if we don't delete, we aren't changing our situation at all
+		} else {
+			cell = ConvertToLeafArgs(uint64(updateKey), data)
+			err = s.cursor.putLeafCell(cell)
+			key = ^roaring.FilterKey(0)
+		}
+		dirty = true
+		return err
+	}
+	for err := s.cursor.Next(); err == nil; err = s.cursor.Next() {
+		elem := &s.cursor.stack.elems[s.cursor.stack.top]
+		leafPage, _, _ := s.cursor.tx.readPage(elem.pgno)
+		readLeafCellInto(&cell, leafPage, elem.index)
+		key = roaring.FilterKey(cell.Key)
+		if key < minKey {
+			continue
+		}
+		s.tx.mu.RUnlock()
+		res := s.rewriter.ConsiderKey(key, int32(cell.BitN))
+		s.tx.mu.RLock()
+		if res.Err != nil {
+			return res.Err
+		}
+		if res.YesKey <= key && res.NoKey <= key {
+			data := intoWritableContainer(cell, s.cursor.tx, &s.header, s.body[:])
+			s.tx.mu.RUnlock()
+			res = s.rewriter.RewriteData(key, data, writeback)
+			s.tx.mu.RLock()
+			if res.Err != nil {
+				return res.Err
+			}
+		}
+		minKey = res.NoKey
+		// if the callback did any writing, we need to reset our cursor,
+		// and if the next key is far away, we should also reset our cursor.
+		//
+		// In practice the "key+64" probably comes out to "we've been told
+		// we're done".
+		if dirty || minKey > (key+64) {
+			dirty = false
+			_, err := s.cursor.Seek(uint64(minKey))
+			if err != nil {
+				return err
+			}
+		}
+	}
+	// notify rewriter that we're done, telling them the last key we
+	// processed.
+	res := s.rewriter.RewriteData(^roaring.FilterKey(0), nil, writeback)
+	return res.Err
+}
+
 // containerIterator wraps Cursor to implement roaring.ContainerIterator.
 type containerIterator struct {
 	cursor *Cursor

roaring/filter.go

@@ -190,6 +190,30 @@ type BitmapFilter interface {
 	ConsiderData(key FilterKey, data *Container) FilterResult
 }
 
+// ContainerWriteback is the type for functions which can feed updated
+// containers back to things from filters.
+type ContainerWriteback func(key FilterKey, data *Container) error
+
+// A BitmapRewriter is like a bitmap filter, but can modify the bitmap
+// it's being called on during the iteration.
+//
+// After the last container is returned, ConsiderData will be called with
+// an unspecified key and a nil container pointer, so the rewriter can
+// write any trailing containers it has. A nil container passed to writeback
+// implies a delete operation on the container. Writeback should only be
+// called with keys greater than any previously given container key, and
+// less than or equal to the current key. So for instance, if
+// ConsiderData is called with key 3, and then with key 5, the call with
+// key 5 may call writeback with key 4, and then key 5, but may not call
+// it with keys 3 or lower, or 6 or higher. When the container provided to
+// the call is nil, any monotonically increasing keys greater than the
+// previous key are allowed. (If there was no previous key, 0 and higher
+// are allowed.)
+type BitmapRewriter interface {
+	ConsiderKey(key FilterKey, n int32) FilterResult
+	RewriteData(key FilterKey, data *Container, writeback ContainerWriteback) FilterResult
+}
+
 // BitmapColumnFilter is a BitmapFilter which checks for containers matching
 // a given column within a row; thus, only the one container per row which
 // matches the column needs to be evaluated, and it's evaluated as matching
@@ -849,6 +873,76 @@ func (b *BitmapMutexDupFilter) Report() map[uint64][]uint64 {
 	return b.extra
 }
 
+// BitmapBitmapTrimmer is like BitmapBitmapFilter, but instead of calling
+// a callback per bit found in the intersection, it calls a callback with the
+// original raw container and the corresponding filter container, and also
+// provides the writeback func it got from the bitmap. So for instance, to
+// implement a "subtract these bits" function, you would difference-in-place
+// the raw container with the filter container, then pass that to the writeback
+// function.
+//
+// It's called a Trimmer because it won't add containers; it won't *add*
+// containers. It calls the callback function for every container, whether or
+// not it matches the filter; this allows an intersect-like filter to work
+// too.
+//
+// Note, however, that the caller's ContainerWriteback function *may* create
+// containers, even though the Trimmer won't have called RewriteData with those
+// keys.
+type BitmapBitmapTrimmer struct {
+	containers []*Container
+	callback   func(key FilterKey, raw, filter *Container, writeback ContainerWriteback) error
+}
+
+var _ BitmapRewriter = &BitmapBitmapTrimmer{}
+
+func (b *BitmapBitmapTrimmer) SetCallback(cb func(FilterKey, *Container, *Container, ContainerWriteback) error) {
+	b.callback = cb
+}
+
+func (b *BitmapBitmapTrimmer) ConsiderKey(key FilterKey, n int32) FilterResult {
+	pos := key & keyMask
+	if b.containers[pos] == nil || n == 0 {
+		return key.RejectOne()
+	}
+	return key.NeedData()
+}
+
+func (b *BitmapBitmapTrimmer) RewriteData(key FilterKey, data *Container, writeback ContainerWriteback) FilterResult {
+	pos := key & keyMask
+	filter := b.containers[pos]
+	err := b.callback(key, data, filter, writeback)
+	if err != nil {
+		return key.Fail(err)
+	}
+	return key.MatchOne()
+}
+
+// NewBitmapBitmapTrimmer creates a filter which calls a callback on every
+// container in a bitmap, with corresponding elements from an initial filter
+// bitmap. It does not call its callback for cases where there's no container
+// in the original bitmap.
+//
+// The input filter is assumed to represent one "row" of a shard's data,
+// which is to say, a range of up to rowWidth consecutive containers starting
+// at some multiple of rowWidth. We coerce that to the 0..rowWidth range
+// because offset-within-row is what we care about.
+func NewBitmapBitmapTrimmer(filter *Bitmap, callback func(FilterKey, *Container, *Container, ContainerWriteback) error) *BitmapBitmapTrimmer {
+	b := &BitmapBitmapTrimmer{
+		callback:   callback,
+		containers: make([]*Container, rowWidth),
+	}
+	iter, _ := filter.Containers.Iterator(0)
+	for iter.Next() {
+		k, v := iter.Value()
+		// Coerce container key into the 0-rowWidth range we'll be
+		// using to compare against containers within each row.
+		k = k & keyMask
+		b.containers[k] = v
+	}
+	return b
+}
+
 // ApplyFilterToIterator is a simplistic implementation that applies a bitmap
 // filter to a ContainerIterator, returning an error if it encounters an error.
 //

stattx.go

@@ -411,6 +411,10 @@ func (c *statTx) ApplyFilter(index, field, view string, shard uint64, ckey uint6
 	return GenericApplyFilter(c, index, field, view, shard, ckey, filter)
 }
 
+func (c *statTx) ApplyRewriter(index, field, view string, shard uint64, ckey uint64, filter roaring.BitmapRewriter) (err error) {
+	return c.b.ApplyRewriter(index, field, view, shard, ckey, filter)
+}
+
 func (c *statTx) ForEach(index, field, view string, shard uint64, fn func(i uint64) error) error {
 	me := kForEach
 

tx.go

@@ -75,6 +75,16 @@ type Tx interface {
 	// must copy it into some other memory.
 	ApplyFilter(index, field, view string, shard uint64, ckey uint64, filter roaring.BitmapFilter) (err error)
 
+	// ApplyRewriter applies a roaring.BitmapRewriter to a specified shard,
+	// starting at the given container key. The filter's ConsiderData
+	// method may be called with transient Container objects which *must
+	// not* be retained or referenced after that function exits. Similarly,
+	// their data must not be retained. If you need the data later, you
+	// must copy it into some other memory. However, it is safe to overwrite
+	// the returned container; for instance, you can DifferenceInPlace on
+	// it.
+	ApplyRewriter(index, field, view string, shard uint64, ckey uint64, filter roaring.BitmapRewriter) (err error)
+
 	// RoaringBitmap retrieves the roaring.Bitmap for the entire shard.
 	RoaringBitmap(index, field, view string, shard uint64) (*roaring.Bitmap, error)