From 8131b2bee914e0c1626b13d4f63353f57ae5cca4 Mon Sep 17 00:00:00 2001 From: Richard R Date: Tue, 9 Jun 2026 14:57:27 -0600 Subject: [PATCH] refactor(epub): unify word-to-region alignment for text highlighting Replace EPUB-specific word highlight mapping with shared token-sequence alignment logic used by HTML and PDF viewers. Remove bespoke alignment code and tests in favor of a single primitive (`locateAlignmentWordSpans`) that robustly maps spoken words to rendered text regions, tolerant of transcription and formatting differences. Update EPUB highlighting to cache per-segment word-region spans for efficient re-use. Add unit tests for the new alignment logic. --- src/hooks/epub/useEPUBHighlighting.ts | 50 +++++++---- src/lib/client/epub/epub-canonical-window.ts | 5 +- .../client/epub/epub-rendered-text-maps.ts | 4 + src/lib/client/epub/epub-word-highlight.ts | 22 ----- src/lib/client/highlight-token-alignment.ts | 41 ++++++++- src/lib/client/html/highlight.ts | 84 +++++++++++-------- tests/unit/epub-word-highlight.vitest.spec.ts | 50 ----------- .../unit/highlight-word-locate.vitest.spec.ts | 73 ++++++++++++++++ 8 files changed, 200 insertions(+), 129 deletions(-) delete mode 100644 src/lib/client/epub/epub-word-highlight.ts delete mode 100644 tests/unit/epub-word-highlight.vitest.spec.ts create mode 100644 tests/unit/highlight-word-locate.vitest.spec.ts diff --git a/src/hooks/epub/useEPUBHighlighting.ts b/src/hooks/epub/useEPUBHighlighting.ts index 4e09753..d186670 100644 --- a/src/hooks/epub/useEPUBHighlighting.ts +++ b/src/hooks/epub/useEPUBHighlighting.ts @@ -1,16 +1,17 @@ 'use client'; -import { useCallback, useEffect, type MutableRefObject, type RefObject } from 'react'; +import { useCallback, useEffect, useRef, type MutableRefObject, type RefObject } from 'react'; import type { Rendition } from 'epubjs'; -import { - resolveAlignmentWordSourceRange, -} from '@/lib/client/epub/epub-word-highlight'; import { createRangeFromMappedOffsets, resolveVisibleSegmentRange, type EpubRenderedTextMap, } from '@/lib/client/epub/epub-rendered-text-maps'; +import { + locateAlignmentWordSpans, + type AlignmentCharSpan, +} from '@/lib/client/highlight-token-alignment'; import type { CanonicalTtsSegment } from '@/lib/shared/tts-segment-plan'; import type { TTSSentenceAlignment } from '@/types/tts'; @@ -42,6 +43,10 @@ export function useEPUBHighlighting({ currentWordHighlightCfiRef, renderedTextMapsRef, }: UseEpubHighlightingParams): UseEpubHighlightingResult { + // Cache the per-segment word→region map so we don't re-align on every whisper + // tick. Keyed by the segment + resolved region + alignment. + const wordRangeCacheRef = useRef<{ key: string; spans: Array } | null>(null); + const clearWordHighlights = useCallback(() => { if (!renditionRef.current) return; if (currentWordHighlightCfiRef.current) { @@ -105,19 +110,32 @@ export function useEPUBHighlighting({ const resolved = resolveVisibleSegmentRange(renderedTextMapsRef.current, segment); if (!resolved || segment.startAnchor.sourceKey !== resolved.map.sourceKey) return; - // Native path only: the alignment's char offsets are authoritative for the - // spoken word (they live in the same canonical space as the segment text and - // the rendered char map). Clamp to the portion of the segment that is - // actually visible in this map so a word straddling a page/spread boundary - // still highlights its visible part instead of dropping the word entirely. - const alignmentRange = resolveAlignmentWordSourceRange(segment, words[wordIndex]); - if (!alignmentRange) return; + // Map each spoken word onto the rendered region with the shared token- + // sequence aligner (same primitive as the HTML and PDF viewers). The region + // text is the *rendered* text, so a returned span's offsets are already + // indices into the char map — no canonical-vs-rendered coordinate drift. + // Spans are relative to resolved.startOffset. + const regionText = resolved.map.text.slice(resolved.startOffset, resolved.endOffset); + const cacheKey = [ + segment.key, + resolved.map.sourceKey, + resolved.startOffset, + resolved.endOffset, + words.length, + ].join('::'); + if (wordRangeCacheRef.current?.key !== cacheKey) { + wordRangeCacheRef.current = { + key: cacheKey, + spans: locateAlignmentWordSpans(words, regionText), + }; + } - const clampedStart = Math.max(alignmentRange.sourceStart, resolved.startOffset); - const clampedEnd = Math.min(alignmentRange.sourceEnd, resolved.endOffset); - if (clampedEnd <= clampedStart) return; + const span = wordRangeCacheRef.current.spans[wordIndex]; + if (!span) return; - const wordRange = createRangeFromMappedOffsets(resolved.map, clampedStart, clampedEnd); + const absStart = resolved.startOffset + span.start; + const absEnd = resolved.startOffset + span.end; + const wordRange = createRangeFromMappedOffsets(resolved.map, absStart, absEnd); if (!wordRange) return; try { @@ -136,7 +154,7 @@ export function useEPUBHighlighting({ } ); } catch (error) { - console.error('Error highlighting EPUB word from alignment offsets:', error); + console.error('Error highlighting EPUB word:', error); } }, [ clearWordHighlights, diff --git a/src/lib/client/epub/epub-canonical-window.ts b/src/lib/client/epub/epub-canonical-window.ts index f003c57..b19ee29 100644 --- a/src/lib/client/epub/epub-canonical-window.ts +++ b/src/lib/client/epub/epub-canonical-window.ts @@ -247,9 +247,8 @@ export function materializeWindowSegments( next.startAnchor = { sourceKey: viewport.sourceKey, offset: clampToViewport(seg.startAnchor.offset) }; next.endAnchor = { sourceKey: viewport.sourceKey, offset: clampToViewport(seg.endAnchor.offset) }; // Keep ownerSourceKey in lock-step with the rewritten anchors: the word - // highlighter (resolveAlignmentWordSourceRange) requires - // startAnchor.sourceKey === ownerSourceKey to treat the segment as - // anchored in the rendered page. + // highlighter requires startAnchor.sourceKey === ownerSourceKey to treat + // the segment as anchored in the rendered page. next.ownerSourceKey = viewport.sourceKey; } out.push(next); diff --git a/src/lib/client/epub/epub-rendered-text-maps.ts b/src/lib/client/epub/epub-rendered-text-maps.ts index 82f3072..07dd131 100644 --- a/src/lib/client/epub/epub-rendered-text-maps.ts +++ b/src/lib/client/epub/epub-rendered-text-maps.ts @@ -14,6 +14,9 @@ type EpubMappedChar = MappedChar; export type EpubRenderedTextMap = { sourceKey: string; chars: EpubMappedPosition[]; + // Normalized rendered text; `text[i]` is the character at `chars[i]`. Used to + // locate spoken words by content within a resolved segment region. + text: string; content: { cfiFromRange: (range: Range) => string; }; @@ -94,6 +97,7 @@ export const buildRenderedTextMaps = ( maps.push({ sourceKey, chars: normalized.map((token) => token.pos), + text: normalized.map((token) => token.char).join(''), content, }); } catch { diff --git a/src/lib/client/epub/epub-word-highlight.ts b/src/lib/client/epub/epub-word-highlight.ts deleted file mode 100644 index ba5914b..0000000 --- a/src/lib/client/epub/epub-word-highlight.ts +++ /dev/null @@ -1,22 +0,0 @@ -import type { CanonicalTtsSegment } from '@/lib/shared/tts-segment-plan'; -import type { TTSSentenceWord } from '@/types/tts'; - -/** - * Resolve a spoken word's char offsets (from the Whisper alignment) into the - * canonical source-offset space the rendered text map is keyed by. The - * alignment's `charStart`/`charEnd` are offsets into the segment text, which is - * in the same normalized space as `segment.startAnchor.offset`. - */ -export const resolveAlignmentWordSourceRange = ( - segment: CanonicalTtsSegment, - word: TTSSentenceWord, -): { sourceStart: number; sourceEnd: number } | null => { - const { charStart, charEnd } = word; - if (!Number.isInteger(charStart) || !Number.isInteger(charEnd)) return null; - if (charStart < 0 || charEnd <= charStart || charEnd > segment.text.length) return null; - - return { - sourceStart: segment.startAnchor.offset + charStart, - sourceEnd: segment.startAnchor.offset + charEnd, - }; -}; diff --git a/src/lib/client/highlight-token-alignment.ts b/src/lib/client/highlight-token-alignment.ts index dd560bf..3b9c192 100644 --- a/src/lib/client/highlight-token-alignment.ts +++ b/src/lib/client/highlight-token-alignment.ts @@ -1,6 +1,6 @@ import { CmpStr } from 'cmpstr'; -import { normalizeUnicodeToken } from '@/lib/shared/language'; +import { normalizeUnicodeToken, segmentWords } from '@/lib/shared/language'; import type { TTSSentenceAlignment } from '@/types/tts'; const cmp = CmpStr.create().setMetric('dice').setFlags('itw'); @@ -10,6 +10,12 @@ export interface HighlightTokenRange { end: number; } +/** A spoken word's resolved span as half-open `[start, end)` char offsets into a region of text. */ +export interface AlignmentCharSpan { + start: number; + end: number; +} + export interface HighlightTokenMatchResult extends HighlightTokenRange { rating: number; lengthDiff: number; @@ -253,3 +259,36 @@ export function buildAlignmentTokenRanges( return ranges; } + +/** + * Map each spoken (Whisper-aligned) word onto a char span of an already-resolved + * region of rendered text. This is the single word→DOM mapping primitive shared + * by every viewer (EPUB, HTML/MD/TXT, PDF-style): the caller resolves the region + * (a sentence wrap, a visible segment range) and supplies its text; we return, + * for each word, the `[start, end)` char offsets to highlight within that text. + * + * The region text is tokenized into words and the spoken words are globally + * aligned against those tokens via {@link buildAlignmentTokenRanges}. Working in + * token space (not raw char offsets) means the result is robust to divergent + * transcription, punctuation, casing, and whitespace; `fillGaps` guarantees + * every word resolves to a neighboring token rather than dropping out. + */ +export function locateAlignmentWordSpans( + words: TTSSentenceAlignment['words'], + regionText: string, + language?: string | null, +): Array { + const tokens = segmentWords(regionText, language); + if (!words.length) return []; + if (!tokens.length) return words.map(() => null); + + const tokenRanges = buildAlignmentTokenRanges( + words, + tokens.map((token) => token.text), + { fillGaps: true }, + ); + + return tokenRanges.map((range) => + range ? { start: tokens[range.start].start, end: tokens[range.end].end } : null, + ); +} diff --git a/src/lib/client/html/highlight.ts b/src/lib/client/html/highlight.ts index f9a35f6..23779df 100644 --- a/src/lib/client/html/highlight.ts +++ b/src/lib/client/html/highlight.ts @@ -11,19 +11,27 @@ * sentence * - `WORD` — saturated background on the currently-spoken word * - * Word-to-DOM alignment is native: the Whisper alignment gives authoritative - * `charStart`/`charEnd` offsets into the sentence text, and we map those offsets - * directly onto a normalized char→DOM map of the located sentence wrap. No fuzzy - * token matching for words — that path used to jump the highlight to a later - * word in the sentence. + * Word-to-DOM alignment uses token-sequence alignment — the same primitive the + * PDF viewer uses. The located sentence wrap is reduced to a normalized char→DOM + * map; we tokenize that wrap text into words and globally align the Whisper + * words against those tokens (`buildAlignmentTokenRanges`). Each word then maps + * to a `[start, end)` char span in the wrap, which `wrapCharRange` turns into a + * DOM span. This tolerates divergent transcription, punctuation, whitespace, and + * markdown inline-element concatenation, and `fillGaps` guarantees every word + * resolves to a neighboring token rather than vanishing. */ import type { TTSSentenceAlignment } from '@/types/tts'; import { segmentWords } from '@/lib/shared/language'; import { findBestHighlightTokenMatch, + locateAlignmentWordSpans, normalizeHighlightToken, + type AlignmentCharSpan, } from '@/lib/client/highlight-token-alignment'; -import { normalizeMappedChars, type MappedChar } from '@/lib/client/highlight-char-map'; +import { + normalizeMappedChars, + type MappedChar, +} from '@/lib/client/highlight-char-map'; export const HTML_SENTENCE_CLASS = 'openreader-html-highlight-sentence'; export const HTML_WORD_CLASS = 'openreader-html-highlight-word'; @@ -55,13 +63,13 @@ interface SentenceState { // stable across word wrap/unwrap cycles because clear() calls // `parent.normalize()` which restores the original text-node structure. chars: CharPosition[]; - // Normalized text of the wrap (chars joined), used to locate the sentence so - // alignment char offsets can be rebased onto the wrap. + // Normalized text of the wrap (chars joined), tokenized to align each spoken + // word against the rendered words. text: string; - // For an alignment we've already seen: the offset of `sentence` inside `text` - // (the wrap window may be slightly wider than the sentence). null = not found. + // For an alignment we've already seen: each word's [start, end) char span + // within `text`/`chars` (null entries = words that aligned to no token). alignment: TTSSentenceAlignment | null; - base: number | null; + wordRanges: Array | null; } let sentenceState: SentenceState | null = null; @@ -157,9 +165,18 @@ function collectDomTokens( /** * Walk inside the current sentence wrap spans and build a normalized char→DOM * map. Every surviving character of the normalized text remembers the exact - * Text node + offset it came from, so alignment char offsets map straight to a - * DOM range. Normalization matches `preprocessSentenceForAudio` (the canonical - * space the alignment offsets live in). + * Text node + offset it came from, so an aligned word's char span maps straight + * to a DOM range. Normalization matches `preprocessSentenceForAudio` (the + * canonical space the alignment lives in). + * + * Crucially, a synthetic space is inserted between adjacent text nodes when the + * boundary isn't already whitespace. ReactMarkdown renders inline formatting as + * sibling nodes ("The quick brown" → "The"|"quick"|" brown") + * and the inter-word space lives at a node edge that the wrap drops — without + * this, the words would concatenate into "Thequick", collapsing distinct words + * into one token and destroying per-word highlight granularity. The synthetic + * space sits on a word boundary, so it is never inside an aligned word's span + * and never becomes a highlight target. */ function collectWrapCharMap(wraps: HTMLSpanElement[]): { chars: CharPosition[]; text: string } { const raw: MappedChar[] = []; @@ -169,6 +186,10 @@ function collectWrapCharMap(wraps: HTMLSpanElement[]): { chars: CharPosition[]; while (current) { const t = current as Text; const value = t.nodeValue || ''; + const lastChar = raw.length ? raw[raw.length - 1].char : ''; + if (lastChar && !/\s/.test(lastChar) && value.length && !/\s/.test(value[0])) { + raw.push({ char: ' ', pos: { node: t, offset: 0 } }); + } for (let offset = 0; offset < value.length; offset += 1) { raw.push({ char: value[offset], pos: { node: t, offset } }); } @@ -298,7 +319,7 @@ export function highlightHtmlSentence( chars, text, alignment: null, - base: null, + wordRanges: null, }; return true; } @@ -321,32 +342,21 @@ export function highlightHtmlWord( const words = alignment.words || []; if (!words.length || wordIndex >= words.length) return false; - // Locate the sentence inside the wrap's normalized text once per alignment. - // The sentence wrap is found by fuzzy token windowing so it may be slightly - // wider than the spoken sentence; rebasing keeps the alignment char offsets - // accurate regardless. - if (sentenceState.alignment !== alignment || sentenceState.base === null) { - const found = sentenceState.text.indexOf(alignment.sentence); - if (found < 0) { - // Sentence not located in the wrap — fail closed (leave only the sentence - // highlight) rather than snapping the word highlight to offset 0. Don't - // commit the alignment so the next tick retries the lookup. - sentenceState.base = null; - return false; - } + // Map each spoken word to a char span of the wrap's normalized text with the + // shared token-sequence aligner (same primitive as the EPUB and PDF viewers). + // The sentence wrap is located by fuzzy token windowing, so absolute char + // offsets can't be trusted; token alignment re-syncs every word against the + // rendered words and won't jump to a later/duplicate word. + if (sentenceState.alignment !== alignment || !sentenceState.wordRanges) { sentenceState.alignment = alignment; - sentenceState.base = found; + sentenceState.wordRanges = locateAlignmentWordSpans(words, sentenceState.text); } - const base = sentenceState.base; - if (base === null) return false; + const range = sentenceState.wordRanges[wordIndex]; + if (!range) return false; - const word = words[wordIndex]; - const { charStart, charEnd } = word; - if (!Number.isInteger(charStart) || !Number.isInteger(charEnd) || charEnd <= charStart) return false; - - const start = Math.max(0, base + charStart); - const end = Math.min(sentenceState.chars.length, base + charEnd); + const start = Math.max(0, range.start); + const end = Math.min(sentenceState.chars.length, range.end); if (end <= start) return false; wordWraps = wrapCharRange(sentenceState.chars, start, end, HTML_WORD_CLASS); diff --git a/tests/unit/epub-word-highlight.vitest.spec.ts b/tests/unit/epub-word-highlight.vitest.spec.ts deleted file mode 100644 index 23f6d77..0000000 --- a/tests/unit/epub-word-highlight.vitest.spec.ts +++ /dev/null @@ -1,50 +0,0 @@ -import { describe, expect, test } from 'vitest'; - -import { - resolveAlignmentWordSourceRange, -} from '../../src/lib/client/epub/epub-word-highlight'; -import type { CanonicalTtsSegment } from '../../src/lib/shared/tts-segment-plan'; -import type { TTSSentenceAlignment } from '../../src/types/tts'; - -const segment = (text: string, offset = 0): CanonicalTtsSegment => ({ - key: `segment:${offset}:${text}`, - ordinal: 0, - text, - ownerSourceKey: 'str:epubcfi(/6/2)', - ownerLocator: { location: 'epubcfi(/6/2)', readerType: 'epub' }, - startAnchor: { sourceKey: 'str:epubcfi(/6/2)', offset }, - endAnchor: { sourceKey: 'str:epubcfi(/6/2)', offset: offset + text.length }, - spansSourceBoundary: false, -}); - -describe('EPUB word highlight mapping', () => { - test('resolves Japanese alignment chunks directly from character offsets', () => { - const japanese = segment('これは日本語です。', 25); - const word: TTSSentenceAlignment['words'][number] = { - text: 'これは', - startSec: 0, - endSec: 0.5, - charStart: 0, - charEnd: 3, - }; - - expect(resolveAlignmentWordSourceRange(japanese, word)).toEqual({ - sourceStart: 25, - sourceEnd: 28, - }); - }); - - test('rejects out-of-range alignment character offsets so the word is skipped', () => { - const japanese = segment('これは日本語です。', 25); - const word: TTSSentenceAlignment['words'][number] = { - text: '範囲外', - startSec: 0, - endSec: 0.5, - charStart: 20, - charEnd: 23, - }; - - expect(resolveAlignmentWordSourceRange(japanese, word)).toBeNull(); - }); - -}); diff --git a/tests/unit/highlight-word-locate.vitest.spec.ts b/tests/unit/highlight-word-locate.vitest.spec.ts new file mode 100644 index 0000000..61a42b6 --- /dev/null +++ b/tests/unit/highlight-word-locate.vitest.spec.ts @@ -0,0 +1,73 @@ +import { describe, expect, test } from 'vitest'; + +import { locateAlignmentWordSpans } from '../../src/lib/client/highlight-token-alignment'; + +const words = (...texts: string[]) => + texts.map((text) => ({ text, startSec: 0, endSec: 0, charStart: 0, charEnd: 0 })); + +// Extract the highlighted substring for a span (or null) so assertions read in +// terms of what the user would actually see highlighted. +const slice = (region: string, span: { start: number; end: number } | null): string | null => + span ? region.slice(span.start, span.end) : null; + +describe('locateAlignmentWordSpans', () => { + test('maps each word to its own span when the region has separators (TXT)', () => { + const region = 'The quick brown fox'; + const spans = locateAlignmentWordSpans(words('The', 'quick', 'brown', 'fox'), region); + + expect(spans.map((s) => slice(region, s))).toEqual(['The', 'quick', 'brown', 'fox']); + }); + + test('degrades gracefully when inline DOM concatenated two words (MD)', () => { + // "The quick brown fox" with the inter-word space dropped + // at the node boundary collapses into a single "Thequick" token. Both words + // must still highlight (the merged region), and crucially NOTHING is null — + // no "nothing highlights", no "random word". + const region = 'Thequick brown fox'; + const spans = locateAlignmentWordSpans(words('The', 'quick', 'brown', 'fox'), region); + + expect(spans.every((s) => s !== null)).toBe(true); + expect(slice(region, spans[0])).toBe('Thequick'); + expect(slice(region, spans[1])).toBe('Thequick'); + expect(slice(region, spans[2])).toBe('brown'); + expect(slice(region, spans[3])).toBe('fox'); + }); + + test('tolerates punctuation/quote divergence between spoken words and region', () => { + const region = '“Hello,” she said'; + const spans = locateAlignmentWordSpans(words('hello', 'she', 'said'), region); + + expect(slice(region, spans[0])).toBe('Hello'); + expect(slice(region, spans[1])).toBe('she'); + expect(slice(region, spans[2])).toBe('said'); + }); + + test('fillGaps: a word absent from the region inherits a neighbor (never null)', () => { + // Whisper emitted "beta" but it is not in the rendered text. It must not + // leave a hole — fillGaps borrows the neighboring token so the highlight + // keeps moving instead of disappearing. + const region = 'alpha gamma delta'; + const spans = locateAlignmentWordSpans(words('alpha', 'beta', 'gamma', 'delta'), region); + + expect(spans.every((s) => s !== null)).toBe(true); + expect(slice(region, spans[0])).toBe('alpha'); + expect(slice(region, spans[2])).toBe('gamma'); + expect(slice(region, spans[3])).toBe('delta'); + // The orphan word borrows a neighbor rather than vanishing. + expect(['alpha', 'gamma']).toContain(slice(region, spans[1])); + }); + + test('is monotonic across repeated words (second "the" resolves later)', () => { + const region = 'the cat the dog'; + const spans = locateAlignmentWordSpans(words('the', 'cat', 'the', 'dog'), region); + + expect(spans[0]).toEqual({ start: 0, end: 3 }); + expect(spans[2]).toEqual({ start: 8, end: 11 }); + expect(spans[3]).toEqual({ start: 12, end: 15 }); + }); + + test('returns all-null for an empty region and empty for no words', () => { + expect(locateAlignmentWordSpans(words('a', 'b'), '')).toEqual([null, null]); + expect(locateAlignmentWordSpans([], 'anything')).toEqual([]); + }); +});