concat23Trees.ts

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import { AstNode, AstNodeKind, ListAstNode } from './ast.js';

/**
 * Concatenates a list of (2,3) AstNode's into a single (2,3) AstNode.
 * This mutates the items of the input array!
 * If all items have the same height, this method has runtime O(items.length).
 * Otherwise, it has runtime O(items.length * max(log(items.length), items.max(i => i.height))).
*/
export function concat23Trees(items: AstNode[]): AstNode | null {
	if (items.length === 0) {
		return null;
	}
	if (items.length === 1) {
		return items[0];
	}

	let i = 0;
	/**
	 * Reads nodes of same height and concatenates them to a single node.
	*/
	function readNode(): AstNode | null {
		if (i >= items.length) {
			return null;
		}
		const start = i;
		const height = items[start].listHeight;

		i++;
		while (i < items.length && items[i].listHeight === height) {
			i++;
		}

		if (i - start >= 2) {
			return concat23TreesOfSameHeight(start === 0 && i === items.length ? items : items.slice(start, i), false);
		} else {
			return items[start];
		}
	}

	// The items might not have the same height.
	// We merge all items by using a binary concat operator.
	let first = readNode()!; // There must be a first item
	let second = readNode();
	if (!second) {
		return first;
	}

	for (let item = readNode(); item; item = readNode()) {
		// Prefer concatenating smaller trees, as the runtime of concat depends on the tree height.
		if (heightDiff(first, second) <= heightDiff(second, item)) {
			first = concat(first, second);
			second = item;
		} else {
			second = concat(second, item);
		}
	}

	const result = concat(first, second);
	return result;
}

export function concat23TreesOfSameHeight(items: AstNode[], createImmutableLists: boolean = false): AstNode | null {
	if (items.length === 0) {
		return null;
	}
	if (items.length === 1) {
		return items[0];
	}

	let length = items.length;
	// All trees have same height, just create parent nodes.
	while (length > 3) {
		const newLength = length >> 1;
		for (let i = 0; i < newLength; i++) {
			const j = i << 1;
			items[i] = ListAstNode.create23(items[j], items[j + 1], j + 3 === length ? items[j + 2] : null, createImmutableLists);
		}
		length = newLength;
	}
	return ListAstNode.create23(items[0], items[1], length >= 3 ? items[2] : null, createImmutableLists);
}

function heightDiff(node1: AstNode, node2: AstNode): number {
	return Math.abs(node1.listHeight - node2.listHeight);
}

function concat(node1: AstNode, node2: AstNode): AstNode {
	if (node1.listHeight === node2.listHeight) {
		return ListAstNode.create23(node1, node2, null, false);
	}
	else if (node1.listHeight > node2.listHeight) {
		// node1 is the tree we want to insert into
		return append(node1 as ListAstNode, node2);
	} else {
		return prepend(node2 as ListAstNode, node1);
	}
}

/**
 * Appends the given node to the end of this (2,3) tree.
 * Returns the new root.
*/
function append(list: ListAstNode, nodeToAppend: AstNode): AstNode {
	list = list.toMutable() as ListAstNode;
	let curNode: AstNode = list;
	const parents: ListAstNode[] = [];
	let nodeToAppendOfCorrectHeight: AstNode | undefined;
	while (true) {
		// assert nodeToInsert.listHeight <= curNode.listHeight
		if (nodeToAppend.listHeight === curNode.listHeight) {
			nodeToAppendOfCorrectHeight = nodeToAppend;
			break;
		}
		// assert 0 <= nodeToInsert.listHeight < curNode.listHeight
		if (curNode.kind !== AstNodeKind.List) {
			throw new Error('unexpected');
		}
		parents.push(curNode);
		// assert 2 <= curNode.childrenLength <= 3
		curNode = curNode.makeLastElementMutable()!;
	}
	// assert nodeToAppendOfCorrectHeight!.listHeight === curNode.listHeight
	for (let i = parents.length - 1; i >= 0; i--) {
		const parent = parents[i];
		if (nodeToAppendOfCorrectHeight) {
			// Can we take the element?
			if (parent.childrenLength >= 3) {
				// assert parent.childrenLength === 3 && parent.listHeight === nodeToAppendOfCorrectHeight.listHeight + 1

				// we need to split to maintain (2,3)-tree property.
				// Send the third element + the new element to the parent.
				nodeToAppendOfCorrectHeight = ListAstNode.create23(parent.unappendChild()!, nodeToAppendOfCorrectHeight, null, false);
			} else {
				parent.appendChildOfSameHeight(nodeToAppendOfCorrectHeight);
				nodeToAppendOfCorrectHeight = undefined;
			}
		} else {
			parent.handleChildrenChanged();
		}
	}
	if (nodeToAppendOfCorrectHeight) {
		return ListAstNode.create23(list, nodeToAppendOfCorrectHeight, null, false);
	} else {
		return list;
	}
}

/**
 * Prepends the given node to the end of this (2,3) tree.
 * Returns the new root.
*/
function prepend(list: ListAstNode, nodeToAppend: AstNode): AstNode {
	list = list.toMutable() as ListAstNode;
	let curNode: AstNode = list;
	const parents: ListAstNode[] = [];
	// assert nodeToInsert.listHeight <= curNode.listHeight
	while (nodeToAppend.listHeight !== curNode.listHeight) {
		// assert 0 <= nodeToInsert.listHeight < curNode.listHeight
		if (curNode.kind !== AstNodeKind.List) {
			throw new Error('unexpected');
		}
		parents.push(curNode);
		// assert 2 <= curNode.childrenFast.length <= 3
		curNode = curNode.makeFirstElementMutable()!;
	}
	let nodeToPrependOfCorrectHeight: AstNode | undefined = nodeToAppend;
	// assert nodeToAppendOfCorrectHeight!.listHeight === curNode.listHeight
	for (let i = parents.length - 1; i >= 0; i--) {
		const parent = parents[i];
		if (nodeToPrependOfCorrectHeight) {
			// Can we take the element?
			if (parent.childrenLength >= 3) {
				// assert parent.childrenLength === 3 && parent.listHeight === nodeToAppendOfCorrectHeight.listHeight + 1

				// we need to split to maintain (2,3)-tree property.
				// Send the third element + the new element to the parent.
				nodeToPrependOfCorrectHeight = ListAstNode.create23(nodeToPrependOfCorrectHeight, parent.unprependChild()!, null, false);
			} else {
				parent.prependChildOfSameHeight(nodeToPrependOfCorrectHeight);
				nodeToPrependOfCorrectHeight = undefined;
			}
		} else {
			parent.handleChildrenChanged();
		}
	}
	if (nodeToPrependOfCorrectHeight) {
		return ListAstNode.create23(nodeToPrependOfCorrectHeight, list, null, false);
	} else {
		return list;
	}
}