feat(btree): add borrowFromRight, mergeChildren and BTreeNode

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

@@ -313,3 +313,80 @@ class BPlusTree private constructor(private val bufferPool: PageStore, rootPageI
         bufferPool.unpinPage(childNode.pageId, true)
     }
 
+    private fun borrowFromRight(parent: BTreeNode.Internal, idx: Int, rightPage: Page, rightNode: BTreeNode) {
+        val (childPage, childNode) = fetch(parent.children[idx])
+        when {
+            rightNode is BTreeNode.Leaf && childNode is BTreeNode.Leaf -> {
+                val borrowedKey = rightNode.keys.removeAt(0)
+                val borrowedValue = rightNode.values.removeAt(0)
+                childNode.keys.add(borrowedKey)
+                childNode.values.add(borrowedValue)
+                parent.keys[idx] = rightNode.keys[0]
+            }
+            rightNode is BTreeNode.Internal && childNode is BTreeNode.Internal -> {
+                val movedChild = rightNode.children.removeAt(0)
+                val separatorDown = rightNode.keys.removeAt(0)
+                childNode.children.add(movedChild)
+                childNode.keys.add(parent.keys[idx])
+                parent.keys[idx] = separatorDown
+            }
+            else -> throw BTreeCorruptionException("sibling type mismatch during redistribution")
+        }
+        NodeCodec.writeTo(rightPage, rightNode)
+        NodeCodec.writeTo(childPage, childNode)
+        bufferPool.unpinPage(childNode.pageId, true)
+    }
+
+    private fun mergeChildren(parent: BTreeNode.Internal, leftIdx: Int) {
+        val leftId = parent.children[leftIdx]
+        val rightId = parent.children[leftIdx + 1]
+        val (leftPage, leftNode) = fetch(leftId)
+        val (_, rightNode) = fetch(rightId)
+        when {
+            leftNode is BTreeNode.Leaf && rightNode is BTreeNode.Leaf -> {
+                leftNode.keys.addAll(rightNode.keys)
+                leftNode.values.addAll(rightNode.values)
+                leftNode.nextLeaf = rightNode.nextLeaf
+            }
+            leftNode is BTreeNode.Internal && rightNode is BTreeNode.Internal -> {
+                leftNode.keys.add(parent.keys[leftIdx])
+                leftNode.keys.addAll(rightNode.keys)
+                leftNode.children.addAll(rightNode.children)
+            }
+            else -> throw BTreeCorruptionException("sibling type mismatch during merge")
+        }
+        NodeCodec.writeTo(leftPage, leftNode)
+        bufferPool.unpinPage(leftId, true)
+        bufferPool.unpinPage(rightId, false)
+        bufferPool.deletePage(rightId)
+        parent.keys.removeAt(leftIdx)
+        parent.children.removeAt(leftIdx + 1)
+    }
+
+    private fun isUnderflowNode(node: BTreeNode): Boolean = NodeCodec.serialize(node).size < UNDERFLOW_THRESHOLD
+
+    private fun fetch(pageId: Long): Pair<Page, BTreeNode> {
+        val page = bufferPool.fetchPage(pageId)
+        return page to NodeCodec.readFrom(page)
+    }
+
+    private fun lowerBound(keys: List<ByteArray>, key: ByteArray): Int {
+        var lo = 0
+        var hi = keys.size
+        while (lo < hi) {
+            val mid = (lo + hi) ushr 1
+            if (KeyCodec.compare(keys[mid], key) < 0) lo = mid + 1 else hi = mid
+        }
+        return lo
+    }
+
+    private fun childIndexFor(keys: List<ByteArray>, key: ByteArray): Int {
+        var lo = 0
+        var hi = keys.size
+        while (lo < hi) {
+            val mid = (lo + hi) ushr 1
+            if (KeyCodec.compare(keys[mid], key) <= 0) lo = mid + 1 else hi = mid
+        }
+        return lo
+    }
+}

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

@@ -0,0 +1,107 @@
+package kestrel.store.btree
+
+import kestrel.store.codec.ByteReader
+import kestrel.store.codec.ByteWriter
+import kestrel.store.codec.CodecException
+import kestrel.store.page.Page
+
+sealed class BTreeNode {
+    abstract val pageId: Long
+
+    class Leaf(
+        override val pageId: Long,
+        val keys: MutableList<ByteArray>,
+        val values: MutableList<Long>,
+        var nextLeaf: Long
+    ) : BTreeNode()
+
+    /** `children.size` is always `keys.size + 1`: `children[0]` holds keys below `keys[0]`. */
+    class Internal(
+        override val pageId: Long,
+        val keys: MutableList<ByteArray>,
+        val children: MutableList<Long>
+    ) : BTreeNode()
+}
+
+/** Serializes [BTreeNode]s onto/from a fixed-size [Page]. */
+object NodeCodec {
+    private const val TYPE_LEAF = 0
+    private const val TYPE_INTERNAL = 1
+    const val NO_PAGE = -1L
+
+    fun serialize(node: BTreeNode): ByteArray {
+        val writer = ByteWriter(Page.PAGE_SIZE)
+        when (node) {
+            is BTreeNode.Leaf -> {
+                writer.writeByte(TYPE_LEAF)
+                writeCount(writer, node.keys.size)
+                writer.writeFixedLong(node.nextLeaf)
+                for (i in node.keys.indices) {
+                    writer.writeLengthPrefixed(node.keys[i])
+                    writer.writeFixedLong(node.values[i])
+                }
+            }
+            is BTreeNode.Internal -> {
+                writer.writeByte(TYPE_INTERNAL)
+                writeCount(writer, node.keys.size)
+                writer.writeFixedLong(node.children[0])
+                for (i in node.keys.indices) {
+                    writer.writeLengthPrefixed(node.keys[i])
+                    writer.writeFixedLong(node.children[i + 1])
+                }
+            }
+        }
+        return writer.toByteArray()
+    }
+
+    fun writeTo(page: Page, node: BTreeNode) {
+        val bytes = serialize(node)
+        if (bytes.size > Page.PAGE_SIZE) {
+            throw CodecException("btree node ${node.pageId} serializes to ${bytes.size} bytes, exceeds page size")
+        }
+        page.data.fill(0)
+        System.arraycopy(bytes, 0, page.data, 0, bytes.size)
+        page.isDirty = true
+    }
+
+    fun deserialize(pageId: Long, data: ByteArray): BTreeNode {
+        val reader = ByteReader(data)
+        val type = reader.readByte()
+        val count = readCount(reader)
+        return when (type) {
+            TYPE_LEAF -> {
+                val nextLeaf = reader.readFixedLong()
+                val keys = ArrayList<ByteArray>(count)
+                val values = ArrayList<Long>(count)
+                repeat(count) {
+                    keys.add(reader.readLengthPrefixed())
+                    values.add(reader.readFixedLong())
+                }
+                BTreeNode.Leaf(pageId, keys, values, nextLeaf)
+            }
+            TYPE_INTERNAL -> {
+                val firstChild = reader.readFixedLong()
+                val keys = ArrayList<ByteArray>(count)
+                val children = ArrayList<Long>(count + 1)
+                children.add(firstChild)
+                repeat(count) {
+                    keys.add(reader.readLengthPrefixed())
+                    children.add(reader.readFixedLong())
+                }
+                BTreeNode.Internal(pageId, keys, children)
+            }
+            else -> throw CodecException("unknown btree node type $type")
+        }
+    }
+
+    fun readFrom(page: Page): BTreeNode = deserialize(page.id, page.data)
+
+    fun fitsInPage(node: BTreeNode): Boolean = serialize(node).size <= Page.PAGE_SIZE
+
+    private fun writeCount(writer: ByteWriter, count: Int) {
+        writer.writeByte((count ushr 8) and 0xFF)
+        writer.writeByte(count and 0xFF)
+    }
+
+    private fun readCount(reader: ByteReader): Int = (reader.readByte() shl 8) or reader.readByte()
+}