tree.getElementPath

joaomoreno · joaomoreno · commit 967d43e22ad6 · 2018-06-22T15:33:59.000+02:00
diff --git a/src/vs/base/common/arrays.ts b/src/vs/base/common/arrays.ts
@@ -24,6 +24,10 @@ export function tail2<T>(arr: T[]): [T[], T] {
 	return [arr.slice(0, arr.length - 1), arr[arr.length - 1]];
 }
 
+export function last<T>(arr: T[]): T {
+	return arr[arr.length - 1];
+}
+
 export function equals<T>(one: ReadonlyArray<T>, other: ReadonlyArray<T>, itemEquals: (a: T, b: T) => boolean = (a, b) => a === b): boolean {
 	if (one.length !== other.length) {
 		return false;
diff --git a/src/vs/base/common/tree.ts b/src/vs/base/common/tree.ts
@@ -5,53 +5,65 @@
 
 'use strict';
 
-import { clamp } from './numbers';
+import { tail2, last } from 'vs/base/common/arrays';
 
 export interface ITreeNode<T> {
 	readonly element: T;
 	readonly children: ITreeNode<T>[];
 }
 
-function clone<T>(nodes: ITreeNode<T>[]): ITreeNode<T>[] {
-	return nodes.map(({ element, children }) => ({ element, children: clone(children) }));
-}
-
 export class Tree<T> {
 
-	private root: ITreeNode<T>[] = [];
+	private root: ITreeNode<T> = { element: undefined, children: [] };
 
 	splice(start: number[], deleteCount: number, nodes: ITreeNode<T>[] = []): ITreeNode<T>[] {
 		if (start.length === 0) {
-			throw new Error('invalid tree location');
+			throw new Error('Invalid tree location');
 		}
 
-		const children = this.findChildren(start);
-		const index = start[start.length - 1];
-		return children.splice(index, deleteCount, ...clone(nodes));
+		const [rest, index] = tail2(start);
+		const parentPath = this.getNodePath(rest);
+		const parent = last(parentPath);
+
+		return parent.children.splice(index, deleteCount, ...nodes);
 	}
 
-	getElement(location: number[]): T | undefined {
-		const node = this.findElement(location);
-		return node && node.element;
+	getElement(location: number[]): T {
+		if (location.length === 0) {
+			throw new Error('Invalid tree location');
+		}
+
+		const nodePath = this.getNodePath(location);
+		const node = last(nodePath);
+
+		return node.element;
 	}
 
-	getNodes(): ITreeNode<T>[] {
-		return clone(this.root);
+	getElementPath(location: number[]): T[] {
+		if (location.length === 0) {
+			throw new Error('Invalid tree location');
+		}
+
+		const nodePath = this.getNodePath(location);
+		return nodePath.map(node => node.element);
 	}
 
-	private findElement(location: number[]): ITreeNode<T> {
-		const children = this.findChildren(location);
-		const lastIndex = clamp(location[location.length - 1], 0, children.length);
-		return children[lastIndex];
+	getNodes(): ITreeNode<T>[] {
+		return this.root.children;
 	}
 
-	private findChildren(location: number[], children: ITreeNode<T>[] = this.root): ITreeNode<T>[] {
-		if (location.length === 1) {
-			return children;
+	private getNodePath(location: number[], node: ITreeNode<T> = this.root): ITreeNode<T>[] {
+		if (location.length === 0) {
+			return [node];
+		}
+
+		let [index, ...rest] = location;
+
+		if (index < 0 || index >= node.children.length) {
+			throw new Error('Invalid location');
 		}
 
-		let [i, ...rest] = location;
-		i = clamp(i, 0, children.length);
-		return this.findChildren(rest, children[i].children);
+		const child = node.children[index];
+		return [node, ...this.getNodePath(rest, child)];
 	}
 }
