bnode.go

package main

import (
	"encoding/binary"
	"bytes"
)

// 4 byte header (2byte for type + 2byte for key counts)
const HEADER = 4

// Each node is exactly 4KB — matching the OS memory page size.
// disk I/O and virtual memory both operate in 4KB pages, so one node = one page = one disk read.
const BTREE_PAGE_SIZE = 4096

const BTREE_MAX_KEY_SIZE = 1000
const BTREE_MAX_VAL_SIZE = 3000

type BNode []byte

type BTree struct {
	root uint64
	get func(uint64) []byte  // read a page from disk
	new func([]byte) uint64  // allocate a new page
	del func(uint64)         // deallocate a page
}

const (
	BNODE_NODE = 1    // keys + child pointers
	BNODE_LEAF = 2    // keys + values
)

// [0:2]               btype
// [2:4]               nkeys
// [4 : 4+8*nkeys]     child pointers  (8 bytes each)
// [above : +2*nkeys]  offset table    (2 bytes each)   offset-table: array that stores "where does each KV pair start?"
// [rest]              KV pairs        (variable size)

// --- Header Getters ---
func (node BNode) btype() uint16 {
	return binary.LittleEndian.Uint16(node[0:2])
}

func (node BNode) nkeys() uint16 {
	return binary.LittleEndian.Uint16(node[2:4])
}

// --- Header Setters ---
func (node BNode) setHeader(btype uint16, nkeys uint16) {
	binary.LittleEndian.PutUint16(node[0:2], btype)
	binary.LittleEndian.PutUint16(node[2:4], nkeys)
}

// getPtr retrieves the child pointer at a specific index
func (node BNode) getPtr(idx uint16) uint64 {
	if idx >= node.nkeys() {
		panic("index out of bounds")
	}
	// Pointers start exactly after the 4-byte header
	pos := HEADER + 8*idx
	return binary.LittleEndian.Uint64(node[pos:])
}

// setPtr writes a child pointer to a specific index
func (node BNode) setPtr(idx uint16, val uint64) {
	pos := HEADER + 8*idx
	binary.LittleEndian.PutUint64(node[pos:], val)
}

//offsetPos returns the absolute byte position of the offset for a given index
func offsetPos(node BNode, idx uint16) uint16 {
	if idx < 1 || idx > node.nkeys() {
		panic("index out of bounds")
	}
	// Header (4) + Pointers (8 * nkeys) + Offsets (2 * (idx - 1))
	return HEADER + 8*node.nkeys() + 2*(idx-1)
}

// getOffset returns the relative offset for the KV pair at idx
func (node BNode) getOffset(idx uint16) uint16 {
	if idx == 0 {
		return 0 // The first KV pair always starts at offset 0
	}
	return binary.LittleEndian.Uint16(node[offsetPos(node, idx):])
}

// setOffset writes the relative offset for the KV pair at idx
func (node BNode) setOffset(idx uint16, offset uint16) {
	binary.LittleEndian.PutUint16(node[offsetPos(node, idx):], offset)
}

// kvPos calculates the absolute byte position of the KV pair at idx
func (node BNode) kvPos(idx uint16) uint16 {
	if idx > node.nkeys() {
		panic("index out of bounds")
	}
	// Header (4) + Pointers (8 * nkeys) + Offsets (2 * nkeys) + Relative Offset
	return HEADER + 8*node.nkeys() + 2*node.nkeys() + node.getOffset(idx)
}

// getKey extracts the byte slice representing the key at idx
func (node BNode) getKey(idx uint16) []byte {
	if idx >= node.nkeys() {
		panic("index out of bounds")
	}
	pos := node.kvPos(idx)
	klen := binary.LittleEndian.Uint16(node[pos:])
	return node[pos+4: pos+4+klen]
}

// getVal extracts the byte slice representing the value at idx
func (node BNode) getVal(idx uint16) []byte {
	if idx >= node.nkeys() {
		panic("index out of bounds")
	}
	pos := node.kvPos(idx)
	klen := binary.LittleEndian.Uint16(node[pos+0:])
	vlen := binary.LittleEndian.Uint16(node[pos+2:])
	return node[pos+4+klen : pos+4+klen+vlen]
}

// nbytes returns the total size of the node in bytes (used space)
func (node BNode) nbytes() uint16 {
	return node.kvPos(node.nkeys())
}

// nodeLookupLE returns the index of the first child node whose range intersects the key.
// It finds an index `i` such that node.getKey(i) <= key.
func nodeLookupLE(node BNode, key []byte) uint16 {
	nkeys := node.nkeys()
	found := uint16(0)

	for i := uint16(1) ; i < nkeys ; i++ {
		cmp := bytes.Compare(node.getKey(i), key)
		if cmp <= 0 {
			found = i
		}
		if cmp >= 0 {
			break
		}
	}
	return  found
}

// nodeAppendKV writes a new pointer, key, and value into the new node at idx
func nodeAppendKV(new BNode, idx uint16, ptr uint64, key []byte, val []byte) {
	new.setPtr(idx, ptr)
	pos := new.kvPos(idx)
	binary.LittleEndian.PutUint16(new[pos+0:], uint16(len(key)))
	binary.LittleEndian.PutUint16(new[pos+2:], uint16(len(val)))
	copy(new[pos+4:], key)
	copy(new[pos+4+uint16(len(key)):], val)
	new.setOffset(idx+1, new.getOffset(idx)+4+uint16(len(key)+len(val)))
}

// nodeAppendRange copies multiple KVs from an old node to a new node
func nodeAppendRange(new BNode, old BNode, dstNew uint16, srcOld uint16, n uint16) {
	for i := uint16(0) ; i < n ; i++ {
		// Extract from old
		ptr := old.getPtr(srcOld + i)
		key := old.getKey(srcOld + i)
		val := old.getVal(srcOld + i)

		// Write to new
		nodeAppendKV(new, dstNew+i, ptr, key, val)
	}
}

// leafInsert adds a new key-value pair to a leaf node
func leafInsert(new BNode, old BNode, idx uint16, key []byte, val []byte) {
	new.setHeader(BNODE_LEAF, old.nkeys()+1)
	nodeAppendRange(new, old, 0, 0, idx)
	nodeAppendKV(new, idx, 0, key, val)
	nodeAppendRange(new, old, idx+1, idx, old.nkeys()-idx)
}

// leafUpdate replaces an existing key's value in a leaf node
func leafUpdate(new BNode, old BNode, idx uint16, key []byte, val []byte) {
	new.setHeader(BNODE_LEAF, old.nkeys())
	nodeAppendRange(new, old, 0, 0, idx)
	nodeAppendKV(new, idx, 0, key, val)
	nodeAppendRange(new, old, idx+1, idx+1, old.nkeys()-(idx+1))
}

// treeInsert recursively inserts a KV pair, returning a newly allocated node
func treeInsert(tree *BTree, node BNode, key []byte, val []byte) BNode {
	new := BNode(make([]byte, 2*BTREE_PAGE_SIZE))
	idx := nodeLookupLE(node, key)

	switch node.btype() {
	case BNODE_LEAF:
		if bytes.Equal(key, node.getKey(idx)) {
			leafUpdate(new, node, idx, key, val)
		} else {
			leafInsert(new, node, idx+1, key, val)
		}
	case BNODE_NODE:
		nodeInsert(tree, new, node, idx, key, val)
	default:
		panic("bad node")
	}
	return new
}

func nodeInsert(tree *BTree, new BNode, node BNode, idx uint16, key []byte, val []byte) {
	kptr := node.getPtr(idx)

	// Recursive insertion into the child node
	knode := treeInsert(tree, tree.get(kptr), key, val)

	// Split the result if it got too big
	nsplit, split := nodeSplit3(knode)

	// Deallocate the old child node
	tree.del(kptr)

	// Update the child links in the current internal node
	nodeReplaceKidN(tree, new, node, idx, split[:nsplit]...)
}

// nodeSplit2 splits a node into two halves. 
// It does not guarantee the halves are under 4KB, it just divides the keys.
func nodeSplit2(left BNode, right BNode, old BNode) {
	half := old.nkeys() / 2

	left.setHeader(old.btype(), half)
	right.setHeader(old.btype(), old.nkeys()-half)

	nodeAppendRange(left, old, 0, 0, half)
	nodeAppendRange(right, old, 0, half, old.nkeys()-half)
}

// nodeSplit3 guarantees that an oversized node is broken down into 
// chunks that strictly obey the BTREE_PAGE_SIZE limit.
func nodeSplit3(old BNode) (uint16, [3]BNode) {
	if old.nbytes() <= BTREE_PAGE_SIZE {
		old = old[:BTREE_PAGE_SIZE]
		return 1, [3]BNode{old}  // perfectly fits no split needed
	}

	left := BNode(make([]byte, 2*BTREE_PAGE_SIZE))
	right := BNode(make([]byte, BTREE_PAGE_SIZE))

	nodeSplit2(left, right, old)   // try splitting into two

	if left.nbytes() <= BTREE_PAGE_SIZE {
		left = left[:BTREE_PAGE_SIZE]
		return 2, [3]BNode{left, right} // two nodes was enough
	}

	leftleft := BNode(make([]byte, BTREE_PAGE_SIZE))
	middle := BNode(make([]byte, BTREE_PAGE_SIZE))

	nodeSplit2(leftleft, middle, left)

	return 3, [3]BNode{leftleft, middle, right}
}

// nodeReplaceKidN replaces a single child pointer with multiple new child pointers.
func nodeReplaceKidN(tree *BTree, new BNode, old BNode, idx uint16, kids ...BNode) {
	inc := uint16(len(kids))

	// The new internal node will have more keys than the old one
	new.setHeader(BNODE_NODE, old.nkeys()+inc-1)

	// Copy everything BEFORE the old pointer
	nodeAppendRange(new, old, 0, 0, idx)

	// Insert the new pointers for freshly split nodes
	for i, node := range kids {
		ptr := tree.new(node)
		nodeAppendKV(new, idx + uint16(i), ptr, node.getKey(0), nil)
	}

	// Copy everything AFTER the old pointer
	nodeAppendRange(new, old, idx+inc, idx+1, old.nkeys()-(idx+1))
}

// leafDelete removes a key from a leaf node by shifting the remaining keys
func leafDelete(new BNode, old BNode, idx uint16) {
	new.setHeader(BNODE_LEAF, old.nkeys()-1)
	nodeAppendRange(new, old, 0, 0, idx)
	nodeAppendRange(new, old, idx, idx+1, old.nkeys()-(idx+1))
}

// nodeMerge combines two smaller nodes into one perfectly sized node
func nodeMerge(new BNode, left BNode, right BNode) {
	new.setHeader(left.btype(), left.nkeys() + right.nkeys())
	nodeAppendRange(new, left, 0, 0, left.nkeys())
	nodeAppendRange(new, right, left.nkeys(), 0, right.nkeys())
}

// nodeReplace2Kid replaces two adjacent child pointers with a single new pointer
func nodeReplace2Kid(new BNode, old BNode, idx uint16, ptr uint64, key []byte) {
	new.setHeader(BNODE_NODE, old.nkeys()-1)
	// Copy everything before the two old pointers
	nodeAppendRange(new, old, 0, 0, idx)
	// Insert the single new merged pointer
	nodeAppendKV(new, idx, ptr, key, nil)
	// Copy everything after the two old pointers
	nodeAppendRange(new, old, idx+1, idx+2, old.nkeys()-(idx+2))
}

// shouldMerge determines if a node is empty enough to merge and finds a valid sibling
// Returns:
//  -1 if merging with the left sibling
//  +1 if merging with the right sibling
//   0 if no merge should happen
func shouldMerge(tree *BTree, node BNode, idx uint16, updated BNode) (int, BNode) {
	if updated.nbytes() > BTREE_PAGE_SIZE / 4 {
		return 0, BNode{}
	}
	if idx > 0 {
		sibling := BNode(tree.get(node.getPtr(idx-1)))
		merged := sibling.nbytes() + updated.nbytes() - HEADER
		if merged <= BTREE_PAGE_SIZE {
			return -1, sibling
		}
	}
	if idx+1 < node.nkeys() {
		sibling := BNode(tree.get(node.getPtr(idx+1)))
		merged := sibling.nbytes() + updated.nbytes() - HEADER
		if merged <= BTREE_PAGE_SIZE {
			return +1, sibling
		}
	}
	return 0, BNode{}
}

// treeDelete recursively finds and removes a key
func treeDelete(tree *BTree, node BNode, key []byte) BNode {
	idx := nodeLookupLE(node, key)

	switch node.btype() {
	case BNODE_LEAF:
		if !bytes.Equal(key, node.getKey(idx)) {
			return BNode{} // node not found
		}
		new := BNode(make([]byte, BTREE_PAGE_SIZE))
		leafDelete(new, node, idx)
		return new
	case BNODE_NODE:
		return nodeDelete(tree, node, idx, key)
	default:
		panic("bad node")
	}
}

// nodeDelete handles recursively deleting from an internal node and managing merges
func nodeDelete(tree *BTree, node BNode, idx uint16, key []byte) BNode {
	kptr := node.getPtr(idx)
	updated := treeDelete(tree, tree.get(kptr), key)

	if len(updated) == 0 {
		return BNode{}
	}

	tree.del(kptr)

	new := BNode(make([]byte, BTREE_PAGE_SIZE))
	mergeDir, sibling := shouldMerge(tree, node, idx, updated)

	switch {
	case mergeDir < 0: // Merge left
		merged := BNode(make([]byte, BTREE_PAGE_SIZE))
		nodeMerge(merged, sibling, updated)
		tree.del(node.getPtr(idx - 1)) // Deallocate old left sibling
		nodeReplace2Kid(new, node, idx-1, tree.new(merged), merged.getKey(0))
		
	case mergeDir > 0: // Merge right
		merged := BNode(make([]byte, BTREE_PAGE_SIZE))
		nodeMerge(merged, updated, sibling)
		tree.del(node.getPtr(idx + 1)) // Deallocate old right sibling
		nodeReplace2Kid(new, node, idx, tree.new(merged), merged.getKey(0))
		
	case mergeDir == 0 && updated.nkeys() == 0:
		// Edge case: Node is completely empty but has no siblings to merge with
		if node.nkeys() != 1 || idx != 0 {
			panic("Invalid state: empty node with no siblings but parent has multiple keys")
		}
		new.setHeader(BNODE_NODE, 0) // The parent becomes empty too
		
	case mergeDir == 0 && updated.nkeys() > 0: 
		// No merge needed, replace the single old pointer with the updated one
		nodeReplaceKidN(tree, new, node, idx, updated)
	}

	return new
}

// Insert adds a new key or updates an existing key.
func (tree *BTree) Insert(key []byte, val []byte) {
	if tree.root == 0 {
		root := BNode(make([]byte, BTREE_PAGE_SIZE))
		root.setHeader(BNODE_LEAF, 2)
		
		nodeAppendKV(root, 0, 0, nil, nil)
		nodeAppendKV(root, 1, 0, key, val)
		tree.root = tree.new(root)
		return
	}

	node := treeInsert(tree, tree.get(tree.root), key, val)
	nsplit, split := nodeSplit3(node)
	tree.del(tree.root)

	if nsplit > 1 {
		root := BNode(make([]byte, BTREE_PAGE_SIZE))
		root.setHeader(BNODE_NODE, nsplit)
		for i, knode := range split[:nsplit] {
			ptr, key := tree.new(knode), knode.getKey(0)
			nodeAppendKV(root, uint16(i), ptr, key, nil)
		}
		tree.root = tree.new(root)
	} else {
		tree.root = tree.new(split[0])
	}
}

// Delete removes a key and returns whether the key was actually found.
func (tree *BTree) Delete(key []byte) bool {
	if tree.root == 0 {
		return false
	}

	updated := treeDelete(tree, tree.get(tree.root), key)
	if len(updated) == 0 {
		return false // Key not found
	}

	tree.del(tree.root) // Deallocate the old root

	if updated.btype() == BNODE_NODE && updated.nkeys() == 1 {
		tree.root = updated.getPtr(0)
	} else {
		tree.root = tree.new(updated)
	}
	
	return true
}

// Get searches the tree for a key and returns the value if found.
func (tree *BTree) Get(key []byte) ([]byte, bool) {
	if tree.root == 0 {
		return nil, false
	}

	ptr := tree.root
	for {
		node := BNode(tree.get(ptr))
		idx := nodeLookupLE(node, key)

		switch node.btype() {
		case BNODE_LEAF:
			if bytes.Equal(key, node.getKey(idx)) {
				return node.getVal(idx), true
			}
			return nil, false // Key not in leaf
		case BNODE_NODE:
			ptr = node.getPtr(idx)
		default:
			panic("bad node type in Get")
		}
	}
}