From b54f41d3b577fcc8a7165e52dbfebf3332df9419 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Fri, 4 Sep 2026 17:30:46 +0200 Subject: [PATCH] feat(btree): add borrowFromRight, mergeChildren and BTreeNode --- .../kotlin/kestrel/store/btree/BPlusTree.kt | 77 +++++++++++++ .../kotlin/kestrel/store/btree/BTreeNode.kt | 107 ++++++++++++++++++ 2 files changed, 184 insertions(+) create mode 100644 src/main/kotlin/kestrel/store/btree/BTreeNode.kt diff --git a/src/main/kotlin/kestrel/store/btree/BPlusTree.kt b/src/main/kotlin/kestrel/store/btree/BPlusTree.kt index 2fa0d62..77b6b6a 100644 --- a/src/main/kotlin/kestrel/store/btree/BPlusTree.kt +++ b/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 { + val page = bufferPool.fetchPage(pageId) + return page to NodeCodec.readFrom(page) + } + + private fun lowerBound(keys: List, 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, 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 + } +} diff --git a/src/main/kotlin/kestrel/store/btree/BTreeNode.kt b/src/main/kotlin/kestrel/store/btree/BTreeNode.kt new file mode 100644 index 0000000..17ae210 --- /dev/null +++ b/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, + val values: MutableList, + 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, + val children: MutableList + ) : 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(count) + val values = ArrayList(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(count) + val children = ArrayList(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() +}