471 lines
12 KiB
TypeScript
471 lines
12 KiB
TypeScript
import { SUPPORTED_LOCALES } from 'emojibase';
|
|
import type { CompactEmoji, Locale, ShortcodesDataset } from 'emojibase';
|
|
|
|
import type { ApiCustomEmojiJSON } from '@/mastodon/api_types/custom_emoji';
|
|
|
|
import { openEmojiDB } from './db-schema';
|
|
import type { Database } from './db-schema';
|
|
import { importEmojiData } from './loader';
|
|
import { localeToSegmenter, toSupportedLocale } from './locale';
|
|
import {
|
|
extractTokens,
|
|
skinHexcodeToEmoji,
|
|
transformCustomEmojiData,
|
|
transformEmojiData,
|
|
} from './normalize';
|
|
import type { AnyEmojiData, CacheKey, CustomEmojiData } from './types';
|
|
import { emojiLogger } from './utils';
|
|
|
|
const loadedLocales = new Set<Locale>();
|
|
|
|
const log = emojiLogger('database');
|
|
|
|
// Loads the database in a way that ensures it's only loaded once.
|
|
const loadDB = (() => {
|
|
let dbPromise: Promise<Database> | null = null;
|
|
|
|
// Actually load the DB.
|
|
async function initDB() {
|
|
const db = await openEmojiDB();
|
|
await syncLocales(db);
|
|
log('Loaded database version %d', db.version);
|
|
return db;
|
|
}
|
|
|
|
// Loads the database, or returns the existing promise if it hasn't resolved yet.
|
|
const loadPromise = async (): Promise<Database> => {
|
|
if (dbPromise) {
|
|
return dbPromise;
|
|
}
|
|
dbPromise = initDB();
|
|
return dbPromise;
|
|
};
|
|
// Special way to reset the database, used for unit testing.
|
|
loadPromise.reset = () => {
|
|
dbPromise = null;
|
|
};
|
|
return loadPromise;
|
|
})();
|
|
|
|
type ScoreMap = Map<string, AnyEmojiData & { score: number }>;
|
|
|
|
export async function search({
|
|
query,
|
|
locale: localeString,
|
|
limit = 0,
|
|
}: {
|
|
query: string;
|
|
locale: string;
|
|
limit?: number;
|
|
}) {
|
|
performance.mark('emoji-search-start');
|
|
|
|
// Get the locale, and extract tokens from the query.
|
|
const locale = await toLoadedLocale(localeString);
|
|
const segmenter = localeToSegmenter(locale);
|
|
const queryTokens = extractTokens(query, segmenter);
|
|
|
|
if (queryTokens.length === 0) {
|
|
log('no tokens extracted from query "%s"', query);
|
|
return [];
|
|
}
|
|
const lastToken = queryTokens.at(-1);
|
|
if (!lastToken) {
|
|
throw new Error('Missing tokens from query');
|
|
}
|
|
|
|
log('searching for tokens %o in locale %s', queryTokens, locale);
|
|
|
|
// Create an array of emoji results
|
|
const db = await loadDB();
|
|
const resultArrays: ScoreMap[] = [];
|
|
const existingCustomShortcodes = new Set<string>();
|
|
|
|
for (let i = 0; i < queryTokens.length; i++) {
|
|
const token = queryTokens[i];
|
|
if (!token) continue;
|
|
|
|
// Only query the range for the last token to allow partial matches.
|
|
const range =
|
|
i === queryTokens.length - 1
|
|
? IDBKeyRange.lowerBound(token)
|
|
: IDBKeyRange.only(token);
|
|
|
|
const [unicodeResults, customResults, shortcodeResults] = await Promise.all(
|
|
[
|
|
db.getAllFromIndex(locale, 'tokens', range),
|
|
db.getAllFromIndex('custom', 'tokens', range),
|
|
db.getAllFromIndex('shortcodes', 'shortcodes', range),
|
|
],
|
|
);
|
|
const resultMap: ScoreMap = new Map();
|
|
|
|
for (const emoji of unicodeResults) {
|
|
const score = getScoreForEmoji(emoji, token);
|
|
if (score === null) {
|
|
continue;
|
|
}
|
|
resultMap.set(emoji.hexcode, { ...emoji, score });
|
|
}
|
|
|
|
for (const emoji of customResults) {
|
|
const score = getScoreForEmoji(emoji, token);
|
|
if (score === null) {
|
|
continue;
|
|
}
|
|
existingCustomShortcodes.add(emoji.shortcode);
|
|
resultMap.set(emoji.shortcode, { ...emoji, score });
|
|
}
|
|
|
|
for (const shortcodeResult of shortcodeResults) {
|
|
if (resultMap.has(shortcodeResult.hexcode)) {
|
|
continue;
|
|
}
|
|
const emoji = await db.get(locale, shortcodeResult.hexcode);
|
|
if (!emoji) {
|
|
continue;
|
|
}
|
|
// Score the emoji with the legacy shortcode, even though it's not part of the emoji.
|
|
const score = getScoreForEmoji(
|
|
{
|
|
...emoji,
|
|
shortcodes: [...shortcodeResult.shortcodes, ...emoji.shortcodes],
|
|
},
|
|
token,
|
|
);
|
|
if (score === null) {
|
|
continue;
|
|
}
|
|
resultMap.set(emoji.hexcode, { ...emoji, score });
|
|
}
|
|
|
|
log('found %d results for token "%s"', resultMap.size, token);
|
|
resultArrays.push(resultMap);
|
|
}
|
|
|
|
// Utilize maps to find the intersection of all result sets.
|
|
const results = Array.from(
|
|
resultArrays
|
|
.reduce((prev, curr) => {
|
|
const intersection: ScoreMap = new Map();
|
|
for (const [code, emoji] of prev) {
|
|
if (curr.has(code)) {
|
|
intersection.set(code, emoji);
|
|
}
|
|
}
|
|
return intersection;
|
|
})
|
|
.values(),
|
|
);
|
|
|
|
// If there are no results, try a cursor-based custom emoji search instead.
|
|
if (results.length === 0 || results.length < limit) {
|
|
const customEmojisFound = await fullCustomSearch(
|
|
query,
|
|
existingCustomShortcodes,
|
|
);
|
|
if (customEmojisFound.length > 0) {
|
|
log(
|
|
'cursor search found %d results for "%s"',
|
|
customEmojisFound.length,
|
|
query,
|
|
);
|
|
results.push(...customEmojisFound);
|
|
}
|
|
}
|
|
|
|
// Sort by score, descending.
|
|
results.sort((a, b) => a.score - b.score);
|
|
|
|
const time = performance.measure('emoji-search-end', 'emoji-search-start');
|
|
log(
|
|
'search for "%s" in locale %s returned %d results and took %dms',
|
|
query,
|
|
locale,
|
|
results.length,
|
|
time.duration,
|
|
);
|
|
if (limit > 0) {
|
|
return results.slice(0, limit);
|
|
}
|
|
return results;
|
|
}
|
|
|
|
function getScoreForEmoji(
|
|
emoji: AnyEmojiData,
|
|
query: string,
|
|
checkTokens = true,
|
|
) {
|
|
const id = 'shortcode' in emoji ? emoji.shortcode : emoji.label;
|
|
if (id === query) {
|
|
return 0;
|
|
}
|
|
|
|
let index = 1;
|
|
const searchTokens = [id];
|
|
if (checkTokens) {
|
|
// Check shortcodes before tokens as they are more important.
|
|
if ('shortcodes' in emoji) {
|
|
searchTokens.push(...emoji.shortcodes);
|
|
}
|
|
searchTokens.push(...emoji.tokens);
|
|
}
|
|
for (const token of searchTokens) {
|
|
const tokenIndex = token.indexOf(query);
|
|
if (tokenIndex !== -1) {
|
|
return index + tokenIndex / token.length;
|
|
}
|
|
index++;
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
async function fullCustomSearch(query: string, existing = new Set<string>()) {
|
|
const db = await loadDB();
|
|
const trx = db.transaction('custom', 'readonly');
|
|
const foundEmojis = new Set<string>();
|
|
|
|
// First iterate over chunks of 1,000 custom emoji keys and find any matches.
|
|
const chunkSize = 1_000;
|
|
let lastKey: string | null = null;
|
|
let keys: string[] = [];
|
|
do {
|
|
const keyRange = lastKey ? IDBKeyRange.lowerBound(lastKey, true) : null;
|
|
keys = await trx.store.getAllKeys(keyRange, chunkSize);
|
|
|
|
if (keys.length === 0) {
|
|
break;
|
|
}
|
|
log('cursor search got batch of %d emojis', keys.length);
|
|
lastKey = keys.at(-1) ?? null;
|
|
|
|
for (const key of keys) {
|
|
if (!foundEmojis.has(key) && !existing.has(key) && key.includes(query)) {
|
|
foundEmojis.add(key);
|
|
}
|
|
}
|
|
} while (keys.length === chunkSize);
|
|
|
|
// Next get the full emojis for all matches.
|
|
const emojis = await Promise.all(
|
|
foundEmojis.keys().map((key) => trx.store.get(key)),
|
|
);
|
|
const results: (CustomEmojiData & { score: number })[] = [];
|
|
for (const emoji of emojis) {
|
|
if (emoji) {
|
|
const score = getScoreForEmoji(emoji, query, false);
|
|
if (score && score > 0) {
|
|
results.push({
|
|
score,
|
|
...emoji,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
return results;
|
|
}
|
|
|
|
export async function putEmojiData(emojis: CompactEmoji[], locale: Locale) {
|
|
loadedLocales.add(locale);
|
|
const db = await loadDB();
|
|
const trx = db.transaction(locale, 'readwrite');
|
|
await trx.store.clear();
|
|
const segmenter = localeToSegmenter(locale);
|
|
await Promise.all(
|
|
emojis
|
|
.sort((a, b) => (a.order ?? 0) - (b.order ?? 0))
|
|
.map((emoji) => trx.store.put(transformEmojiData(emoji, segmenter))),
|
|
);
|
|
await trx.done;
|
|
}
|
|
|
|
export async function putCustomEmojiData({
|
|
emojis,
|
|
clear = false,
|
|
}: {
|
|
emojis: ApiCustomEmojiJSON[];
|
|
clear?: boolean;
|
|
}) {
|
|
const db = await loadDB();
|
|
const trx = db.transaction('custom', 'readwrite');
|
|
|
|
// When importing from the API, clear everything first.
|
|
if (clear) {
|
|
await trx.store.clear();
|
|
log('Cleared existing custom emojis in database');
|
|
}
|
|
|
|
await Promise.all(
|
|
emojis.map((emoji) => trx.store.put(transformCustomEmojiData(emoji))),
|
|
);
|
|
await trx.done;
|
|
|
|
log('Imported %d custom emojis into database', emojis.length);
|
|
}
|
|
|
|
export async function putLegacyShortcodes(shortcodes: ShortcodesDataset) {
|
|
const db = await loadDB();
|
|
const trx = db.transaction('shortcodes', 'readwrite');
|
|
await Promise.all(
|
|
Object.entries(shortcodes).map(([hexcode, codes]) =>
|
|
trx.store.put({
|
|
hexcode,
|
|
shortcodes: Array.isArray(codes) ? codes : [codes],
|
|
}),
|
|
),
|
|
);
|
|
await trx.done;
|
|
}
|
|
|
|
export async function loadCacheValue(key: CacheKey) {
|
|
const db = await loadDB();
|
|
const value = await db.get('etags', key);
|
|
return value;
|
|
}
|
|
|
|
export async function putCacheValue(key: CacheKey, value: string) {
|
|
const db = await loadDB();
|
|
await db.put('etags', value, key);
|
|
}
|
|
|
|
export async function clearCache(key: CacheKey) {
|
|
const db = await loadDB();
|
|
await db.delete('etags', key);
|
|
log('Cleared cache for %s', key);
|
|
}
|
|
|
|
export async function resetDatabase() {
|
|
const db = await loadDB();
|
|
const storeNames = [...db.objectStoreNames];
|
|
await Promise.all(storeNames.map((storeName) => db.clear(storeName)));
|
|
log(storeNames, 'Reset emoji database stores:');
|
|
}
|
|
|
|
export async function loadEmojiByHexcode(
|
|
hexcode: string,
|
|
localeString: string,
|
|
) {
|
|
const db = await loadDB();
|
|
const locale = await toLoadedLocale(localeString);
|
|
const result = await db.get(locale, hexcode);
|
|
if (result) {
|
|
return result;
|
|
}
|
|
|
|
// If the emoji wasn't found, check if it's a skin tone variant.
|
|
const skinResult = await db.getFromIndex(
|
|
locale,
|
|
'skinHexcodes',
|
|
IDBKeyRange.only(hexcode),
|
|
);
|
|
|
|
if (!skinResult) {
|
|
return skinResult;
|
|
}
|
|
|
|
// Reconstruct the full unicode string from the skin tone hexcode.
|
|
return skinHexcodeToEmoji(hexcode, skinResult);
|
|
}
|
|
|
|
export async function loadAllUnicodeEmojis(localeString: string) {
|
|
const locale = await toLoadedLocale(localeString);
|
|
const db = await loadDB();
|
|
return db.getAll(locale);
|
|
}
|
|
|
|
export async function loadCustomEmojiByShortcode(shortcode: string) {
|
|
const db = await loadDB();
|
|
return db.get('custom', shortcode);
|
|
}
|
|
|
|
export async function searchCustomEmojisByShortcodes(shortcodes: string[]) {
|
|
const db = await loadDB();
|
|
const sortedCodes = shortcodes.toSorted();
|
|
const results = await db.getAll(
|
|
'custom',
|
|
IDBKeyRange.bound(sortedCodes.at(0), sortedCodes.at(-1)),
|
|
);
|
|
return results.filter((emoji) => shortcodes.includes(emoji.shortcode));
|
|
}
|
|
|
|
export async function loadAllCustomEmoji() {
|
|
const db = await loadDB();
|
|
const cacheValue = await db.get('etags', 'custom');
|
|
if (!cacheValue) {
|
|
return null;
|
|
}
|
|
return db.getAll('custom');
|
|
}
|
|
|
|
export async function loadLegacyShortcodesByShortcode(shortcode: string) {
|
|
const db = await loadDB();
|
|
return db.getFromIndex(
|
|
'shortcodes',
|
|
'shortcodes',
|
|
IDBKeyRange.only(shortcode),
|
|
);
|
|
}
|
|
|
|
export async function loadAllShortcodes() {
|
|
const db = await loadDB();
|
|
return db.getAll('shortcodes');
|
|
}
|
|
|
|
// Private functions
|
|
|
|
async function syncLocales(db: Database) {
|
|
const locales = await Promise.all(
|
|
SUPPORTED_LOCALES.map(
|
|
async (locale) =>
|
|
[locale, await hasLocale(locale, db)] satisfies [Locale, boolean],
|
|
),
|
|
);
|
|
for (const [locale, loaded] of locales) {
|
|
if (loaded) {
|
|
loadedLocales.add(locale);
|
|
} else {
|
|
loadedLocales.delete(locale);
|
|
}
|
|
}
|
|
log('Loaded %d locales: %o', loadedLocales.size, loadedLocales);
|
|
}
|
|
|
|
async function toLoadedLocale(localeString: string) {
|
|
const locale = toSupportedLocale(localeString);
|
|
if (localeString !== locale) {
|
|
log(`Locale ${locale} is different from provided ${localeString}`);
|
|
}
|
|
if (!loadedLocales.has(locale)) {
|
|
log('Locale %s not loaded, importing...', locale);
|
|
await importEmojiData(locale);
|
|
return locale;
|
|
}
|
|
return locale;
|
|
}
|
|
|
|
export class LocaleNotLoadedError extends Error {
|
|
constructor(locale: Locale) {
|
|
super(`Locale ${locale} is not loaded in emoji database`);
|
|
this.name = 'LocaleNotLoadedError';
|
|
}
|
|
}
|
|
|
|
async function hasLocale(locale: Locale, db: Database): Promise<boolean> {
|
|
if (loadedLocales.has(locale)) {
|
|
return true;
|
|
}
|
|
const rowCount = await db.count(locale);
|
|
return !!rowCount;
|
|
}
|
|
|
|
// Testing helpers
|
|
export async function testGet() {
|
|
const db = await loadDB();
|
|
return { db, loadedLocales };
|
|
}
|
|
export function testClear() {
|
|
loadedLocales.clear();
|
|
loadDB.reset();
|
|
}
|