diff --git a/trie-map.d.ts b/trie-map.d.ts index af889890..b0833042 100644 --- a/trie-map.d.ts +++ b/trie-map.d.ts @@ -13,6 +13,7 @@ export default class TrieMap implements Iterable<[K, V]> { // Methods clear(): void; set(prefix: K, value: V): this; + update(prefix: K, updateFunction: (oldValue: V | undefined) => V): this get(prefix: K): V; delete(prefix: K): boolean; has(prefix: K): boolean; diff --git a/trie-map.js b/trie-map.js index 1d99d059..e395b428 100644 --- a/trie-map.js +++ b/trie-map.js @@ -66,6 +66,32 @@ TrieMap.prototype.set = function(prefix, value) { return this; }; +/** + * Method used to update the value of the given prefix in the trie. + * + * @param {string|array} prefix - Prefix to follow. + * @param {(oldValue: any | undefined) => any} updateFunction - Update value visitor callback. + * @return {TrieMap} + */ +TrieMap.prototype.set = function(prefix, updateFunction) { + var node = this.root, + token; + + for (var i = 0, l = prefix.length; i < l; i++) { + token = prefix[i]; + + node = node[token] || (node[token] = {}); + } + + // Do we need to increase size? + if (!(SENTINEL in node)) + this.size++; + + node[SENTINEL] = updateFunction(node[SENTINEL]); + + return this; +}; + /** * Method used to return the value sitting at the end of the given prefix or * undefined if none exist.