feat(btree): add insertRecursive and splitLeaf to BPlusTree

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

@@ -151,3 +151,165 @@ class BPlusTree private constructor(private val bufferPool: PageStore, rootPageI
 
     private data class SplitResult(val promotedKey: ByteArray, val newPageId: Long)
 
+    private fun insertRecursive(pageId: Long, key: ByteArray, value: Long): SplitResult? {
+        val (page, node) = fetch(pageId)
+        var dirty = false
+        val result: SplitResult?
+        when (node) {
+            is BTreeNode.Leaf -> {
+                val idx = lowerBound(node.keys, key)
+                if (idx < node.keys.size && KeyCodec.compare(node.keys[idx], key) == 0) {
+                    node.values[idx] = value
+                } else {
+                    node.keys.add(idx, key)
+                    node.values.add(idx, value)
+                }
+                dirty = true
+                result = if (NodeCodec.fitsInPage(node)) {
+                    NodeCodec.writeTo(page, node)
+                    null
+                } else {
+                    splitLeaf(page, node)
+                }
+            }
+            is BTreeNode.Internal -> {
+                val idx = childIndexFor(node.keys, key)
+                val childSplit = insertRecursive(node.children[idx], key, value)
+                if (childSplit == null) {
+                    result = null
+                } else {
+                    node.keys.add(idx, childSplit.promotedKey)
+                    node.children.add(idx + 1, childSplit.newPageId)
+                    dirty = true
+                    result = if (NodeCodec.fitsInPage(node)) {
+                        NodeCodec.writeTo(page, node)
+                        null
+                    } else {
+                        splitInternal(page, node)
+                    }
+                }
+            }
+        }
+        bufferPool.unpinPage(pageId, dirty)
+        return result
+    }
+
+    private fun splitLeaf(page: Page, leaf: BTreeNode.Leaf): SplitResult {
+        val mid = leaf.keys.size / 2
+        val rightKeys = ArrayList(leaf.keys.subList(mid, leaf.keys.size))
+        val rightValues = ArrayList(leaf.values.subList(mid, leaf.values.size))
+        val leftKeys = ArrayList(leaf.keys.subList(0, mid))
+        val leftValues = ArrayList(leaf.values.subList(0, mid))
+        val rightPage = bufferPool.newPage()
+        val rightNode = BTreeNode.Leaf(rightPage.id, rightKeys, rightValues, leaf.nextLeaf)
+        val leftNode = BTreeNode.Leaf(page.id, leftKeys, leftValues, rightPage.id)
+        NodeCodec.writeTo(page, leftNode)
+        NodeCodec.writeTo(rightPage, rightNode)
+        bufferPool.unpinPage(rightPage.id, true)
+        return SplitResult(rightKeys[0], rightPage.id)
+    }
+
+    private fun splitInternal(page: Page, node: BTreeNode.Internal): SplitResult {
+        val mid = node.keys.size / 2
+        val promotedKey = node.keys[mid]
+        val leftKeys = ArrayList(node.keys.subList(0, mid))
+        val leftChildren = ArrayList(node.children.subList(0, mid + 1))
+        val rightKeys = ArrayList(node.keys.subList(mid + 1, node.keys.size))
+        val rightChildren = ArrayList(node.children.subList(mid + 1, node.children.size))
+        val rightPage = bufferPool.newPage()
+        val rightNode = BTreeNode.Internal(rightPage.id, rightKeys, rightChildren)
+        val leftNode = BTreeNode.Internal(page.id, leftKeys, leftChildren)
+        NodeCodec.writeTo(page, leftNode)
+        NodeCodec.writeTo(rightPage, rightNode)
+        bufferPool.unpinPage(rightPage.id, true)
+        return SplitResult(promotedKey, rightPage.id)
+    }
+
+    private fun deleteRecursive(pageId: Long, key: ByteArray): Boolean {
+        val (page, node) = fetch(pageId)
+        val found: Boolean
+        when (node) {
+            is BTreeNode.Leaf -> {
+                val idx = lowerBound(node.keys, key)
+                found = idx < node.keys.size && KeyCodec.compare(node.keys[idx], key) == 0
+                if (found) {
+                    node.keys.removeAt(idx)
+                    node.values.removeAt(idx)
+                    NodeCodec.writeTo(page, node)
+                }
+                bufferPool.unpinPage(pageId, found)
+            }
+            is BTreeNode.Internal -> {
+                val idx = childIndexFor(node.keys, key)
+                val childId = node.children[idx]
+                found = deleteRecursive(childId, key)
+                if (found) {
+                    val (_, childNode) = fetch(childId)
+                    val underflow = isUnderflowNode(childNode)
+                    bufferPool.unpinPage(childId, false)
+                    if (underflow) {
+                        rebalanceChild(node, idx)
+                        NodeCodec.writeTo(page, node)
+                    }
+                }
+                bufferPool.unpinPage(pageId, found)
+            }
+        }
+        return found
+    }
+
+    private fun rebalanceChild(parent: BTreeNode.Internal, idx: Int) {
+        val leftIdx = idx - 1
+        val rightIdx = idx + 1
+
+        if (leftIdx >= 0) {
+            val (leftPage, leftNode) = fetch(parent.children[leftIdx])
+            if (canLend(leftNode)) {
+                borrowFromLeft(parent, idx, leftPage, leftNode)
+                bufferPool.unpinPage(leftNode.pageId, true)
+                return
+            }
+            bufferPool.unpinPage(leftNode.pageId, false)
+        }
+        if (rightIdx < parent.children.size) {
+            val (rightPage, rightNode) = fetch(parent.children[rightIdx])
+            if (canLend(rightNode)) {
+                borrowFromRight(parent, idx, rightPage, rightNode)
+                bufferPool.unpinPage(rightNode.pageId, true)
+                return
+            }
+            bufferPool.unpinPage(rightNode.pageId, false)
+        }
+        if (leftIdx >= 0) mergeChildren(parent, leftIdx) else mergeChildren(parent, idx)
+    }
+
+    private fun canLend(node: BTreeNode): Boolean = when (node) {
+        is BTreeNode.Leaf -> node.keys.size > 1
+        is BTreeNode.Internal -> node.keys.size > 1
+    }
+
+    private fun borrowFromLeft(parent: BTreeNode.Internal, idx: Int, leftPage: Page, leftNode: BTreeNode) {
+        val (childPage, childNode) = fetch(parent.children[idx])
+        when {
+            leftNode is BTreeNode.Leaf && childNode is BTreeNode.Leaf -> {
+                val lastIdx = leftNode.keys.size - 1
+                val borrowedKey = leftNode.keys.removeAt(lastIdx)
+                val borrowedValue = leftNode.values.removeAt(lastIdx)
+                childNode.keys.add(0, borrowedKey)
+                childNode.values.add(0, borrowedValue)
+                parent.keys[idx - 1] = childNode.keys[0]
+            }
+            leftNode is BTreeNode.Internal && childNode is BTreeNode.Internal -> {
+                val movedChild = leftNode.children.removeAt(leftNode.children.size - 1)
+                val separatorDown = leftNode.keys.removeAt(leftNode.keys.size - 1)
+                childNode.children.add(0, movedChild)
+                childNode.keys.add(0, parent.keys[idx - 1])
+                parent.keys[idx - 1] = separatorDown
+            }
+            else -> throw BTreeCorruptionException("sibling type mismatch during redistribution")
+        }
+        NodeCodec.writeTo(leftPage, leftNode)
+        NodeCodec.writeTo(childPage, childNode)
+        bufferPool.unpinPage(childNode.pageId, true)
+    }
+