petal/web/src/components/Editor/SearchHighlight.ts

import { Extension } from '@tiptap/core'
import { Plugin, PluginKey } from '@tiptap/pm/state'
import type { EditorState, Transaction } from '@tiptap/pm/state'
import { Decoration, DecorationSet } from '@tiptap/pm/view'
import type { Node as PMNode } from '@tiptap/pm/model'
import { mapOffset } from './SuggestionHighlight'
import { holdRedraw } from './Composition'

// SearchHighlight powers the in-document Find & Replace bar. Like the suggestion
// layer it uses ProseMirror *decorations* (not stored marks), so matches are
// recomputed from the live document on every edit and never leave a stale
// highlight behind. Matches are found per-textblock (a query never spans a
// paragraph break), mirroring how the rest of the editor anchors text.

export interface Match {
  from: number
  to: number
}

interface PluginState {
  query: string
  caseSensitive: boolean
  matches: Match[]
  active: number // index into matches, or -1 when there are none
  decorations: DecorationSet
  // Held back while an IME composition was in flight — see Composition.ts.
  // `matches` is held with the decorations rather than recomputed on its own:
  // the Find bar's "3 / 7" and the wash on the page are one answer, and half of
  // it moving while the other half waits would be worse than both waiting.
  stale: boolean
}

export const searchPluginKey = new PluginKey<PluginState>('petalSearch')

// findMatches collects every occurrence of `query` across the document's
// textblocks and maps each to an absolute ProseMirror range.
function findMatches(doc: PMNode, query: string, caseSensitive: boolean): Match[] {
  if (!query) return []
  const needle = caseSensitive ? query : query.toLowerCase()
  const matches: Match[] = []
  doc.descendants((node, pos) => {
    if (!node.isTextblock) return true
    const haystackRaw = node.textContent
    const haystack = caseSensitive ? haystackRaw : haystackRaw.toLowerCase()
    let idx = haystack.indexOf(needle)
    while (idx >= 0) {
      matches.push({
        from: mapOffset(node, pos, idx),
        to: mapOffset(node, pos, idx + query.length),
      })
      idx = haystack.indexOf(needle, idx + query.length)
    }
    return false // don't descend into inline children
  })
  return matches
}

function build(doc: PMNode, query: string, caseSensitive: boolean, preferred: number): PluginState {
  const matches = findMatches(doc, query, caseSensitive)
  const active = matches.length === 0 ? -1 : Math.max(0, Math.min(preferred, matches.length - 1))
  const decos = matches.map((m, i) =>
    Decoration.inline(m.from, m.to, {
      class: i === active ? 'petal-find-match petal-find-match-active' : 'petal-find-match',
    }),
  )
  return { query, caseSensitive, matches, active, decorations: DecorationSet.create(doc, decos), stale: false }
}

const EMPTY: PluginState = {
  query: '',
  caseSensitive: false,
  matches: [],
  active: -1,
  decorations: DecorationSet.empty,
  stale: false,
}

// setSearch updates the query / case-sensitivity and recomputes matches. Passing
// an empty query clears the layer.
export function setSearch(
  state: EditorState,
  dispatch: (tr: Transaction) => void,
  query: string,
  caseSensitive: boolean,
) {
  dispatch(state.tr.setMeta(searchPluginKey, { kind: 'search', query, caseSensitive }))
}

// setActive moves the highlighted (current) match by index.
export function setActive(state: EditorState, dispatch: (tr: Transaction) => void, index: number) {
  dispatch(state.tr.setMeta(searchPluginKey, { kind: 'active', index }))
}

// clearSearch tears the layer down (on close).
export function clearSearch(state: EditorState, dispatch: (tr: Transaction) => void) {
  dispatch(state.tr.setMeta(searchPluginKey, { kind: 'clear' }))
}

// getSearchState reads the live plugin state (match list + active index) so the
// Find bar can render "n / m" and drive navigation/replacement.
export function getSearchState(state: EditorState): PluginState {
  return searchPluginKey.getState(state) ?? EMPTY
}

type Meta =
  | { kind: 'search'; query: string; caseSensitive: boolean }
  | { kind: 'active'; index: number }
  | { kind: 'clear' }

export function searchPlugin(): Plugin<PluginState> {
  return new Plugin<PluginState>({
    key: searchPluginKey,
    state: {
      init: () => EMPTY,
      apply(tr, value, oldState, newState): PluginState {
        const meta = tr.getMeta(searchPluginKey) as Meta | undefined
        // Clearing the layer is the one thing a composition never holds: it
        // removes decorations rather than adding them, and it is what closing
        // the Find bar does.
        if (meta?.kind === 'clear') return EMPTY

        const held = holdRedraw(tr, oldState)
        const query = meta?.kind === 'search' ? meta.query : value.query
        const caseSensitive = meta?.kind === 'search' ? meta.caseSensitive : value.caseSensitive

        if (meta?.kind === 'active') {
          if (value.matches.length === 0) return value
          const active = ((meta.index % value.matches.length) + value.matches.length) % value.matches.length
          if (held) return { ...value, active, stale: true }
          const decos = value.matches.map((m, i) =>
            Decoration.inline(m.from, m.to, {
              class: i === active ? 'petal-find-match petal-find-match-active' : 'petal-find-match',
            }),
          )
          return { ...value, active, stale: false, decorations: DecorationSet.create(newState.doc, decos) }
        }

        // A new query, or any document change: re-anchor so highlights track
        // edits and replaces. Once due, it stays due until it happens.
        const due = value.stale || meta?.kind === 'search' || (tr.docChanged && !!value.query)
        if (!due) return value
        if (held) {
          return {
            ...value,
            query,
            caseSensitive,
            stale: true,
            decorations: tr.docChanged ? value.decorations.map(tr.mapping, tr.doc) : value.decorations,
          }
        }
        const preferred = meta?.kind === 'search' && value.active < 0 ? 0 : value.active
        return build(newState.doc, query, caseSensitive, preferred)
      },
    },
    props: {
      decorations(state) {
        return searchPluginKey.getState(state)?.decorations
      },
    },
  })
}

export const SearchHighlight = Extension.create({
  name: 'searchHighlight',

  addProseMirrorPlugins() {
    return [searchPlugin()]
  },
})