// 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< 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<= 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<= 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) parents := make([]branchCell, 0, len(groups)) 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. parents := make([]branchCell, 0, len(groups)) 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 }