Autocomplete a search bar

"""
Suggest as the user types: a prefix tree, ordered by how often a term is searched.

Rung N0. Standard library only, and the whole index is a nest of dictionaries
that fits in the memory of the process serving the search bar.

Two decisions carry the approach.

First, the tree is keyed on a normalised spelling, accents folded and case
dropped, while each leaf keeps the original one. Someone typing "ec" finds
"écharpe", and still reads it spelled properly in the drop-down.

Second, the ordering is a plain sort on the usage count. Suggesting is not
retrieving: ten candidates under a prefix is a common case, and sorting ten
items at every keystroke costs nothing worth optimising.
"""

import unicodedata

# Marks the terms that end at a node. A character can never collide with it.
END = "\0"


def normalise(text: str) -> str:
    """Fold case and strip accents, so that "ec" reaches "écharpe"."""
    decomposed = unicodedata.normalize("NFD", text.casefold())
    return "".join(c for c in decomposed if not unicodedata.combining(c))


def build(entries) -> dict:
    """
    Build the tree from pairs of (term, how often it was searched).

    Terms sharing a normalised spelling are kept side by side at the same
    leaf, rather than one silently replacing the other.
    """
    root: dict = {}
    for term, count in entries:
        node = root
        for char in normalise(term):
            node = node.setdefault(char, {})
        node.setdefault(END, []).append((count, term))
    return root


def _descend(root: dict, prefix: str):
    """Walk down to the node holding everything that starts with `prefix`."""
    node = root
    for char in prefix:
        node = node.get(char)
        if node is None:
            return None
    return node


def _collect(node: dict):
    """Every (count, term) stored under a node, in no particular order."""
    for key, value in node.items():
        if key == END:
            yield from value
        else:
            yield from _collect(value)


def suggest(root: dict, prefix: str, limit: int = 5) -> list[str]:
    """
    The most searched terms starting with `prefix`, most searched first.

    An empty prefix returns the most searched terms overall, which is what an
    empty search bar should offer. An unknown prefix returns nothing: the tree
    answers about the characters it was given, not about the ones it guesses.
    """
    node = _descend(root, normalise(prefix))
    if node is None:
        return []
    found = sorted(_collect(node), key=lambda pair: (-pair[0], pair[1]))
    return [term for _, term in found[:limit]]

/**
 * Suggest as the user types: a prefix tree, ordered by how often a term is searched.
 *
 * Rung N0. No dependency, and the whole index is a nest of maps that fits in
 * the memory of the process serving the search bar.
 *
 * Two decisions carry the approach.
 *
 * First, the tree is keyed on a normalised spelling, accents folded and case
 * dropped, while each leaf keeps the original one. Someone typing "ec" finds
 * "écharpe", and still reads it spelled properly in the drop-down.
 *
 * Second, the ordering is a plain sort on the usage count. Suggesting is not
 * retrieving: ten candidates under a prefix is a common case, and sorting ten
 * items at every keystroke costs nothing worth optimising.
 */

// Marks the terms that end at a node. A character can never collide with it.
const END = Symbol('term');

/** Fold case and strip accents, so that "ec" reaches "écharpe". */
export function normalise(text) {
  return text.normalize('NFD').replace(/\p{M}/gu, '').toLowerCase();
}

/**
 * Build the tree from pairs of [term, how often it was searched].
 *
 * Terms sharing a normalised spelling are kept side by side at the same leaf,
 * rather than one silently replacing the other.
 */
export function build(entries) {
  const root = new Map();
  for (const [term, count] of entries) {
    let node = root;
    for (const char of normalise(term)) {
      if (!node.has(char)) node.set(char, new Map());
      node = node.get(char);
    }
    if (!node.has(END)) node.set(END, []);
    node.get(END).push([count, term]);
  }
  return root;
}

/** Walk down to the node holding everything that starts with `prefix`. */
function descend(root, prefix) {
  let node = root;
  for (const char of prefix) {
    node = node.get(char);
    if (node === undefined) return undefined;
  }
  return node;
}

/** Every [count, term] stored under a node, in no particular order. */
function collect(node, found = []) {
  for (const [key, value] of node) {
    if (key === END) found.push(...value);
    else collect(value, found);
  }
  return found;
}

/**
 * The most searched terms starting with `prefix`, most searched first.
 *
 * An empty prefix returns the most searched terms overall, which is what an
 * empty search bar should offer. An unknown prefix returns nothing: the tree
 * answers about the characters it was given, not about the ones it guesses.
 */
export function suggest(root, prefix, limit = 5) {
  const node = descend(root, normalise(prefix));
  if (node === undefined) return [];
  const found = collect(node);
  found.sort((a, b) => b[0] - a[0] || a[1].localeCompare(b[1]));
  return found.slice(0, limit).map(([, term]) => term);
}

"""
Reorder the suggestions with what people actually clicked.

Rung N1. The prefix tree of N0 ranks candidates by how often a term is
searched. That count says what people looked for, not what they picked once
the drop-down opened. Clicks say the second, and they are already in the logs.

The model is a count, not a gradient: for every prefix that was ever typed,
how many times each suggestion was chosen. It trains in one pass over the log
and is read back with a dictionary lookup, which is what a suggestion budget
of a keystroke allows.

The candidates come in as an argument: this rung reorders a list, it does not
retrieve it.
"""

import unicodedata


def normalise(text: str) -> str:
    """Same folding as the prefix tree, so both rungs agree on what was typed."""
    decomposed = unicodedata.normalize("NFD", text.casefold())
    return "".join(c for c in decomposed if not unicodedata.combining(c))


def learn(clicks) -> dict:
    """
    Count clicks from pairs of (what was typed, which suggestion was clicked).

    One click teaches something about every prefix of what was typed: whoever
    chose "chaussettes de sport" after typing "chau" also tells us what to
    show at "c" and at "cha".
    """
    model: dict = {}
    for typed, term in clicks:
        typed = normalise(typed)
        for length in range(len(typed) + 1):
            key = (typed[:length], term)
            model[key] = model.get(key, 0) + 1
    return model


def _evidence(model: dict, prefix: str, term: str) -> int:
    """
    Clicks recorded for the longest prefix of the query that saw this term.

    Backing off matters: a rare prefix has too few clicks of its own, but it
    shares its first letters with hundreds of past queries that do. The longer
    the matching prefix, the more specific the evidence, hence the weight.
    """
    for length in range(len(prefix), -1, -1):
        clicked = model.get((prefix[:length], term), 0)
        if clicked:
            return clicked * (length + 1)
    return 0


def rerank(model: dict, prefix: str, candidates, limit: int = 5) -> list[str]:
    """
    Sort candidates by past clicks, keeping their incoming order as tie-break.

    Candidates arrive ordered by search frequency, as the previous rung left
    them. A term nobody ever clicked keeps that order: the model only moves
    what it has evidence about, which is what makes it safe to ship on a log
    that is still thin.
    """
    prefix = normalise(prefix)
    ranked = sorted(
        enumerate(candidates),
        key=lambda pair: (-_evidence(model, prefix, pair[1]), pair[0]),
    )
    return [term for _, term in ranked[:limit]]

/**
 * Reorder the suggestions with what people actually clicked.
 *
 * Rung N1. The prefix tree of N0 ranks candidates by how often a term is
 * searched. That count says what people looked for, not what they picked once
 * the drop-down opened. Clicks say the second, and they are already in the
 * logs.
 *
 * The model is a count, not a gradient: for every prefix that was ever typed,
 * how many times each suggestion was chosen. It trains in one pass over the
 * log and is read back with a map lookup, which is what a suggestion budget of
 * a few milliseconds per keystroke allows.
 *
 * The candidates come in as an argument: this rung reorders a list, it does
 * not retrieve it.
 */

/** Same folding as the prefix tree, so both rungs agree on what was typed. */
export function normalise(text) {
  return text.normalize('NFD').replace(/\p{M}/gu, '').toLowerCase();
}

// A map key has to be a single value, and a tab never occurs inside a prefix.
const key = (prefix, term) => `${prefix}\t${term}`;

/**
 * Count clicks from pairs of [what was typed, which suggestion was clicked].
 *
 * One click teaches something about every prefix of what was typed: whoever
 * chose "chaussettes de sport" after typing "chau" also tells us what to show
 * at "c" and at "cha".
 */
export function learn(clicks) {
  const model = new Map();
  for (const [typed, term] of clicks) {
    const prefix = normalise(typed);
    for (let length = 0; length <= prefix.length; length += 1) {
      const at = key(prefix.slice(0, length), term);
      model.set(at, (model.get(at) ?? 0) + 1);
    }
  }
  return model;
}

/**
 * Clicks recorded for the longest prefix of the query that saw this term.
 *
 * Backing off matters: a rare prefix has too few clicks of its own, but it
 * shares its first letters with hundreds of past queries that do. The longer
 * the matching prefix, the more specific the evidence, hence the weight.
 */
function evidence(model, prefix, term) {
  for (let length = prefix.length; length >= 0; length -= 1) {
    const clicked = model.get(key(prefix.slice(0, length), term)) ?? 0;
    if (clicked) return clicked * (length + 1);
  }
  return 0;
}

/**
 * Sort candidates by past clicks, keeping their incoming order as tie-break.
 *
 * Candidates arrive ordered by search frequency, as the previous rung left
 * them. A term nobody ever clicked keeps that order: the model only moves what
 * it has evidence about, which is what makes it safe to ship on a log that is
 * still thin.
 */
export function rerank(model, prefix, candidates, limit = 5) {
  const typed = normalise(prefix);
  const scored = candidates.map((term, rank) => [evidence(model, typed, term), rank, term]);
  scored.sort((a, b) => b[0] - a[0] || a[1] - b[1]);
  return scored.slice(0, limit).map(([, , term]) => term);
}