improve TrieMap in prep for microsoft#10583

jrieken · jrieken · commit 3ecfa314f0bd · 2016-11-04T17:18:58.000+01:00
diff --git a/src/vs/base/common/map.ts b/src/vs/base/common/map.ts
@@ -310,7 +310,7 @@ interface Node<E> {
  */
 export class TrieMap<E> {
 
-	static PathSplitter = s => s.split(/[\\/]/);
+	static PathSplitter = s => s.split(/[\\/]/).filter(s => !!s);
 
 	private _splitter: (s: string) => string[];
 	private _root: Node<E> = { children: Object.create(null) };
@@ -345,6 +345,22 @@ export class TrieMap<E> {
 		node.element = element;
 	}
 
+	lookUp(path: string): E {
+		const parts = this._splitter(path);
+
+		let {children} = this._root;
+		let node: Node<E>;
+		for (const part of parts) {
+			node = children[part];
+			if (!node) {
+				return;
+			}
+			children = node.children;
+		}
+
+		return node.element;
+	}
+
 	findSubstr(path: string): E {
 		const parts = this._splitter(path);
 
@@ -367,4 +383,22 @@ export class TrieMap<E> {
 			return lastNode.element;
 		}
 	}
+
+	findSuperstr(path: string): TrieMap<E> {
+		const parts = this._splitter(path);
+
+		let {children} = this._root;
+		let node: Node<E>;
+		for (const part of parts) {
+			node = children[part];
+			if (!node) {
+				return;
+			}
+			children = node.children;
+		}
+
+		const result = new TrieMap<E>(this._splitter);
+		result._root = node;
+		return result;
+	}
 }
diff --git a/src/vs/base/test/common/map.test.ts b/src/vs/base/test/common/map.test.ts
@@ -290,4 +290,34 @@ suite('Map', () => {
 		assert.equal(map.findSubstr('/user/foo/bar/far/boo'), 1);
 
 	});
+
+	test('TrieMap - lookup', function () {
+
+		const map = new TrieMap<number>(TrieMap.PathSplitter);
+		map.insert('/user/foo/bar', 1);
+		map.insert('/user/foo', 2);
+		map.insert('/user/foo/flip/flop', 3);
+
+		assert.equal(map.lookUp('/foo'), undefined);
+		assert.equal(map.lookUp('/user'), undefined);
+		assert.equal(map.lookUp('/user/foo'), 2);
+		assert.equal(map.lookUp('/user/foo/bar'), 1);
+		assert.equal(map.lookUp('/user/foo/bar/boo'), undefined);
+	});
+
+	test('TrieMap - superstr', function () {
+
+		const map = new TrieMap<number>(TrieMap.PathSplitter);
+		map.insert('/user/foo/bar', 1);
+		map.insert('/user/foo', 2);
+		map.insert('/user/foo/flip/flop', 3);
+
+		const supMap = map.findSuperstr('/user');
+
+		assert.equal(supMap.lookUp('foo'), 2);
+		assert.equal(supMap.lookUp('foo/bar'), 1);
+		assert.equal(supMap.lookUp('foo/flip/flop'), 3);
+		assert.equal(supMap.lookUp('foo/flip/flop/bar'), undefined);
+		assert.equal(supMap.lookUp('user'), undefined);
+	});
 });
