featurebase/rbf/cursor.go
2022-09-06 09:39:22 -07:00

1501 lines
41 KiB
Go

// Copyright 2022 Molecula Corp. (DBA FeatureBase).
// SPDX-License-Identifier: Apache-2.0
package rbf
import (
"encoding/binary"
"fmt"
"io"
"sort"
"unsafe"
"github.com/featurebasedb/featurebase/v3/roaring"
"github.com/pkg/errors"
)
const (
BitmapN = (1 << 16) / 64
)
// Cursor represents an object for traversing forward and backward in the
// keyspace. Users should initialize a cursor with either First(), Last(), or
// Seek() and then move forward or backward using Next() or Prev(), respectively.
//
// If you mutate the b-tree that a cursor is attached to then you'll need to
// re-initialize the cursor. This may be changed in the future though.
//
// Cursors can be reused in some cases; if you're keeping a cursor around,
// be sure to nil out the tx, or it will hold a reference to it.
type Cursor struct {
tx *Tx
buffered bool // if true, Next() and Prev() do not move the cursor position
// buffers
array [ArrayMaxSize + 1]uint16
rle [RLEMaxSize + 1]roaring.Interval16
leafCells [PageSize / 8]leafCell
// stack holds branches
stack searchStack
}
type searchStack struct {
top int // current depth in elems
elems [32]stackElem
}
func runAdd(runs []roaring.Interval16, v uint16) ([]roaring.Interval16, bool) {
i := sort.Search(len(runs),
func(i int) bool { return runs[i].Last >= v })
if i == len(runs) {
i--
}
iv := runs[i]
if v >= iv.Start && iv.Last >= v {
return nil, false
}
if iv.Last < v {
if iv.Last == v-1 {
runs[i].Last++
} else {
runs = append(runs, roaring.Interval16{Start: v, Last: v})
}
} else if v+1 == iv.Start {
// combining two intervals
if i > 0 && runs[i-1].Last == v-1 {
runs[i-1].Last = iv.Last
runs = append(runs[:i], runs[i+1:]...)
//TODO check if too big
return runs, true
}
// just before an interval
runs[i].Start--
} else if i > 0 && v-1 == runs[i-1].Last {
// just after an interval
runs[i-1].Last++
} else {
// alone
newIv := roaring.Interval16{Start: v, Last: v}
runs = append(runs[:i], append([]roaring.Interval16{newIv}, runs[i:]...)...)
}
return runs, true
}
// maybeConvertOversizedRunToBitmap is only called by Cursor.Add().
// bitN must be the correct new bit count.
func maybeConvertOversizedRunToBitmap(runs []roaring.Interval16, bitN int, key uint64) leafCell {
if len(runs) >= RLEMaxSize {
//convertToBitmap
bitmap := make([]uint64, BitmapN)
for _, iv := range runs {
w1, w2 := iv.Start/64, iv.Last/64
b1, b2 := iv.Start&63, iv.Last&63
// a mask for everything under bit X looks like
// (1 << x) - 1. Say b1 is 4; our mask will want
// to have the bottom 4 bits be zero, so we shift
// left 4, getting 10000, then subtract 1, and
// get 01111, which is the mask to *remove*.
m1 := (uint64(1) << b1) - 1
// inclusive mask: same thing, then shift left 1 and
// or in 1. So for 4, we'd get 011111, which is the
// mask to *keep*.
m2 := (((uint64(1) << b2) - 1) << 1) | 1
if w1 == w2 {
// If we only had bit 4 in the range, this would
// end up being 011111 &^ 01111, or 010000.
bitmap[w1] |= (m2 &^ m1)
continue
}
// for w2, the "To" field, we want to set the bottom N
// bits. For w1, the "From" word, we want to set all *but*
// the bottom N bits.
bitmap[w2] |= m2
bitmap[w1] |= ^m1
words := bitmap[w1+1 : w2]
// set every bit between them
for i := range words {
words[i] = ^uint64(0)
}
}
// note that ElemN should be left 0 for ContainerTypeBitmap.
return leafCell{Key: key, BitN: int(bitN), Type: ContainerTypeBitmap, Data: fromArray64(bitmap)}
}
return leafCell{Key: key, ElemN: len(runs), BitN: int(bitN), Type: ContainerTypeRLE, Data: fromInterval16(runs)}
}
// Add sets a bit on the underlying bitmap.
// If changed is true, this cursor will need to be reinitialized before use.
func (c *Cursor) Add(v uint64) (changed bool, err error) {
hi, lo := highbits(v), lowbits(v)
// Move cursor to the key of the container.
// Insert new container if it doesn't exist.
if exact, err := c.Seek(hi); err != nil {
return false, err
} else if !exact {
return true, c.putLeafCell(leafCell{Key: hi, Type: ContainerTypeArray, ElemN: 1, BitN: 1, Data: fromArray16([]uint16{lo})})
}
// If the container exists and bit is not set then update the page.
elem := &c.stack.elems[c.stack.top]
leafPage, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return false, err
}
cell := readLeafCell(leafPage, elem.index)
switch cell.Type {
case ContainerTypeArray:
// Exit if value exists in array container.
a := toArray16(cell.Data)
i, ok := arrayIndex(a, lo)
if ok {
return false, nil
}
// Copy container data and insert new value.
other := c.array[:len(a)+1]
copy(other, a[:i])
other[i] = lo
copy(other[i+1:], a[i:])
return true, c.putLeafCell(leafCell{Key: cell.Key, Type: ContainerTypeArray, ElemN: len(other), BitN: cell.BitN + 1, Data: fromArray16(other)})
case ContainerTypeRLE:
runs := toInterval16(cell.Data)
copy(c.rle[:], runs)
run, added := runAdd(c.rle[:len(runs)], lo)
if added {
leaf := maybeConvertOversizedRunToBitmap(run, cell.BitN+1, cell.Key)
return true, c.putLeafCell(leaf)
}
return false, nil
case ContainerTypeBitmapPtr:
// Exit if bit set in bitmap container.
pgno, bm, err := c.tx.leafCellBitmap(toPgno(cell.Data))
if err != nil {
return false, errors.Wrap(err, "cursor.Add")
}
a := cloneArray64(bm)
if a[lo/64]&(1<<uint64(lo%64)) != 0 {
return false, nil
}
// Insert new value and rewrite page.
a[lo/64] |= 1 << uint64(lo%64)
if err := c.tx.writeBitmapPage(pgno, fromArray64(a)); err != nil {
return false, err
}
// Update with new BitN.
cell.BitN++
return true, c.putLeafCell(cell)
default:
return false, fmt.Errorf("rbf.Cursor.Add(): invalid container type: %d", cell.Type)
}
}
// Remove unsets a bit on the underlying bitmap.
// If changed is true, the cursor will need to be reinitialized before use.
func (c *Cursor) Remove(v uint64) (changed bool, err error) {
hi, lo := highbits(v), lowbits(v)
// Move cursor to the key of the container.
// Exit if container does not exist.
if exact, err := c.Seek(hi); err != nil {
return false, err
} else if !exact {
return false, nil
}
// If the container exists and bit is not set then update the page.
elem := &c.stack.elems[c.stack.top]
leafPage, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return false, err
}
cell := readLeafCell(leafPage, elem.index)
switch cell.Type {
case ContainerTypeArray:
// Exit if value does not exists in array container.
a := toArray16(cell.Data)
i, ok := arrayIndex(a, lo)
if !ok {
return false, nil
} else if len(a) == 1 {
return true, c.deleteLeafCell(cell.Key)
}
// Copy container data and remove new value.
other := c.array[:len(a)-1]
copy(other[:i], a[:i])
copy(other[i:], a[i+1:])
return true, c.putLeafCell(leafCell{Key: cell.Key, Type: ContainerTypeArray, ElemN: len(other), BitN: cell.BitN - 1, Data: fromArray16(other)})
case ContainerTypeRLE:
r := toInterval16(cell.Data)
i, contains := roaring.BinSearchRuns(lo, r)
if !contains {
return false, nil
}
// INVAR: lo is in run[i]
copy(c.rle[:], r)
runs := c.rle[:len(r)]
if lo == runs[i].Last && lo == runs[i].Start {
runs = append(runs[:i], runs[i+1:]...)
} else if lo == runs[i].Last {
runs[i].Last--
} else if lo == c.rle[i].Start {
runs[i].Start++
} else if lo > runs[i].Start {
// INVAR: Start < lo < Last.
// We remove lo, so split into two runs:
last := runs[i].Last
runs[i].Last = lo - 1
// INVAR: runs[:i] is correct, but still need to insert the new interval at i+1.
runs = append(runs, roaring.Interval16{})
// copy the tail first
copy(runs[i+2:], runs[i+1:])
// overwrite with the new interval.
runs[i+1] = roaring.Interval16{Start: lo + 1, Last: last}
}
if len(runs) == 0 {
return true, c.deleteLeafCell(cell.Key)
}
return true, c.putLeafCell(leafCell{Key: cell.Key, Type: ContainerTypeRLE, ElemN: len(runs), BitN: cell.BitN - 1, Data: fromInterval16(runs)})
case ContainerTypeBitmapPtr:
pgno, bm, err := c.tx.leafCellBitmap(toPgno(cell.Data))
if err != nil {
return false, errors.Wrap(err, "cursor.add")
}
a := cloneArray64(bm)
if a[lo/64]&(1<<uint64(lo%64)) == 0 {
// not present.
return false, nil
}
// clear the bit
a[lo/64] &^= 1 << uint64(lo%64)
cell.BitN--
if cell.BitN == 0 {
if err := c.tx.freePgno(pgno); err != nil {
return false, err
}
return true, c.deleteLeafCell(cell.Key)
}
// shrink if we've gotten small.
if cell.BitN <= ArrayMaxSize {
cbm := roaring.NewContainerBitmap(cell.BitN, a)
// convert to array
cbm = roaring.Optimize(cbm)
leafCell1 := ConvertToLeafArgs(cell.Key, cbm)
// ConvertToLeafArgs returns leafCell1 with BitN and ElemN updated.
return true, c.putLeafCell(leafCell1)
}
// rewrite page, still as a bitmap.
if err := c.tx.writeBitmapPage(pgno, fromArray64(a)); err != nil {
return false, err
}
return true, c.putLeafCell(cell)
default:
return false, fmt.Errorf("rbf.Cursor.Add(): invalid container type: %d", cell.Type)
}
}
// Contains returns true if a bit is set on the underlying bitmap.
func (c *Cursor) Contains(v uint64) (exists bool, err error) {
hi, lo := highbits(v), lowbits(v)
// Move cursor to the key of the container.
if exact, err := c.Seek(hi); err != nil {
return false, err
} else if !exact {
return false, nil
}
// If the container exists then check for low bits existence.
elem := &c.stack.elems[c.stack.top]
leafPage, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return false, err
}
cell := readLeafCell(leafPage, elem.index)
switch cell.Type {
case ContainerTypeArray:
a := toArray16(cell.Data)
_, ok := arrayIndex(a, lo)
return ok, nil
case ContainerTypeRLE:
a := toInterval16(cell.Data)
i := int32(sort.Search(len(a),
func(i int) bool { return a[i].Last >= lo }))
if i < int32(len(a)) {
return (lo >= a[i].Start) && (lo <= a[i].Last), nil
}
return false, nil
case ContainerTypeBitmapPtr:
_, a, err := c.tx.leafCellBitmap(toPgno(cell.Data))
if err != nil {
return false, errors.Wrap(err, "cursor.Contains")
}
return a[lo/64]&(1<<uint64(lo%64)) != 0, err
default:
return false, fmt.Errorf("rbf.Cursor.Contains(): invalid container type: %d", cell.Type)
}
}
func fromPgno(val uint32) []byte {
buf := make([]byte, 4)
binary.LittleEndian.PutUint32(buf, val)
// return (*[4]byte)(unsafe.Pointer(&val))[:]
return buf
}
func toPgno(val []byte) uint32 {
return binary.LittleEndian.Uint32(val)
}
// putLeafCell inserts or updates a cell into the currently position leaf page
// in the stack. If in.Key exists in the page, the cell is updated. Otherwise,
// cell is inserted.
//
// This method has two paths. First, a fast path is used if we know the insert
// will not cause the page to grow past the max page size. This path simply
// shifts bytes around in the page to make room for the new cell. The second
// path is the slow path where a page may split into two pages. This causes the
// entire page to be deserialized so that pages can more easily be split.
//
// The slow path may cause a cascade of changes to parent pages as a page split
// creates a new entry in the parent. If that entry overflows the parent then
// the parent branch page will split which could cascade up to the root. If the
// root page splits, a new branch page is created but retains the original root
// page number so that the root records do not need to be updated.
func (c *Cursor) putLeafCell(in leafCell) (err error) {
elem := &c.stack.elems[c.stack.top]
leafPage, isHeap, err := c.tx.readPage(elem.pgno) // the last read leaf page
if err != nil {
return err
}
cellN := readCellN(leafPage)
// Determine if the insert/update will overflow the page.
// If it doesn't then we can do an optimized write where we don't deserialize.
isInsert := elem.index >= cellN || pageKeyAt(leafPage, elem.index) != in.Key
newEstPageSize := leafPageSize(leafPage)
if isInsert {
newEstPageSize += in.Size() + leafCellIndexElemSize
} else {
newEstPageSize += in.Size() - len(readLeafCellBytesAtOffset(leafPage, readCellOffset(leafPage, elem.index)))
}
// Use the fast path if we are not splitting pages and the container types are the same.
useFast := newEstPageSize+16 <= PageSize
if useFast && !isInsert {
if prev := readLeafCell(leafPage, elem.index); prev.Type != in.Type {
useFast = false
}
}
// Use an optimized routine to insert the leaf cell if we won't overflow.
// We pad the estimate with 16 bytes because we do 8-byte alignment of
// both the cell and the index.
if useFast {
return c.putLeafCellFast(in, isInsert)
}
cells := readLeafCells(leafPage, c.leafCells[:])
cell := in
if isInsert {
//new cell
if in.Type == ContainerTypeBitmap {
//allocated bitmap()
bitmapPgno, err := c.tx.allocatePgno()
if err != nil {
return err
}
cell.Data = fromPgno(bitmapPgno)
cell.Type = ContainerTypeBitmapPtr
cell.BitN = in.BitN
cell.ElemN = in.ElemN
}
// Shift cells over if this is an insertion.
cells = append(cells, leafCell{})
copy(cells[elem.index+1:], cells[elem.index:])
} else {
// FB-1239: Free bitmap page if replaced container is a bitmap pointer.
prev := cells[elem.index]
if prev.Type == ContainerTypeBitmapPtr {
if in.Type == ContainerTypeBitmapPtr {
if toPgno(in.Data) != toPgno(prev.Data) { // bptr-to-bptr with different bitmap pages
if err := c.tx.freePgno(toPgno(prev.Data)); err != nil {
return err
}
}
} else if in.Type != ContainerTypeBitmap {
if err := c.tx.freePgno(toPgno(prev.Data)); err != nil {
return err
}
}
}
if in.Type == ContainerTypeBitmap {
cell = cells[elem.index]
if cell.Type != ContainerTypeBitmapPtr {
bitmapPgno, err := c.tx.allocatePgno()
if err != nil {
return errors.Wrap(err, "cursor.putLeafCell")
}
cell.Type = ContainerTypeBitmapPtr
cell.Data = fromPgno(bitmapPgno)
cell.ElemN = in.ElemN
}
// update the BitN regardless
cell.BitN = in.BitN
}
}
if in.Type == ContainerTypeArray && in.ElemN > ArrayMaxSize {
//convert to bitmap
in.Type = ContainerTypeBitmap
a := make([]uint64, PageSize/8)
for _, v := range toArray16(in.Data) {
a[v/64] |= 1 << uint64(v%64)
}
in.Data = fromArray64(a)
cell.Type = ContainerTypeBitmapPtr
bitmapPgno, err := c.tx.allocatePgno()
if err != nil {
return err
}
cell.Data = fromPgno(bitmapPgno)
}
cells[elem.index] = cell
// Split into multiple pages if page size is exceeded.
groups := [][]leafCell{cells}
sz := leafCellsPageSize(cells)
if sz >= PageSize {
groups = splitLeafCells(cells)
}
// Write each group to a separate page.
newRoot := (len(groups) > 1) && (c.stack.top == 0)
var parents []branchCell
origPgno := elem.pgno
// newRoot if split occured and bottom of the stack
for i, group := range groups {
// First page should overwrite the original.
// Subsequent pages should allocate new pages.
parent := branchCell{LeftKey: group[0].Key} //<<< this is the key spot for making sure that key is correct
if i == 0 && !newRoot {
parent.ChildPgno = origPgno
} else {
if parent.ChildPgno, err = c.tx.allocatePgno(); err != nil {
return fmt.Errorf("cannot allocate leaf: %w", err)
}
}
// if the cell is a bitmap write out its page
if in.Type == ContainerTypeBitmap {
var bm [PageSize]byte
copy(bm[:], fromArray64(toArray64(in.Data)))
if err = c.tx.writeBitmapPage(toPgno(cell.Data), bm[:]); err != nil {
return errors.Wrap(err, "putLeafCell writing bitmap page")
}
}
var buf [PageSize]byte
// Write cells to page.
writePageNo(buf[:], parent.ChildPgno)
writeFlags(buf[:], PageTypeLeaf)
writeCellN(buf[:], len(group))
offset := dataOffset(len(group))
x := 0
for j, cell := range group {
writeLeafCell(buf[:], j, offset, cell)
offset += align8(cell.Size())
x++
}
if err := c.tx.writePage(buf[:]); err != nil {
return err
}
parents = append(parents, parent)
}
// Free the source page once we've finished with it if it is on heap.
if isHeap {
freePage(leafPage)
}
// TODO(BBJ): Update page in buffer & cursor stack.
// If this is not a split then exit now.
if len(groups) == 1 {
return nil
}
// Initialize a new root if we are currently the root page.
if c.stack.top == 0 {
assert(newRoot) // leaf write must be root when stack at root
return c.writeRoot(origPgno, parents)
}
assert(!newRoot) // leaf write must NOT be root when stack not at root
// Otherwise update existing parent.
return c.putBranchCells(c.stack.top-1, parents)
}
// putLeafCellFast quickly insert or updates a cell on a leaf page.
// It works by shifting bytes around instead of deserializing. This must not overflow.
func (c *Cursor) putLeafCellFast(in leafCell, isInsert bool) (err error) {
elem := &c.stack.elems[c.stack.top]
src, isHeap, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
srcCellN := readCellN(src)
// Determine the cell count of the new page.
dstCellN := srcCellN
if isInsert {
dstCellN++
}
// Write page header.
dst := allocPage()
writePageNo(dst, readPageNo(src))
writeFlags(dst, PageTypeLeaf)
writeCellN(dst, dstCellN)
// Copy data before index.
offset := dataOffset(dstCellN)
if elem.index > 0 {
srcStart := dataOffset(srcCellN)
srcEnd := readCellEndingOffset(src, elem.index-1)
shiftN := offset - srcStart
copy(dst[offset:], src[srcStart:srcEnd])
offset += align8(srcEnd - srcStart)
// Rewrite initial index slots by position moved.
for i := 0; i < elem.index; i++ {
writeCellOffset(dst, i, readCellOffset(src, i)+shiftN)
}
}
// Insert new row.
writeLeafCell(dst[:], elem.index, offset, in)
offset += align8(in.Size())
// Copy data after inserted element.
if (isInsert && elem.index < srcCellN) || (!isInsert && elem.index < srcCellN-1) {
var srcStart int
if isInsert {
srcStart = readCellOffset(src, elem.index)
} else {
srcStart = readCellOffset(src, elem.index+1)
}
srcEnd := readCellEndingOffset(src, srcCellN-1)
copy(dst[offset:], src[srcStart:srcEnd])
// Rewrite ending index slots by position moved.
shiftN := offset - srcStart
for i := elem.index + 1; i < dstCellN; i++ {
srci := i
if isInsert {
srci--
}
writeCellOffset(dst, i, readCellOffset(src, srci)+shiftN)
}
}
// Write new page to dirty page cache.
if err := c.tx.writePage(dst); err != nil {
return err
}
// Free page if on heap.
if isHeap {
freePage(src)
}
return nil
}
// deleteLeafCell removes a cell from the currently positioned page & index.
// If the removal causes the leaf page to have no more elements then its entry
// is removed from the parent. If the removal changes the first entry in the
// leaf page then the entry will be updated in the parent branch page.
func (c *Cursor) deleteLeafCell(key uint64) (err error) {
elem := &c.stack.elems[c.stack.top]
leafPage, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
cells := readLeafCells(leafPage, c.leafCells[:])
oldPageKey := cells[0].Key
cell := readLeafCell(leafPage, elem.index)
if cell.Type == ContainerTypeBitmapPtr {
if err := c.tx.freePgno(toPgno(cell.Data)); err != nil {
return err
}
}
// If no more cells exist and we have a parent, remove from parent.
if c.stack.top > 0 && len(cells) == 1 {
if err := c.tx.freePgno(elem.pgno); err != nil {
return err
}
return c.deleteBranchCell(c.stack.top-1, cells[0].Key)
}
// Remove matching cell from list.
copy(cells[elem.index:], cells[elem.index+1:])
cells[len(cells)-1] = leafCell{}
cells = cells[:len(cells)-1]
// Write cells to page.
buf := allocPage()
writePageNo(buf[:], elem.pgno)
writeFlags(buf[:], PageTypeLeaf)
writeCellN(buf[:], len(cells))
offset := dataOffset(len(cells))
for j, cell := range cells {
writeLeafCell(buf[:], j, offset, cell)
offset += align8(cell.Size())
}
if err := c.tx.writePage(buf[:]); err != nil {
return err
}
// Update the parent's reference key if it's changed.
if c.stack.top > 0 && oldPageKey != cells[0].Key {
return c.updateBranchCell(c.stack.top-1, cells[0].Key)
}
return nil
}
// putBranchCells updates a branch page with one or more cells. If the cells
// cause the branch page to split then a new entry will be created in the
// parent page. If this branch page is a root then a new root will be created.
func (c *Cursor) putBranchCells(stackIndex int, newCells []branchCell) (err error) {
elem := &c.stack.elems[stackIndex]
// Read branch page from disk. The current buffer is the leaf page.
page, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
cells := readBranchCells(page)
if len(cells) == 0 {
cells = make([]branchCell, 1)
}
// Update current cell & insert additional cells after it.
cells[elem.index] = newCells[0]
if len(newCells) > 1 {
cells = append(cells, make([]branchCell, len(newCells)-1)...)
copy(cells[elem.index+len(newCells):], cells[elem.index+1:])
copy(cells[elem.index+1:], newCells[1:])
}
// Split into multiple pages if page size is exceeded.
groups := [][]branchCell{cells}
if branchCellsPageSize(cells) > PageSize {
groups = splitBranchCells(cells)
}
// Write each group to a separate page.
var parents []branchCell
origPgno := readPageNo(page)
newRoot := len(groups) > 1 && stackIndex == 0
for i, group := range groups {
// First page should overwrite the original.
// Subsequent pages should allocate new pages.
parent := branchCell{LeftKey: group[0].LeftKey}
if i == 0 && !newRoot {
parent.ChildPgno = origPgno
} else {
if parent.ChildPgno, err = c.tx.allocatePgno(); err != nil {
return fmt.Errorf("cannot allocate branch: %w", err)
}
}
parents = append(parents, parent)
// Write cells to page.
var buf [PageSize]byte
writePageNo(buf[:], parents[i].ChildPgno)
writeFlags(buf[:], PageTypeBranch)
writeCellN(buf[:], len(group))
offset := dataOffset(len(group))
for j, cell := range group {
writeBranchCell(buf[:], j, offset, cell)
offset += align8(branchCellSize)
}
if err := c.tx.writePage(buf[:]); err != nil {
return err
}
}
// TODO(BBJ): Update page in buffer & cursor stack.
// If this is not a split, then exit now.
if len(groups) == 1 {
return nil
}
// Initialize a new root if we are currently the root page.
if stackIndex == 0 {
assert(newRoot) // branch write must be root when stack at root
return c.writeRoot(origPgno, parents)
}
assert(!newRoot) // branch write must NOT be root when stack at root
// Otherwise update existing parent.
return c.putBranchCells(stackIndex-1, parents)
}
// updateBranchCell updates the key for cell in the branch. Branch elements
// are fixed size so this cannot cause a page split.
func (c *Cursor) updateBranchCell(stackIndex int, newKey uint64) (err error) {
elem := &c.stack.elems[stackIndex]
// Read branch page from disk. The current buffer is the leaf page.
page, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
cells := readBranchCells(page)
oldPageKey := cells[0].LeftKey
// Update key in branch cell.
cells[elem.index].LeftKey = newKey
// Write cells to page.
var buf [PageSize]byte
writePageNo(buf[:], elem.pgno)
writeFlags(buf[:], PageTypeBranch)
writeCellN(buf[:], len(cells))
offset := dataOffset(len(cells))
for j, cell := range cells {
writeBranchCell(buf[:], j, offset, cell)
offset += align8(branchCellSize)
}
if err := c.tx.writePage(buf[:]); err != nil {
return err
}
if stackIndex > 0 && oldPageKey != cells[0].LeftKey {
return c.updateBranchCell(stackIndex-1, cells[0].LeftKey)
}
return nil
}
// deleteBranchCell removes a cell from a branch page. If no more elements
// exist in a branch page then it is removed. If an empty branch page is the
// root page then it is converted to an empty leaf page as branch pages are not
// allowed to be empty.
func (c *Cursor) deleteBranchCell(stackIndex int, key uint64) (err error) {
elem := &c.stack.elems[stackIndex]
// Read branch page from disk. The current buffer is the leaf page.
page, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
cells := readBranchCells(page)
oldPageKey := cells[0].LeftKey
// Remove cell from branch.
copy(cells[elem.index:], cells[elem.index+1:])
cells[len(cells)-1] = branchCell{}
cells = cells[:len(cells)-1]
// Branches are not allowed to have zero element so we must remove the page
// or, in the case of the root page, convert to a leaf page.
if len(cells) == 0 {
// If this is the root page, convert to leaf page.
if stackIndex == 0 {
var buf [PageSize]byte
writePageNo(buf[:], elem.pgno)
writeFlags(buf[:], PageTypeLeaf)
writeCellN(buf[:], len(cells))
return c.tx.writePage(buf[:])
}
// If this is a non-root page, free and remove from parent.
if err := c.tx.freePgno(elem.pgno); err != nil {
return err
}
return c.deleteBranchCell(stackIndex-1, oldPageKey)
}
// If the root only has one node, replace it with its child.
if stackIndex == 0 && len(cells) == 1 {
target, _, err := c.tx.readPage(cells[0].ChildPgno)
if err != nil {
return err
}
buf := allocPage()
copy(buf, target)
writePageNo(buf[:], elem.pgno)
if err := c.tx.freePgno(cells[0].ChildPgno); err != nil {
return err
}
return c.tx.writePage(buf[:])
}
// Write cells to page.
var buf [PageSize]byte
writePageNo(buf[:], elem.pgno)
writeFlags(buf[:], PageTypeBranch)
writeCellN(buf[:], len(cells))
offset := dataOffset(len(cells))
for j, cell := range cells {
writeBranchCell(buf[:], j, offset, cell)
offset += align8(branchCellSize)
}
assert(readCellN(buf[:]) > 0) // must have at least one cell
if err := c.tx.writePage(buf[:]); err != nil {
return err
}
if stackIndex > 0 && len(cells) > 0 && oldPageKey != cells[0].LeftKey {
return c.updateBranchCell(stackIndex-1, cells[0].LeftKey)
}
return nil
}
// writeRoot writes a new branch page at the root with the given cells.
// The root of a b-tree should always stay the same, even in the case of a
// split. This allows us to not rewrite the root record for the b-tree.
func (c *Cursor) writeRoot(pgno uint32, cells []branchCell) error {
var buf [PageSize]byte
writePageNo(buf[:], pgno)
writeFlags(buf[:], PageTypeBranch)
writeCellN(buf[:], len(cells))
offset := dataOffset(len(cells))
for i := range cells {
writeBranchCell(buf[:], i, offset, cells[i])
offset += align8(branchCellSize)
}
return c.tx.writePage(buf[:])
}
// splitLeafCells splits cells into roughly equal parts. It's a naive
// implementation that splits cells whenever a page is 60% full.
func splitLeafCells(cells []leafCell) [][]leafCell {
slices := make([][]leafCell, 1, 2)
var dataSize int
for _, cell := range cells {
if cell.Type == ContainerTypeBitmap {
panic("no! all ContainerTypeBitmap should be ContainerTypeBitmapPtr by now")
}
// Determine number of cells on current slice & cell size.
cellN := len(slices[len(slices)-1])
sz := align8(leafCellHeaderSize + len(cell.Data))
// If there is at least one cell on the slice & we've exceeded
// half a page then create a new group of cells.
thresh := int(float64(PageSize) * globalBranchFillPct)
if cellN != 0 && (dataOffset(cellN+1)+dataSize+sz) > thresh {
slices, dataSize = append(slices, nil), 0
} else if cellN != 0 && cell.Type == ContainerTypeArray && cell.ElemN > ArrayMaxSize {
slices, dataSize = append(slices, nil), 0
sz = PageSize
}
// Append to current slice & increase total cell data size.
slices[len(slices)-1] = append(slices[len(slices)-1], cell)
dataSize += sz
}
return slices
}
var globalBranchFillPct = 0.60
// splitBranchCells splits cells into roughly equal parts. It's a naive
// implementation that splits cells whenever a page is 60% full.
func splitBranchCells(cells []branchCell) [][]branchCell {
slices := make([][]branchCell, 1, 2)
var dataSize int
for _, cell := range cells {
// Determine number of cells on current slice & cell size.
cellN := len(slices[len(slices)-1])
sz := align8(branchCellSize)
// If there is at least one cell on the slice & we've exceeded
// half a page then create a new group of cells.
thresh := int(float64(PageSize) * globalBranchFillPct)
if cellN != 0 && (dataOffset(cellN+1)+dataSize+sz) > thresh {
slices, dataSize = append(slices, nil), 0
}
// Append to current slice & increase total cell data size.
slices[len(slices)-1] = append(slices[len(slices)-1], cell)
dataSize += sz
}
return slices
}
// pageKeyAt returns the key at the given index of the page.
func pageKeyAt(page []byte, index int) uint64 {
offset := readCellOffset(page, index)
return *(*uint64)(unsafe.Pointer(&page[offset]))
}
// First moves to the first element of the btree. Returns io.EOF if no elements
// exist. The first call to Next() will not move the position but subsequent
// calls will move the position forward until it reaches the end and return io.EOF.
func (c *Cursor) First() error {
c.buffered = true
for c.stack.top = 0; ; c.stack.top++ {
elem := &c.stack.elems[c.stack.top]
buf, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
switch typ := readFlags(buf); typ {
case PageTypeBranch:
elem.index = 0
if n := readCellN(buf); elem.index >= n { // branch cell index must less than cell count
return fmt.Errorf("branch cell index out of range: pgno=%d i=%d n=%d", elem.pgno, elem.index, n)
}
// Read cell pgno into the next stack level.
cell := readBranchCell(buf, elem.index)
c.stack.elems[c.stack.top+1] = stackElem{
pgno: cell.ChildPgno,
key: cell.LeftKey,
}
case PageTypeLeaf:
elem.index = 0
if readCellN(buf) == 0 {
return io.EOF // root leaf with no elements
}
return nil
default:
return fmt.Errorf("rbf.Cursor.First(): invalid page type: pgno=%d type=%d", elem.pgno, typ)
}
}
}
// Last moves to the last element of the btree.
func (c *Cursor) Last() error {
// c.stack.elems[0].pgno = c.root
c.buffered = true
for c.stack.top = 0; ; c.stack.top++ {
elem := &c.stack.elems[c.stack.top]
buf, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
switch typ := readFlags(buf); typ {
case PageTypeBranch:
elem.index = readCellN(buf) - 1
// Read cell pgno into the next stack level.
cell := readBranchCell(buf, elem.index)
c.stack.elems[c.stack.top+1] = stackElem{
pgno: cell.ChildPgno,
key: cell.LeftKey,
}
case PageTypeLeaf:
elem.index = readCellN(buf) - 1
if readCellN(buf) == 0 {
return io.EOF // root leaf with no elements
}
return nil
default:
return fmt.Errorf("rbf.Cursor.Last(): invalid page type: pgno=%d type=%d", elem.pgno, typ)
}
}
}
// Seek moves to the specified container of the btree.
// If the container does not exist then it moves to the next container after the key.
// TODO: what happens if there are no more containers?!?!
func (c *Cursor) Seek(key uint64) (exact bool, err error) {
// c.stack.elems[0].pgno = c.bitmap.root
c.buffered = true
for c.stack.top = 0; ; c.stack.top++ {
elem := &c.stack.elems[c.stack.top]
assert(elem.pgno != 0) // cursor should never point to page zero (meta)
buf, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return false, err
}
switch typ := readFlags(buf); typ {
case PageTypeBranch:
n := readCellN(buf)
index, xact := search(n, func(i int) int {
if v := readBranchCellKey(buf, i); key == v {
return 0
} else if key < v {
return -1
}
return 1
})
//if not found (xact) the cell
if !xact && index > 0 {
index--
}
elem.index = index
// Read cell pgno into the next stack level.
cell := readBranchCell(buf, elem.index)
c.stack.elems[c.stack.top+1] = stackElem{
pgno: cell.ChildPgno,
key: cell.LeftKey,
}
case PageTypeLeaf:
n := readCellN(buf)
index, xact := search(n, func(i int) int {
if v := readLeafCellKey(buf, i); key == v {
return 0
} else if key < v {
return -1
}
return 1
})
elem.index = index
return xact, nil
default:
return false, fmt.Errorf("rbf.Cursor.Seek(): invalid page type: pgno=%d type=%d", elem.pgno, typ)
}
}
}
// Next moves to the next element of the btree. Returns EOF if no more elements exist.
func (c *Cursor) Next() error {
elem := &c.stack.elems[c.stack.top]
leafPage, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
if c.buffered {
c.buffered = false
// Move to next available element if we are past the last cell in the page.
if elem.index >= readCellN(leafPage) {
return c.goNextPage()
}
return nil
}
// Move forward to the next leaf element if available.
if elem.index < readCellN(leafPage)-1 {
elem.index++
return nil
}
return c.goNextPage()
}
// Prev moves to the previous element of the btree. Returns io.EOF if at the
// beginning of the b-tree.
func (c *Cursor) Prev() error {
if c.buffered {
c.buffered = false
return nil
}
// Move forward to the next leaf element if available.
if elem := &c.stack.elems[c.stack.top]; elem.index > 0 {
elem.index--
return nil
}
// Move up the stack until we can move forward one element.
for c.stack.top--; c.stack.top >= 0; c.stack.top-- {
elem := &c.stack.elems[c.stack.top]
if elem.index > 0 {
elem.index--
break
}
}
// No more elements, return EOF.
if c.stack.top == -1 {
c.stack.top = 0
return io.EOF
}
// Traverse back down the stack to find the first element in each page.
for ; ; c.stack.top++ {
elem := &c.stack.elems[c.stack.top]
buf, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
switch typ := readFlags(buf); typ {
case PageTypeBranch:
cell := readBranchCell(buf, elem.index)
c.stack.elems[c.stack.top+1] = stackElem{
pgno: cell.ChildPgno,
key: cell.LeftKey,
}
case PageTypeLeaf:
elem.index = readCellN(buf) - 1
return nil
default:
return fmt.Errorf("rbf.Cursor.Prev(): invalid page type: pgno=%d type=%d", elem.pgno, typ)
}
}
}
// Key returns the key for the container the cursor is currently pointing to.
func (c *Cursor) Key() uint64 {
elem := &c.stack.elems[c.stack.top]
leafPage, _, _ := c.tx.readPage(elem.pgno)
if readCellN(leafPage[:]) == 0 {
return 0
}
return readLeafCellKey(leafPage, elem.index)
}
// Values returns the values for the container the cursor is currently pointing to.
func (c *Cursor) Values() []uint16 {
elem := &c.stack.elems[c.stack.top]
leafPage, _, _ := c.tx.readPage(elem.pgno)
if readCellN(leafPage[:]) == 0 {
return nil
}
cell := readLeafCell(leafPage, elem.index)
return cell.Values(c.tx)
}
// stackElem represents a single element on the cursor stack.
type stackElem struct {
pgno uint32 // current page number
index int // cell index
key uint64 // element key
}
func (se *stackElem) equal(se2 *stackElem) bool {
if se.pgno != se2.pgno {
return false
}
if se.index != se2.index {
return false
}
if se.key != se2.key {
return false
}
return true
}
func (se *stackElem) String() string {
return fmt.Sprintf("stackElem{pgno:%v index:%v key:%v}", int(se.pgno) /*,tx.pageTypeDesc(se.pgno)*/, se.index, int(se.key))
}
func (se *stackElem) clear() {
se.pgno = 0
se.index = 0
se.key = 0
}
// suppress unused warnings - these can be useful for debugging
var _ = (&stackElem{}).clear
var _ = (&stackElem{}).String
var _ = (&stackElem{}).equal
// goNextPage moves up the stack until an element has additional elements
// available after the current position. It then traverses back down the stack
// to position at the next leaf page element.
func (c *Cursor) goNextPage() error {
for c.stack.top--; c.stack.top >= 0; c.stack.top-- {
elem := &c.stack.elems[c.stack.top]
if buf, _, err := c.tx.readPage(elem.pgno); err != nil {
return err
} else if n := readCellN(buf); elem.index+1 < n {
elem.index++
break
}
}
// No more elements, return EOF.
if c.stack.top == -1 {
c.stack.top = 0
return io.EOF
}
// Traverse back down the stack to find the first element in each page.
for ; ; c.stack.top++ {
elem := &c.stack.elems[c.stack.top]
buf, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return err
}
switch typ := readFlags(buf); typ {
case PageTypeBranch:
cell := readBranchCell(buf, elem.index)
c.stack.elems[c.stack.top+1] = stackElem{
pgno: cell.ChildPgno,
key: cell.LeftKey,
}
case PageTypeLeaf:
elem.index = 0
return nil
default:
return fmt.Errorf("rbf.Cursor.Next(): invalid page type: pgno=%d type=%d", elem.pgno, typ)
}
}
}
func ConvertToLeafArgs(key uint64, c *roaring.Container) (result leafCell) {
result.Key = key
result.BitN = int(c.N())
result.Type = ContainerTypeNone
if c.N() == 0 {
return
}
switch roaring.ContainerType(c) {
case 1: //array
a := roaring.AsArray(c)
if len(a) > ArrayMaxSize {
roaring.ConvertArrayToBitmap(c)
result.Type = ContainerTypeBitmap
result.Data = fromArray64(roaring.AsBitmap(c))
// result.ElemN is 0 or undefined for bitmap
return
}
result.Type = ContainerTypeArray
result.Data = fromArray16(a)
result.ElemN = int(c.N())
return
case 2: //bitmap
result.Type = ContainerTypeBitmap
result.Data = fromArray64(roaring.AsBitmap(c))
// result.ElemN is 0 or undefined for bitmap
return
case 3: //run
r := roaring.AsRuns(c)
if len(r) > RLEMaxSize {
roaring.ConvertRunToBitmap(c)
result.Type = ContainerTypeBitmap
result.Data = fromArray64(roaring.AsBitmap(c))
// result.ElemN is 0 or undefined for bitmap
return
}
result.ElemN = len(r) // ElemN is the number of runs
result.Type = ContainerTypeRLE
result.Data = fromInterval16(r)
return
}
return
}
func (c *Cursor) merge(key uint64, data *roaring.Container) (bool, error) {
elem := &c.stack.elems[c.stack.top]
leafPage, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return false, err
}
cell := readLeafCell(leafPage, elem.index)
var container *roaring.Container
switch cell.Type {
case ContainerTypeArray:
d := toArray16(cell.Data)
container = roaring.NewContainerArray(d)
case ContainerTypeBitmapPtr:
_, d, err := c.tx.leafCellBitmap(toPgno(cell.Data))
if err != nil {
return false, errors.Wrap(err, "cursor.merge")
}
container = roaring.NewContainerBitmap(-1, d)
case ContainerTypeRLE:
d := toInterval16(cell.Data)
container = roaring.NewContainerRun(d)
}
res := roaring.Union(data, container)
if res.N() != data.N() {
leaf := ConvertToLeafArgs(key, res)
err := c.putLeafCell(leaf)
return true, err
}
return false, nil
}
func (c *Cursor) AddRoaring(bm *roaring.Bitmap) (changed bool, err error) {
itr, _ := bm.Containers.Iterator(0)
for itr.Next() {
hi, cont := itr.Value()
leaf := ConvertToLeafArgs(hi, cont)
if leaf.BitN == 0 {
continue
}
// Move cursor to the key of the container.
// Insert new container if it doesn't exist.
if exact, err := c.Seek(hi); err != nil {
return false, err
} else if !exact {
err = c.putLeafCell(leaf)
if err != nil {
return false, err
}
changed = true
continue
}
// If the container exists and bit is not set then update the page.
u, err := c.merge(hi, cont)
if err != nil {
return false, err
}
if u {
changed = true
}
}
return changed, nil
}
func (c *Cursor) RemoveRoaring(bm *roaring.Bitmap) (changed bool, err error) {
itr, _ := bm.Containers.Iterator(0)
for itr.Next() {
hi, cont := itr.Value()
if cont.N() == 0 {
continue
}
// Move cursor to the key of the container.
// Insert new container if it doesn't exist.
if exact, err := c.Seek(hi); err != nil {
return false, err
} else if exact {
f, err := c.difference(hi, cont)
if err != nil {
return f, err
}
if f {
changed = true
}
}
}
return
}
func (c *Cursor) difference(key uint64, data *roaring.Container) (bool, error) {
elem := &c.stack.elems[c.stack.top]
leafPage, _, err := c.tx.readPage(elem.pgno)
if err != nil {
return false, err
}
cell := readLeafCell(leafPage, elem.index)
var container *roaring.Container
switch cell.Type {
case ContainerTypeArray:
d := toArray16(cell.Data)
container = roaring.NewContainerArray(d)
case ContainerTypeBitmapPtr:
_, d, err := c.tx.leafCellBitmap(toPgno(cell.Data))
if err != nil {
return false, errors.Wrap(err, "cursor.difference")
}
container = roaring.NewContainerBitmap(-1, d)
case ContainerTypeRLE:
d := toInterval16(cell.Data)
container = roaring.NewContainerRun(d)
}
res := roaring.Difference(container, data)
if res.N() == 0 {
return true, c.deleteLeafCell(cell.Key)
}
if res.N() != container.N() {
leaf := ConvertToLeafArgs(key, res)
err := c.putLeafCell(leaf)
return true, err
}
return false, nil
}
// Close closes a cursor out, invalidating it for future use, and puts it
// back in a pool to reduce allocations. Contrast with unpooledClose, which
// you probably shouldn't use.
func (c *Cursor) Close() {
assert(c != nil)
c.unpooledClose()
cursorSyncPool.Put(c)
}
// unpooledClose is the part of a close operation which does not put the
// cursor back in the pool. It is useful only for cases where we're reusing
// a cursor, such as a Db's freelistCursor.
func (c *Cursor) unpooledClose() {
c.tx = nil
}
func keysFromParents(parents []branchCell) (ckeys []int) {
for _, par := range parents {
ckeys = append(ckeys, int(par.LeftKey))
}
return
}
var _ = (&Cursor{}).showCursorStack
func (c *Cursor) showCursorStack() (r string) {
r = fmt.Sprintf("top = %v\n", c.stack.top)
for i := 0; i <= c.stack.top; i++ {
r += fmt.Sprintf(" [%02v] %v\n", i, c.stack.elems[i].String())
}
return
}