B+ trees and indexes.

type Page = {
	leaf: boolean;
	keys: number[];
	entries: IndexEntry[]; // leaves only, one per key
	children: Page[]; // internal pages only, one more than keys
	next: Page | null; // leaves only: the leaf with the next keys
};

const leafPage = (): Page => ({ leaf: true, keys: [], entries: [], children: [], next: null });

// A B+ tree keeps every record in a leaf page, sorted, and links each leaf to the next.
// Internal pages hold separators: keys[i] is the smallest key in children[i + 1]. A page
// that overflows splits in two and hands a separator to its parent; when the root splits,
// a new root grows above it, so every leaf stays at the same depth.
export class BPlusIndex {
	#root: Page = leafPage();
	#size = 0;
	#pages = 1;
	#height = 1;
	#pagesRead = 0;

	get size(): number {
		return this.#size;
	}
	get pages(): number {
		return this.#pages;
	}
	/** Levels from the root down to the leaves. */
	get height(): number {
		return this.#height;
	}
	/** Pages read by lookups and ranges since the index was built. */
	get pagesRead(): number {
		return this.#pagesRead;
	}

	insert(entry: IndexEntry, steps?: IndexStep[]): boolean {
		const record = validateEntry(entry, 0);
		const path: Page[] = [];
		let page = this.#root;
		while (!page.leaf) {
			path.push(page);
			page = page.children[childFor(page, record.key)];
		}
		const at = lowerBound(page.keys, record.key);
		if (page.keys[at] === record.key) return false;
		page.keys.splice(at, 0, record.key);
		page.entries.splice(at, 0, record);
		this.#size++;
		steps?.push({ kind: 'insert', key: record.key, keys: [...page.keys] });

		let level = path.length;
		let overflow: { separator: number; right: Page } | null =
			page.keys.length > PAGE_SIZE ? this.#splitLeaf(page, level, steps) : null;
		while (overflow) {
			const parent = path.pop();
			level--;
			if (!parent) {
				this.#root = {
					leaf: false,
					keys: [overflow.separator],
					entries: [],
					children: [page, overflow.right],
					next: null
				};
				this.#pages++;
				this.#height++;
				steps?.push({ kind: 'grow', keys: [overflow.separator] });
				break;
			}
			const slot = childFor(parent, overflow.separator);
			parent.keys.splice(slot, 0, overflow.separator);
			parent.children.splice(slot + 1, 0, overflow.right);
			page = parent;
			overflow = parent.keys.length > PAGE_SIZE ? this.#splitInternal(parent, level, steps) : null;
		}
		steps?.push({ kind: 'shape', levels: this.shape() });
		return true;
	}

	get(key: number, steps?: IndexStep[]): IndexEntry | null {
		const leaf = this.#leafFor(key, steps);
		const at = lowerBound(leaf.keys, key);
		if (leaf.keys[at] !== key) return null;
		steps?.push({ kind: 'match', key });
		return { ...leaf.entries[at] };
	}

	/** Every record with lower ≤ key ≤ upper, read down to one leaf, then along the leaf links. */
	range(lower: number, upper: number, steps?: IndexStep[]): IndexEntry[] {
		validateRange(lower, upper);
		const found: IndexEntry[] = [];
		let leaf: Page | null = this.#leafFor(lower, steps);
		while (leaf) {
			for (const entry of leaf.entries) {
				if (entry.key > upper) return found;
				if (entry.key >= lower) {
					found.push({ ...entry });
					steps?.push({ kind: 'match', key: entry.key });
				}
			}
			leaf = leaf.next;
			if (leaf) {
				this.#pagesRead++;
				steps?.push({ kind: 'next', keys: [...leaf.keys] });
			}
		}
		return found;
	}

	/** Each level's pages, as their keys, root first. */
	shape(): number[][][] {
		const levels: number[][][] = [];
		let row = [this.#root];
		while (row.length) {
			levels.push(row.map((page) => [...page.keys]));
			row = row.flatMap((page) => page.children);
		}
		return levels;
	}

	#leafFor(key: number, steps?: IndexStep[]): Page {
		let page = this.#root;
		let level = 0;
		for (;;) {
			this.#pagesRead++;
			steps?.push({ kind: 'page', level, keys: [...page.keys] });
			if (page.leaf) return page;
			page = page.children[childFor(page, key)];
			level++;
		}
	}

	#splitLeaf(page: Page, level: number, steps?: IndexStep[]): { separator: number; right: Page } {
		const half = Math.ceil(page.keys.length / 2);
		const right: Page = {
			leaf: true,
			keys: page.keys.splice(half),
			entries: page.entries.splice(half),
			children: [],
			next: page.next
		};
		page.next = right;
		this.#pages++;
		// A leaf split copies the right page's first key up: the record stays in the leaf.
		const separator = right.keys[0];
		steps?.push({ kind: 'split', level, left: [...page.keys], right: [...right.keys], separator });
		return { separator, right };
	}

	#splitInternal(
		page: Page,
		level: number,
		steps?: IndexStep[]
	): { separator: number; right: Page } {
		const middle = Math.floor(page.keys.length / 2);
		// An internal split moves the middle separator up; it is not kept on either side.
		const separator = page.keys[middle];
		const right: Page = {
			leaf: false,
			keys: page.keys.splice(middle + 1),
			entries: [],
			children: page.children.splice(middle + 1),
			next: null
		};
		page.keys.pop();
		this.#pages++;
		steps?.push({ kind: 'split', level, left: [...page.keys], right: [...right.keys], separator });
		return { separator, right };
	}
}

// The child that can hold key: the first separator greater than key marks its right edge.
function childFor(page: Page, key: number): number {
	let child = 0;
	while (child < page.keys.length && key >= page.keys[child]) child++;
	return child;
}

function lowerBound(keys: readonly number[], key: number): number {
	let low = 0;
	let high = keys.length;
	while (low < high) {
		const mid = low + Math.floor((high - low) / 2);
		if (keys[mid] < key) low = mid + 1;
		else high = mid;
	}
	return low;
}

type page struct {
	leaf     bool
	keys     []int
	entries  []IndexEntry // leaves only, one per key
	children []*page      // internal pages only, one more than keys
	next     *page        // leaves only: the leaf with the next keys
}

// BPlusIndex keeps every record in a leaf page, sorted, and links each leaf to the next.
// Internal pages hold separators: keys[i] is the smallest key in children[i+1]. A page that
// overflows splits in two and hands a separator to its parent; when the root splits, a new
// root grows above it, so every leaf stays at the same depth.
type BPlusIndex struct {
	root      *page
	size      int
	pages     int
	height    int
	pagesRead int
}

func NewBPlusIndex() *BPlusIndex {
	return &BPlusIndex{root: &page{leaf: true}, pages: 1, height: 1}
}

func (b *BPlusIndex) Size() int      { return b.size }
func (b *BPlusIndex) Pages() int     { return b.pages }
func (b *BPlusIndex) Height() int    { return b.height }
func (b *BPlusIndex) PagesRead() int { return b.pagesRead }

// Insert adds a record and reports false when its key is already indexed.
func (b *BPlusIndex) Insert(entry IndexEntry) (bool, error) {
	if err := validateEntry(entry, 0); err != nil {
		return false, err
	}
	path := []*page{}
	current := b.root
	for !current.leaf {
		path = append(path, current)
		current = current.children[childFor(current, entry.Key)]
	}
	at := sort.SearchInts(current.keys, entry.Key)
	if at < len(current.keys) && current.keys[at] == entry.Key {
		return false, nil
	}
	current.keys = insertAt(current.keys, at, entry.Key)
	current.entries = append(current.entries, IndexEntry{})
	copy(current.entries[at+1:], current.entries[at:])
	current.entries[at] = entry
	b.size++

	var separator int
	var right *page
	split := false
	if len(current.keys) > PageSize {
		separator, right = b.splitLeaf(current)
		split = true
	}
	for split {
		if len(path) == 0 {
			b.root = &page{keys: []int{separator}, children: []*page{current, right}}
			b.pages++
			b.height++
			break
		}
		parent := path[len(path)-1]
		path = path[:len(path)-1]
		slot := childFor(parent, separator)
		parent.keys = insertAt(parent.keys, slot, separator)
		parent.children = append(parent.children, nil)
		copy(parent.children[slot+2:], parent.children[slot+1:])
		parent.children[slot+1] = right
		current = parent
		split = false
		if len(parent.keys) > PageSize {
			separator, right = b.splitInternal(parent)
			split = true
		}
	}
	return true, nil
}

// Get finds one record by key.
func (b *BPlusIndex) Get(key int) (IndexEntry, bool) {
	leaf := b.leafFor(key)
	at := sort.SearchInts(leaf.keys, key)
	if at == len(leaf.keys) || leaf.keys[at] != key {
		return IndexEntry{}, false
	}
	return leaf.entries[at], true
}

// Range returns every record with lower ≤ key ≤ upper, read down to one leaf, then along the
// leaf links.
func (b *BPlusIndex) Range(lower, upper int) ([]IndexEntry, error) {
	if err := validateRange(lower, upper); err != nil {
		return nil, err
	}
	found := []IndexEntry{}
	for leaf := b.leafFor(lower); leaf != nil; {
		for _, entry := range leaf.entries {
			if entry.Key > upper {
				return found, nil
			}
			if entry.Key >= lower {
				found = append(found, entry)
			}
		}
		leaf = leaf.next
		if leaf != nil {
			b.pagesRead++
		}
	}
	return found, nil
}

// Shape lists each level's pages as their keys, root first.
func (b *BPlusIndex) Shape() [][][]int {
	levels := [][][]int{}
	row := []*page{b.root}
	for len(row) > 0 {
		level := [][]int{}
		next := []*page{}
		for _, p := range row {
			level = append(level, append([]int(nil), p.keys...))
			next = append(next, p.children...)
		}
		levels = append(levels, level)
		row = next
	}
	return levels
}

func (b *BPlusIndex) leafFor(key int) *page {
	current := b.root
	for {
		b.pagesRead++
		if current.leaf {
			return current
		}
		current = current.children[childFor(current, key)]
	}
}

func (b *BPlusIndex) splitLeaf(p *page) (int, *page) {
	half := (len(p.keys) + 1) / 2
	right := &page{
		leaf:    true,
		keys:    append([]int(nil), p.keys[half:]...),
		entries: append([]IndexEntry(nil), p.entries[half:]...),
		next:    p.next,
	}
	p.keys = p.keys[:half:half]
	p.entries = p.entries[:half:half]
	p.next = right
	b.pages++
	// A leaf split copies the right page's first key up: the record stays in the leaf.
	return right.keys[0], right
}

func (b *BPlusIndex) splitInternal(p *page) (int, *page) {
	middle := len(p.keys) / 2
	// An internal split moves the middle separator up; it is not kept on either side.
	separator := p.keys[middle]
	right := &page{
		keys:     append([]int(nil), p.keys[middle+1:]...),
		children: append([]*page(nil), p.children[middle+1:]...),
	}
	p.keys = p.keys[:middle:middle]
	p.children = p.children[: middle+1 : middle+1]
	b.pages++
	return separator, right
}

// childFor is the child that can hold key: the first separator greater than key marks its
// right edge.
func childFor(p *page, key int) int {
	child := 0
	for child < len(p.keys) && key >= p.keys[child] {
		child++
	}
	return child
}

func insertAt(keys []int, at, key int) []int {
	keys = append(keys, 0)
	copy(keys[at+1:], keys[at:])
	keys[at] = key
	return keys
}