From 3e72df475e16a147074c2ab9a06e0728c961e32c Mon Sep 17 00:00:00 2001 From: DualFroz Date: Fri, 4 Sep 2026 17:19:24 +0200 Subject: [PATCH] feat(btree): add insertRecursive and splitLeaf to BPlusTree --- .../kotlin/kestrel/store/btree/BPlusTree.kt | 162 ++++++++++++++++++ 1 file changed, 162 insertions(+) diff --git a/src/main/kotlin/kestrel/store/btree/BPlusTree.kt b/src/main/kotlin/kestrel/store/btree/BPlusTree.kt index 3c9083d..2fa0d62 100644 --- a/src/main/kotlin/kestrel/store/btree/BPlusTree.kt +++ b/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) + } +