feat(btree): add BTreeCorruptionException and BPlusTree

src/main/kotlin/kestrel/store/btree/BPlusTree.kt

@@ -0,0 +1,153 @@
+package kestrel.store.btree
+
+import kestrel.KestrelException
+import kestrel.store.codec.KeyCodec
+import kestrel.store.page.Page
+import kestrel.store.page.PageStore
+
+class BTreeCorruptionException(message: String) : KestrelException(message)
+
+/**
+ * A disk-backed B+Tree mapping order-preserving byte keys (see [KeyCodec]) to
+ * row pointers. Nodes live on [Page]s fetched through a [BufferPool]; leaves
+ * are linked for ordered range scans. Insert splits an overflowing node and
+ * propagates a promoted separator key upward, growing the root when needed.
+ * Delete removes the entry and, when a node drops below
+ * [UNDERFLOW_THRESHOLD], redistributes an entry from a sibling with spare
+ * capacity or merges with one, shrinking the root when it is left with a
+ * single child.
+ */
+class BPlusTree private constructor(private val bufferPool: PageStore, rootPageId: Long) {
+    companion object {
+        private const val UNDERFLOW_THRESHOLD = Page.PAGE_SIZE / 4
+
+        fun create(bufferPool: PageStore): BPlusTree {
+            val page = bufferPool.newPage()
+            NodeCodec.writeTo(page, BTreeNode.Leaf(page.id, mutableListOf(), mutableListOf(), NodeCodec.NO_PAGE))
+            bufferPool.unpinPage(page.id, true)
+            return BPlusTree(bufferPool, page.id)
+        }
+
+        fun open(bufferPool: PageStore, rootPageId: Long): BPlusTree = BPlusTree(bufferPool, rootPageId)
+    }
+
+    var rootPageId: Long = rootPageId
+        private set
+
+    fun search(key: ByteArray): Long? {
+        var pageId = rootPageId
+        while (true) {
+            val (_, node) = fetch(pageId)
+            when (node) {
+                is BTreeNode.Internal -> {
+                    val nextId = node.children[childIndexFor(node.keys, key)]
+                    bufferPool.unpinPage(pageId, false)
+                    pageId = nextId
+                }
+                is BTreeNode.Leaf -> {
+                    val idx = lowerBound(node.keys, key)
+                    val result = if (idx < node.keys.size && KeyCodec.compare(node.keys[idx], key) == 0) {
+                        node.values[idx]
+                    } else {
+                        null
+                    }
+                    bufferPool.unpinPage(pageId, false)
+                    return result
+                }
+            }
+        }
+    }
+
+    fun insert(key: ByteArray, value: Long) {
+        val split = insertRecursive(rootPageId, key, value)
+        if (split != null) {
+            val newRoot = bufferPool.newPage()
+            NodeCodec.writeTo(
+                newRoot,
+                BTreeNode.Internal(newRoot.id, mutableListOf(split.promotedKey), mutableListOf(rootPageId, split.newPageId))
+            )
+            bufferPool.unpinPage(newRoot.id, true)
+            rootPageId = newRoot.id
+        }
+    }
+
+    fun delete(key: ByteArray): Boolean {
+        val found = deleteRecursive(rootPageId, key)
+        val (_, rootNode) = fetch(rootPageId)
+        if (rootNode is BTreeNode.Internal && rootNode.keys.isEmpty()) {
+            val onlyChild = rootNode.children[0]
+            bufferPool.unpinPage(rootPageId, false)
+            bufferPool.deletePage(rootPageId)
+            rootPageId = onlyChild
+        } else {
+            bufferPool.unpinPage(rootPageId, false)
+        }
+        return found
+    }
+
+    /** Every page belonging to this tree: every internal node and every leaf, reached by a full traversal from the root. */
+    fun allPageIds(): List<Long> {
+        val ids = mutableListOf<Long>()
+        val pending = ArrayDeque<Long>()
+        pending.add(rootPageId)
+        while (pending.isNotEmpty()) {
+            val pageId = pending.removeFirst()
+            ids.add(pageId)
+            val (_, node) = fetch(pageId)
+            bufferPool.unpinPage(pageId, false)
+            if (node is BTreeNode.Internal) pending.addAll(node.children)
+        }
+        return ids
+    }
+
+    /** Ordered scan over `[from, to]` (or `[from, to)` when [toInclusive] is false). Null bounds are open. */
+    fun rangeScan(from: ByteArray? = null, to: ByteArray? = null, toInclusive: Boolean = false): Iterator<Pair<ByteArray, Long>> {
+        val startLeafId = leftmostLeafFor(from)
+        return iterator {
+            var currentLeafId = startLeafId
+            var firstLeaf = true
+            outer@ while (currentLeafId != NodeCodec.NO_PAGE) {
+                val (_, leaf) = fetch(currentLeafId)
+                leaf as BTreeNode.Leaf
+                var index = if (firstLeaf && from != null) lowerBound(leaf.keys, from) else 0
+                firstLeaf = false
+                while (index < leaf.keys.size) {
+                    val currentKey = leaf.keys[index]
+                    if (to != null) {
+                        val cmp = KeyCodec.compare(currentKey, to)
+                        if (cmp > 0 || (cmp == 0 && !toInclusive)) {
+                            bufferPool.unpinPage(currentLeafId, false)
+                            break@outer
+                        }
+                    }
+                    yield(currentKey to leaf.values[index])
+                    index++
+                }
+                val next = leaf.nextLeaf
+                bufferPool.unpinPage(currentLeafId, false)
+                currentLeafId = next
+            }
+        }
+    }
+
+    private fun leftmostLeafFor(from: ByteArray?): Long {
+        var pageId = rootPageId
+        while (true) {
+            val (_, node) = fetch(pageId)
+            when (node) {
+                is BTreeNode.Internal -> {
+                    val idx = if (from == null) 0 else childIndexFor(node.keys, from)
+                    val nextId = node.children[idx]
+                    bufferPool.unpinPage(pageId, false)
+                    pageId = nextId
+                }
+                is BTreeNode.Leaf -> {
+                    bufferPool.unpinPage(pageId, false)
+                    return pageId
+                }
+            }
+        }
+    }
+
+    private data class SplitResult(val promotedKey: ByteArray, val newPageId: Long)
+