// Client-side tag completion. import { UniqueHeap } from './unique-heap'; import store from './store'; export interface Result { aliasName: string; name: string; imageCount: number; } /** * Opaque, unique pointer to tag data. */ type TagPointer = number; /** * Numeric index of a tag in its primary order. */ type TagReferenceIndex = number; /** * Compare two UTF-8 strings, C-style. */ function strcmp(a: Uint8Array, b: Uint8Array): number { const aLength = a.length; const bLength = b.length; let index = 0; while (index < aLength && index < bLength && a[index] === b[index]) { index++; } const aValue = index >= aLength ? 0 : a[index]; const bValue = index >= bLength ? 0 : b[index]; return aValue - bValue; } const namespaceSeparator = ':'.charCodeAt(0); /** * Returns the name of a tag without any namespace component. */ function nameInNamespace(s: Uint8Array): Uint8Array { const first = s.indexOf(namespaceSeparator); if (first !== -1) { return s.slice(first + 1); } return s; } /** * See lib/philomena/autocomplete.ex for binary structure details. * * A binary blob is used to avoid the creation of large amounts of garbage on * the JS heap and speed up the execution of the search. */ export class LocalAutocompleter { private encoder: TextEncoder; private decoder: TextDecoder; private data: Uint8Array; private view: DataView; private numTags: number; private referenceStart: number; private secondaryStart: number; private formatVersion: number; /** * Build a new local autocompleter. */ constructor(backingStore: ArrayBuffer) { this.encoder = new TextEncoder(); this.decoder = new TextDecoder(); this.data = new Uint8Array(backingStore); this.view = new DataView(backingStore); this.numTags = this.view.getUint32(backingStore.byteLength - 4, true); this.referenceStart = this.view.getUint32(backingStore.byteLength - 8, true); this.secondaryStart = this.referenceStart + 8 * this.numTags; this.formatVersion = this.view.getUint32(backingStore.byteLength - 12, true); if (this.formatVersion !== 2) { throw new Error('Incompatible autocomplete format version'); } } /** * Return the pointer to tag data for the given reference index. */ private resolveTagReference(i: TagReferenceIndex, resolveAlias: boolean = true): TagPointer { const tagPointer = this.view.getUint32(this.referenceStart + i * 8, true); const imageCount = this.view.getInt32(this.referenceStart + i * 8 + 4, true); if (resolveAlias && imageCount < 0) { // This is actually an alias, so follow it return this.resolveTagReference(-imageCount - 1); } return tagPointer; } /** * Get the images count for the given reference index. */ private getImageCount(i: TagReferenceIndex): number { const imageCount = this.view.getInt32(this.referenceStart + i * 8 + 4, true); if (imageCount < 0) { // This is actually an alias, so follow it return this.getImageCount(-imageCount - 1); } return imageCount; } /** * Return the name buffer of the pointed-to result. */ private referenceToName(i: TagReferenceIndex, resolveAlias: boolean = true): Uint8Array { const pointer = this.resolveTagReference(i, resolveAlias); const nameLength = this.view.getUint8(pointer); return this.data.slice(pointer + 1, pointer + nameLength + 1); } /** * Return whether any associations in the pointed-to result are in comparisonValues. */ private isFilteredByReference(comparisonValues: Set, i: TagReferenceIndex): boolean { const pointer = this.resolveTagReference(i); const nameLength = this.view.getUint8(pointer); const assnLength = this.view.getUint8(pointer + 1 + nameLength); for (let j = 0; j < assnLength; j++) { const assnValue = this.view.getUint32(pointer + 1 + nameLength + 1 + j * 4, true); if (comparisonValues.has(assnValue)) { return true; } } return false; } /** * Return whether Result a is considered less than Result b. */ private compareReferenceToReference(a: TagReferenceIndex, b: TagReferenceIndex): number { const imagesA = this.getImageCount(a); const imagesB = this.getImageCount(b); if (imagesA !== imagesB) { return imagesA - imagesB; } const nameA = this.referenceToName(a); const nameB = this.referenceToName(a); return strcmp(nameA, nameB); } /** * Get a Result object as the ith tag inside the file, secondary ordering. */ private getSecondaryResultAt(i: number): TagReferenceIndex { return this.view.getUint32(this.secondaryStart + i * 4, true); } /** * Perform a binary search to fetch all results matching a condition. */ private scanResults( getResult: (i: number) => TagReferenceIndex, compare: (result: TagReferenceIndex) => number, hasFilteredAssociation: (result: TagReferenceIndex) => boolean, results: UniqueHeap, ) { const filter = !store.get('unfilter_tag_suggestions'); let min = 0; let max = this.numTags; while (min < max - 1) { const med = min + (((max - min) / 2) | 0); const referenceIndex = getResult(med); if (compare(referenceIndex) >= 0) { // too large, go left max = med; } else { // too small, go right min = med; } } // Scan forward until no more matches occur while (min < this.numTags - 1) { const referenceIndex = getResult(++min); if (compare(referenceIndex) !== 0) { break; } // Check if any associations are filtered if (filter && hasFilteredAssociation(referenceIndex)) { continue; } // Nothing was filtered, so add results.append(referenceIndex); } } /** * Find the top K results by image count which match the given string prefix. */ matchPrefix(prefixStr: string, k: number): Result[] { if (prefixStr.length === 0) { return []; } // Set up binary matching context const prefix = this.encoder.encode(prefixStr); const results = new UniqueHeap( this.compareReferenceToReference.bind(this), this.resolveTagReference.bind(this), new Uint32Array(this.numTags), ); // Set up filter context const hiddenTags = new Set(window.booru.hiddenTagList); const hasFilteredAssociation = this.isFilteredByReference.bind(this, hiddenTags); // Find tags ordered by full name const prefixMatch = (i: TagReferenceIndex) => strcmp(this.referenceToName(i, false).slice(0, prefix.length), prefix); const referenceToNameIndex = (i: number) => i; this.scanResults(referenceToNameIndex, prefixMatch, hasFilteredAssociation, results); // Find tags ordered by name in namespace const namespaceMatch = (i: TagReferenceIndex) => strcmp(nameInNamespace(this.referenceToName(i, false)).slice(0, prefix.length), prefix); const referenceToAliasIndex = this.getSecondaryResultAt.bind(this); this.scanResults(referenceToAliasIndex, namespaceMatch, hasFilteredAssociation, results); // Convert top K from heap into result array return results.topK(k).map((i: TagReferenceIndex) => ({ aliasName: this.decoder.decode(this.referenceToName(i, false)), name: this.decoder.decode(this.referenceToName(i)), imageCount: this.getImageCount(i), })); } }