iterator.go

package main

import "bytes"

const (
	CMP_GE = +3 // >=
	CMP_GT = +2 // >
	CMP_LT = -2 // <
	CMP_LE = -3 // <=
)

// BIter is a stateful cursor representing a specific position in the B-Tree.
type BIter struct {
	tree *BTree
	path []BNode  // Pointers to the nodes from root to current leaf
	pos  []uint16 // The exact index we are sitting at within each node
}

// Valid checks if the iterator has fallen off the edge of the tree.
func (iter *BIter) Valid() bool {
	return len(iter.path) > 0 && iter.pos[len(iter.pos)-1] < iter.path[len(iter.path)-1].nkeys()
}

// Deref extracts the Key and Value at the current iterator position.
func (iter *BIter) Deref() ([]byte, []byte) {
	node := iter.path[len(iter.path)-1]
	idx := iter.pos[len(iter.pos)-1]
	return node.getKey(idx), node.getVal(idx)
}

// SeekLE finds the closest position that is Less Than or Equal to the input key.
func (tree *BTree) SeekLE(key []byte) *BIter {
	iter := &BIter{tree: tree}
	for ptr := tree.root; ptr != 0; {
		node := BNode(tree.get(ptr))
		idx := nodeLookupLE(node, key)
		iter.path = append(iter.path, node)
		iter.pos = append(iter.pos, idx)
		
		if node.btype() == BNODE_NODE {
			ptr = node.getPtr(idx)
		} else {
			ptr = 0
		}
	}
	return iter
}

// cmpOK is a helper to verify if the current key satisfies the requested range condition.
func cmpOK(key []byte, cmp int, ref []byte) bool {
	r := bytes.Compare(key, ref)
	switch cmp {
	case CMP_GE:
		return r >= 0
	case CMP_GT:
		return r > 0
	case CMP_LT:
		return r < 0
	case CMP_LE:
		return r <= 0
	}
	panic("bad cmp operator")
}

// Seek locates a starting point in the tree based on the comparison operator.
func (tree *BTree) Seek(key []byte, cmp int) *BIter {
	iter := tree.SeekLE(key)
	if cmp != CMP_LE && iter.Valid() {
		curKey, _ := iter.Deref()
		if !cmpOK(curKey, cmp, key) {
			if cmp > 0 { 
				iter.Next()
			} else {
				iter.Prev()
			}
		}
	}
	return iter
}

// Next moves the cursor forward by 1 KV pair, carrying over node boundaries if necessary.
func (iter *BIter) Next() {
	iterNext(iter, len(iter.path)-1)
}

func iterNext(iter *BIter, level int) {
	if iter.pos[level]+1 < iter.path[level].nkeys() {
		iter.pos[level]++ 
	} else if level > 0 {
		iterNext(iter, level-1) 
	} else {
		iter.pos[len(iter.pos)-1]++
		return
	}
	
	if level+1 < len(iter.pos) {
		node := iter.path[level]
		kid := BNode(iter.tree.get(node.getPtr(iter.pos[level])))
		iter.path[level+1] = kid
		iter.pos[level+1] = 0
	}
}

// Prev moves the cursor backward by 1 KV pair.
func (iter *BIter) Prev() {
	iterPrev(iter, len(iter.path)-1)
}

func iterPrev(iter *BIter, level int) {
	if iter.pos[level] > 0 {
		iter.pos[level]--
	} else if level > 0 {
		iterPrev(iter, level-1)
	} else {
		iter.path = nil
		iter.pos = nil
		return
	}
	
	if level+1 < len(iter.pos) {
		node := iter.path[level]
		kid := BNode(iter.tree.get(node.getPtr(iter.pos[level])))
		iter.path[level+1] = kid
		iter.pos[level+1] = kid.nkeys() - 1
	}
}