feat(btree): implement deletion (d1f3dabd) · Commits · e / tools / webadb

libraries/b-tree/src/index.spec.ts

+77 −42

Original line number	Diff line number	Diff line
		import { describe, expect, it } from '@jest/globals';
		import { BTree } from "./index.js";
		import { BTree, BTreeNode } from "./index.js";

		const LENGTH = 128;

		function shuffle<T>(array: T[]) {
		for (let i = array.length - 1; i > 0; i -= 1) {
		const j = (Math.random() * (i + 1)) \| 0;
		[array[i], array[j]] = [array[j]!, array[i]!];
		}
		}

		describe('BTree', () => {
		describe('insert', () => {
		function validate(log: any[], min = -Infinity) {
		for (const item of log) {
		if (Array.isArray(item)) {
		min = validate(item, min);
		} else if (typeof item === 'number') {
		if (item < min) {
		throw new Error();
		for (let order = 3; order < 10; order += 1) {
		// Math.ceil(order / 2) - 1
		const MIN_KEY_COUNT = ((order + 1) >> 1) - 1;

		function validateNode(node: BTreeNode, root: boolean, min = -Infinity) {
		if (node.height === 0) {
		expect(node.keyCount).toBeGreaterThan(0);
		expect(node.keyCount).toBeLessThan(order);
		for (let i = 0; i < node.keyCount; i += 1) {
		expect(node.keys[i]).toBeGreaterThan(min);
		min = node.keys[i]!;
		}
		min = item;
		return min;
		}

		if (!root) {
		expect(node.keyCount).toBeGreaterThanOrEqual(MIN_KEY_COUNT);
		}
		expect(node.keyCount).toBeLessThan(order);

		for (let i = 0; i < node.keyCount; i += 1) {
		min = validateNode(node.children[i]!, false, min);
		expect(node.keys[i]).toBeGreaterThan(min);
		min = node.keys[i]!;
		}
		min = validateNode(node.children[node.keyCount]!, false, min);
		return min;
		}

		it('should work with incremental values', () => {
		const tree = new BTree(6);
		const values = Array.from({ length: 1024 }, (_, i) => i);
		function validateTree(tree: BTree) {
		if (tree.size === 0) {
		expect(tree.root.keyCount).toBe(0);
		return;
		}

		for (let i = 0; i < values.length; i += 1) {
		tree.insert(values[i]!);
		expect(() => validate(tree.root.log())).not.toThrow();
		validateNode(tree.root, true);
		}
		});

		it('should work with random data', () => {
		const tree = new BTree(6);
		const values = Array.from({ length: 1024 }, () => Math.random() * 1024 \| 0);
		// const values = Array.from({ length: 1024 }, (_, i) => 1024 - i);
		// const values = Array.from({ length: 1024 }, (_, i) => i);
		describe(`order ${order}`, () => {
		it('should generate valid tree with incremental values', () => {
		const tree = new BTree(order);

		for (let i = 0; i < values.length; i += 1) {
		tree.insert(values[i]!);
		expect(() => validate(tree.root.log())).not.toThrow();
		const values = Array.from({ length: LENGTH }, (_, i) => i - LENGTH / 2);
		for (let value of values) {
		tree.insert(value);
		validateTree(tree);
		expect(tree.has(value)).toBe(true);
		}

		for (let value of values) {
		tree.remove(value);
		validateTree(tree);
		expect(tree.has(value)).toBe(false);
		}
		});
		});

		describe('has', () => {
		it('should return true for inserted values', () => {
		const tree = new BTree(6);
		const values = Array.from({ length: 1024 }, () => Math.random() * 512 + 512 \| 0);
		it('should generate valid tree with random values', () => {
		const tree = new BTree(order);

		const values = Array.from({ length: LENGTH }, (_, i) => i - LENGTH / 2);
		shuffle(values);
		for (const value of values) {
		tree.insert(value);
		validateTree(tree);
		expect(tree.has(value)).toBe(true);
		}
		for (let i = -1024; i < 2048; i++) {
		expect(tree.has(i)).toBe(values.includes(i));

		shuffle(values);
		for (const value of values) {
		tree.remove(value);
		validateTree(tree);
		expect(tree.has(value)).toBe(false);
		}
		});
		});
		}
		});

libraries/b-tree/src/index.ts

+207 −62

Original line number	Diff line number	Diff line
		@@ -10,9 +10,16 @@ function insert(array: Int32Array, length: number, value: number, index: number)
		array[index] = value;
		}

		function remove(array: Int32Array, length: number, index: number) {
		if (index < length - 1) {
		array.set(array.subarray(index + 1, length), index);
		}
		}

		export class BTreeNode {
		order: number;
		mid: number;
		minKeyCount: number;

		keys: Int32Array;
		keyCount: number;
		@@ -29,6 +36,8 @@ export class BTreeNode {
		) {
		this.order = order;
		this.mid = this.order >> 1;
		// Math.ceil(order / 2) - 1
		this.minKeyCount = ((this.order + 1) >> 1) - 1;

		this.keys = keys;
		this.keyCount = keyCount;
		@@ -37,57 +46,65 @@ export class BTreeNode {
		this.children = children;
		}

		/**
		* Split the current node into two
		* @param value The key to be inserted.
		* @param index The index of the key to be inserted at.
		* @param child The child (right to the key) to be inserted. May be undefined when current node is a leaf.
		* @returns The new key and child need to be inserted to its parent.
		* The new key is the middle key of the current node, and the child contains the right half of the current node.
		*/
		protected split(value: number, index: number, child?: BTreeNode): BTreeInsertionResult {
		let parentKey: number;
		const siblingKeys = new Int32Array(this.order - 1);
		let siblingsChildren: BTreeNode[];
		let middleKey: number;
		const rightKeys = new Int32Array(this.order - 1);
		let rightChildren: BTreeNode[];

		if (index < this.mid) {
		parentKey = this.keys[this.mid - 1]!;
		siblingKeys.set(this.keys.subarray(this.mid), 0);
		middleKey = this.keys[this.mid - 1]!;
		rightKeys.set(this.keys.subarray(this.mid), 0);

		insert(this.keys, this.mid - 1, value, index);

		if (child) {
		// internal node
		siblingsChildren = this.children.splice(this.mid, this.order - this.mid);
		rightChildren = this.children.splice(this.mid, this.order - this.mid);
		// TODO: this may cause the underlying array to grow (re-alloc and copy)
		// investigate if this is a problem.
		// investigate if this hurts performance.
		this.children.splice(index + 1, 0, child);
		} else {
		// leaf node
		siblingsChildren = new Array(this.order);
		// leaf node, doesn't have children, create am empty array for it.
		rightChildren = new Array(this.order);
		}
		} else {
		if (index === this.mid) {
		parentKey = value;
		siblingKeys.set(this.keys.subarray(this.mid), 0);
		middleKey = value;
		rightKeys.set(this.keys.subarray(this.mid), 0);
		} else {
		parentKey = this.keys[this.mid]!;
		middleKey = this.keys[this.mid]!;
		if (index !== this.mid + 1) {
		siblingKeys.set(this.keys.subarray(this.mid + 1, index), 0);
		rightKeys.set(this.keys.subarray(this.mid + 1, index), 0);
		}
		siblingKeys[index - this.mid - 1] = value;
		siblingKeys.set(this.keys.subarray(index), index - this.mid);
		rightKeys[index - this.mid - 1] = value;
		rightKeys.set(this.keys.subarray(index), index - this.mid);
		}

		if (child) {
		siblingsChildren = this.children.splice(this.mid + 1, this.order - this.mid - 1);
		siblingsChildren.splice(index - this.mid, 0, child);
		rightChildren = this.children.splice(this.mid + 1, this.order - this.mid - 1);
		rightChildren.splice(index - this.mid, 0, child);
		} else {
		siblingsChildren = new Array(this.order);
		rightChildren = new Array(this.order);
		}
		}

		this.keyCount = this.mid;
		return {
		key: parentKey,
		key: middleKey,
		child: new BTreeNode(
		this.order,
		siblingKeys,
		rightKeys,
		this.order - 1 - this.mid,
		this.height,
		siblingsChildren
		rightChildren
		),
		};
		}
		@@ -108,10 +125,22 @@ export class BTreeNode {
		return ~start;
		}

		public insert(value: number): BTreeInsertionResult \| undefined {
		public has(value: number): boolean {
		let index = this.search(value);
		if (index >= 0) {
		return;
		return true;
		}
		if (this.height > 0) {
		index = ~index;
		return this.children[index]!.has(value);
		}
		return false;
		}

		public insert(value: number): BTreeInsertionResult \| boolean {
		let index = this.search(value);
		if (index >= 0) {
		return false;
		}

		index = ~index;
		@@ -123,12 +152,11 @@ export class BTreeNode {

		insert(this.keys, this.keyCount, value, index);
		this.keyCount += 1;
		return;
		return true;
		}

		const child = this.children[index]!;
		const split = child.insert(value);
		if (split) {
		const split = this.children[index]!.insert(value);
		if (typeof split === 'object') {
		if (this.keyCount === this.order - 1) {
		return this.split(split.key, index, split.child);
		}
		@@ -138,50 +166,127 @@ export class BTreeNode {

		this.children.splice(index + 1, 0, split.child);
		}

		return true;
		}

		public remove(value: number): boolean {
		let index = this.search(value);
		if (index >= 0) {
		this.removeAt(index);
		return true;
		}

		if (this.height > 0) {
		index = ~index;
		const removed = this.children[index]!.remove(value);
		if (removed) {
		this.balance(index);
		}
		return removed;
		}

		return false;
		}

		public max(): number {
		if (this.height === 0) {
		return this.keys[this.keyCount - 1]!;
		}
		return this.children[this.keyCount]!.max();
		}

		protected balance(index: number) {
		const child = this.children[index]!;

		if (child.keyCount >= this.minKeyCount) {
		return;
		}

		if (index > 0) {
		const left = this.children[index - 1]!;
		if (left.keyCount > this.minKeyCount) {
		// rotate right
		insert(child.keys, child.keyCount, this.keys[index - 1]!, 0);
		if (this.height > 1) {
		child.children.splice(0, 0, left.children[left.keyCount]!);
		}
		child.keyCount += 1;

		this.keys[index - 1] = left.keys[left.keyCount - 1]!;
		left.keyCount -= 1;
		return;
		}

		log(): any[] {
		const result = [];
		for (let i = 0; i < this.keyCount; i += 1) {
		result.push(this.children[i]?.log());
		result.push(this.keys[i]);
		// merge with left
		left.keys[left.keyCount] = this.keys[index - 1]!;
		left.keyCount += 1;
		left.keys.set(child.keys.subarray(0, child.keyCount), left.keyCount);
		if (this.height > 1) {
		for (let i = 0; i <= child.keyCount; i++) {
		left.children[left.keyCount + i] = child.children[i]!;
		}
		result.push(this.children[this.keyCount]?.log());
		return result;
		}
		left.keyCount += child.keyCount;
		remove(this.keys, this.keyCount, index - 1);
		this.children.splice(index, 1);
		this.keyCount -= 1;
		return;
		}

		export class BTreeRoot extends BTreeNode {
		public constructor(order: number) {
		super(
		order,
		new Int32Array(order - 1),
		0,
		0,
		new Array(order)
		);
		const right = this.children[index + 1]!;
		if (right.keyCount > this.minKeyCount) {
		// rotate left
		child.keys[child.keyCount] = this.keys[index]!;
		if (this.height > 1) {
		child.children[child.keyCount + 1] = right.children.splice(0, 1)[0]!;
		}
		child.keyCount += 1;

		protected override split(value: number, index: number, child?: BTreeNode): BTreeInsertionResult {
		const split = super.split(value, index, child);
		this.keys[index] = right.keys[0]!;

		this.children[0] = new BTreeNode(
		this.order,
		this.keys,
		this.keyCount,
		this.height,
		this.children.slice(),
		);
		remove(right.keys, right.keyCount, 0);
		right.keyCount -= 1;
		return;
		}

		// merge right into child
		child.keys[child.keyCount] = this.keys[index]!;
		child.keyCount += 1;
		child.keys.set(right.keys.subarray(0, right.keyCount), child.keyCount);
		if (this.height > 1) {
		for (let i = 0; i <= right.keyCount; i++) {
		child.children[child.keyCount + i] = right.children[i]!;
		}
		}
		child.keyCount += right.keyCount;
		remove(this.keys, this.keyCount, index);
		this.children.splice(index + 1, 1);
		this.keyCount -= 1;
		}

		protected removeMax(): void {
		if (this.height === 0) {
		this.keyCount -= 1;
		return;
		}

		this.children[1] = split.child;
		const child = this.children[this.keyCount]!;
		child.removeMax();
		this.balance(this.keyCount);
		}

		this.keys = new Int32Array(this.order);
		this.keys[0] = split.key;
		this.keyCount = 1;
		this.height += 1;
		protected removeAt(index: number) {
		if (this.height === 0) {
		remove(this.keys, this.keyCount, index);
		this.keyCount -= 1;
		return;
		}

		return split;
		const max = this.children[index]!.max();
		this.keys[index] = max;
		this.children[index]!.removeMax();
		this.balance(index);
		}
		}

		@@ -189,12 +294,21 @@ export class BTree {
		order: number;
		root: BTreeNode;

		size: number = 0;

		public constructor(order: number) {
		this.order = order;
		this.root = new BTreeRoot(order);
		this.root = new BTreeNode(
		order,
		new Int32Array(order - 1),
		0,
		0,
		new Array(order)
		);
		}

		public has(value: number) {
		// TODO(btree): benchmark this non-recursive version
		let node = this.root;
		while (true) {
		const index = node.search(value);
		@@ -210,6 +324,37 @@ export class BTree {
		}

		public insert(value: number) {
		this.root.insert(value);
		const split = this.root.insert(value);
		if (typeof split === 'object') {
		const keys = new Int32Array(this.order - 1);
		keys[0] = split.key;

		const children = new Array(this.order);
		children[0] = this.root;
		children[1] = split.child;

		this.root = new BTreeNode(
		this.order,
		keys,
		1,
		this.root.height + 1,
		children
		);
		}
		if (split) {
		this.size += 1;
		}
		return !!split;
		}

		public remove(value: number) {
		const removed = this.root.remove(value);
		if (removed) {
		if (this.root.height > 0 && this.root.keyCount === 0) {
		this.root = this.root.children[0]!;
		}
		this.size -= 1;
		}
		return removed;
		}
		}