@picosearch/trie/dist/trie.js

Simple, zero dependency, type-safe implementation of a trie data structure.

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.Trie = exports.TrieMap = void 0;
const constants_1 = require("./constants");
const util_1 = require("./util");
const getNewEmptyNode = () => [
    Object.create(null),
];
class TrieMap {
    root;
    constructor(root = getNewEmptyNode()) {
        this.root = root;
    }
    /**
     * Insert values into the TrieMap for the given key. You can pass multiple values.
     *
     * @param key - The key to insert.
     * @param values - The values associated with the key.
     */
    insert(key, ...values) {
        if (!key?.length)
            return;
        let currentNode = this.root;
        for (const part of key.split('')) {
            if (currentNode[0][part]) {
                currentNode = currentNode[0][part];
                continue;
            }
            currentNode[0][part] = getNewEmptyNode();
            currentNode = currentNode[0][part];
        }
        currentNode[1] ??= [];
        currentNode[1].push(...values);
    }
    /**
     * Lookup the values associated with the given key.
     *
     * @param key - The key to lookup.
     * @returns An array of values associated with the given key, or `null` if the key is not found.
     */
    lookup(query) {
        if (!query?.length)
            return null;
        let currentNode = this.root;
        for (const part of query.split('')) {
            if (!currentNode[0][part]) {
                return null;
            }
            currentNode = currentNode[0][part];
        }
        return currentNode[1] ?? null;
    }
    getFuzzyMatches(query, options) {
        if (!query?.length)
            return [];
        if (options?.maxErrors === 0) {
            // this case equals a regular non-fuzzy search...
            const values = this.lookup(query);
            if (values === null)
                return [];
            return options?.includeValues ? [[query, values]] : [query];
        }
        const maxErrors = options?.maxErrors ?? constants_1.FUZZY_SEARCH_DEFAULT_MAX_ERROR;
        const result = [];
        const stack = [
            {
                prefix: '',
                // the first row in the matrix is just the sequence 0,1,..,sequence.length
                prevRow: Array.from({ length: query.length + 1 }, (_, i) => i),
                node: this.root,
            },
        ];
        // The following traverses the Trie while computing the Edit distance for each path using dynamic
        // programming.
        // see https://repositorio.uchile.cl/bitstream/handle/2250/126168/Navarro_Gonzalo_Guided_tour.pdf (page 17)
        while (true) {
            const currentItem = stack.shift();
            if (!currentItem)
                break;
            const { prefix, prevRow, node } = currentItem;
            for (const item of Object.entries(node[0])) {
                const [char, childNode] = item;
                const newPrefix = prefix + char;
                const newRow = [newPrefix.length];
                let allValuesInRowExceedMaxError = newRow[0] > maxErrors;
                for (let j = 1; j <= query.length; j++) {
                    const charDist = query[j - 1] === char ? 0 : 1;
                    const value = Math.min(prevRow[j - 1] + charDist, prevRow[j] + 1, newRow[j - 1] + 1);
                    newRow.push(value);
                    if (allValuesInRowExceedMaxError && value <= maxErrors) {
                        allValuesInRowExceedMaxError = false;
                    }
                }
                if (childNode?.[1] && newRow[newRow.length - 1] <= maxErrors) {
                    const distance = newRow[newRow.length - 1];
                    (0, util_1.sortedInsert)(result, [newPrefix, distance, childNode[1]], (a, b) => {
                        const diff = a[1] - b[1];
                        return diff === 0 ? a[0].localeCompare(b[0]) : diff;
                    });
                    if (options?.limit && result.length >= options.limit)
                        break;
                }
                if (!allValuesInRowExceedMaxError) {
                    stack.push({
                        prefix: newPrefix,
                        prevRow: newRow,
                        node: childNode,
                    });
                }
            }
        }
        return options?.includeValues
            ? result.map(([word, distance, values]) => [word, values])
            : result.map(([word]) => word);
    }
    toJSON() {
        return JSON.stringify(this.root);
    }
    static fromJSON(jsonStr) {
        return new TrieMap(JSON.parse(jsonStr));
    }
}
exports.TrieMap = TrieMap;
class Trie {
    trieMap;
    constructor(options) {
        this.trieMap = new TrieMap(options?.root ?? getNewEmptyNode());
    }
    /**
     * Insert a key into the Trie.
     *
     * @param key - The key to insert.
     */
    insert(key) {
        if (this.has(key))
            return;
        this.trieMap.insert(key, 1);
    }
    /**
     * Check if the Trie contains the given key.
     *
     * @param key - The key to check.
     * @returns `true` if the Trie contains the key, `false` otherwise.
     */
    has(key) {
        return this.trieMap.lookup(key) !== null;
    }
    /**
     * Get all words that match the given query.
     *
     * @param query - The query to match.
     * @param options - Optional options to control the behavior of the function.
     * @returns An array of words that match the given query.
     */
    getFuzzyMatches(query, options) {
        return this.trieMap.getFuzzyMatches(query, {
            ...options,
            includeValues: false,
        });
    }
    toJSON() {
        return this.trieMap.toJSON();
    }
    static fromJSON(jsonStr) {
        return new Trie({ root: JSON.parse(jsonStr) });
    }
}
exports.Trie = Trie;