Improve note and editor search highlighting

Author: ChristianChiarulli

src/components/editor/note-editor.tsx

@@ -61,7 +61,11 @@ type NoteEditorProps = {
   html: string | null;
   isNew: boolean;
   markdown: string;
+  onEditorFocusChange?(focused: boolean): void;
+  onSearchMatchCountChange?(count: number): void;
   readOnly: boolean;
+  searchHighlightAllMatchesYellow?: boolean;
+  searchActiveMatchIndex?: number | null;
   searchQuery: string;
   toolbarContainer: HTMLElement | null;
   onChange(markdown: string): void;
@@ -80,9 +84,13 @@ function EditorInner({
   isNew,
   markdown,
   readOnly,
+  searchHighlightAllMatchesYellow,
+  searchActiveMatchIndex,
   searchQuery,
   toolbarContainer,
   onChange,
+  onEditorFocusChange,
+  onSearchMatchCountChange,
   onFocusHandled,
   editorRef,
 }: NoteEditorProps & {
@@ -144,9 +152,42 @@ function EditorInner({
     onFocusHandled();
   }, [editor, focusMode, onFocusHandled, readOnly]);
 
+  useEffect(() => {
+    const handleFocusIn = () => {
+      onEditorFocusChange?.(true);
+    };
+
+    const handleFocusOut = (event: FocusEvent) => {
+      const root = editor.getRootElement();
+      const nextTarget = event.relatedTarget;
+
+      if (root && nextTarget instanceof Node && root.contains(nextTarget)) {
+        return;
+      }
+
+      onEditorFocusChange?.(false);
+    };
+
+    return editor.registerRootListener((root, prevRoot) => {
+      if (prevRoot) {
+        prevRoot.removeEventListener("focusin", handleFocusIn);
+        prevRoot.removeEventListener("focusout", handleFocusOut);
+      }
+
+      if (!root) {
+        onEditorFocusChange?.(false);
+        return;
+      }
+
+      root.addEventListener("focusin", handleFocusIn);
+      root.addEventListener("focusout", handleFocusOut);
+      onEditorFocusChange?.(root.contains(document.activeElement));
+    });
+  }, [editor, onEditorFocusChange]);
+
   return (
     <>
-      <div className="comet-editor-content-wrap">
+      <div className="comet-editor-content-wrap relative">
         <ContentEditable
           className="comet-editor-content"
           autoCapitalize="off"
@@ -174,7 +215,12 @@ function EditorInner({
       <MarkdownPastePlugin />
       <YouTubeEmbedPlugin />
       <ImageDropPlugin />
-      <SearchHighlightPlugin searchWords={searchWords} />
+      <SearchHighlightPlugin
+        activeMatchIndex={searchActiveMatchIndex}
+        highlightAllMatchesYellow={searchHighlightAllMatchesYellow}
+        onMatchCountChange={onSearchMatchCountChange}
+        searchWords={searchWords}
+      />
       <TableActionMenuPlugin />
 
       <TableClickOutsidePlugin />

src/components/editor/plugins/search-highlight-plugin.tsx

@@ -1,17 +1,29 @@
-import { useEffect } from "react";
+import { useEffect, useRef } from "react";
 import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext";
 
+const ACTIVE_HIGHLIGHT_NAME = "comet-search-active";
 const HIGHLIGHT_NAME = "comet-search";
 const HIGHLIGHT_STYLE_ID = "comet-search-highlight-style";
 const HIGHLIGHT_STYLES = `
 ::highlight(${HIGHLIGHT_NAME}) {
+  background-color: var(--editor-selection);
+  color: var(--editor-text);
+}
+
+::highlight(${ACTIVE_HIGHLIGHT_NAME}) {
   background-color: rgb(253 224 71);
   color: var(--background);
 }
 `;
 
-function clearHighlight() {
-  (CSS as CSSWithHighlights).highlights?.delete(HIGHLIGHT_NAME);
+type CSSWithHighlights = typeof CSS & {
+  highlights?: Map<string, Highlight>;
+};
+
+function clearHighlights() {
+  const highlights = (CSS as CSSWithHighlights).highlights;
+  highlights?.delete(HIGHLIGHT_NAME);
+  highlights?.delete(ACTIVE_HIGHLIGHT_NAME);
 }
 
 function getOrCreateStyleElement() {
@@ -25,25 +37,30 @@ function getOrCreateStyleElement() {
 }
 
 export default function SearchHighlightPlugin({
+  activeMatchIndex = null,
+  highlightAllMatchesYellow = false,
+  onMatchCountChange,
   searchWords,
 }: {
+  activeMatchIndex?: number | null;
+  highlightAllMatchesYellow?: boolean;
+  onMatchCountChange?(count: number): void;
   searchWords: string[];
 }) {
   const [editor] = useLexicalComposerContext();
+  const shouldScrollRef = useRef(true);
 
-  // Toggle highlight visibility by adding/removing the CSS rule.
-  // Highlight ranges stay registered — they're just invisible without the rule.
   useEffect(() => {
     return editor.registerRootListener((root, prevRoot) => {
       if (prevRoot) {
         prevRoot.removeEventListener("focusin", handleFocusIn);
         prevRoot.removeEventListener("focusout", handleFocusOut);
       }
+
       if (root) {
         root.addEventListener("focusin", handleFocusIn);
         root.addEventListener("focusout", handleFocusOut);
 
-        // Sync to current state
         if (root.contains(document.activeElement)) {
           handleFocusIn();
         } else {
@@ -61,62 +78,136 @@ export default function SearchHighlightPlugin({
     }
   }, [editor]);
 
-  // Maintain highlight ranges — always applied regardless of focus.
   useEffect(() => {
     const highlights = (CSS as CSSWithHighlights).highlights;
-    if (!highlights) return;
+    if (!highlights) {
+      onMatchCountChange?.(0);
+      return;
+    }
 
-    getOrCreateStyleElement().textContent = HIGHLIGHT_STYLES;
+    shouldScrollRef.current = true;
+    const root = editor.getRootElement();
+    getOrCreateStyleElement().textContent =
+      root?.contains(document.activeElement) ? "" : HIGHLIGHT_STYLES;
 
     if (searchWords.length === 0) {
-      highlights.delete(HIGHLIGHT_NAME);
+      clearHighlights();
+      onMatchCountChange?.(0);
       return;
     }
 
     const apply = () => {
-      const root = editor.getRootElement();
-      if (!root) {
-        highlights.delete(HIGHLIGHT_NAME);
+      const currentRoot = editor.getRootElement();
+      if (!currentRoot) {
+        clearHighlights();
+        onMatchCountChange?.(0);
         return;
       }
 
-      const escaped = searchWords.map((w) =>
-        w.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"),
+      const escaped = searchWords.map((word) =>
+        word.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"),
       );
       const regex = new RegExp(`(${escaped.join("|")})`, "gi");
       const ranges: Range[] = [];
-      const walker = document.createTreeWalker(root, NodeFilter.SHOW_TEXT);
+      const walker = document.createTreeWalker(
+        currentRoot,
+        NodeFilter.SHOW_TEXT,
+      );
 
       let textNode: Text | null;
       while ((textNode = walker.nextNode() as Text | null)) {
         const text = textNode.textContent || "";
         let match: RegExpExecArray | null;
         while ((match = regex.exec(text)) !== null) {
-          const range = new Range();
+          const range = document.createRange();
           range.setStart(textNode, match.index);
           range.setEnd(textNode, match.index + match[0].length);
           ranges.push(range);
         }
       }
 
-      if (ranges.length > 0) {
-        highlights.set(HIGHLIGHT_NAME, new Highlight(...ranges));
-      } else {
+      onMatchCountChange?.(ranges.length);
+      if (ranges.length === 0) {
+        clearHighlights();
+        return;
+      }
+
+      if (highlightAllMatchesYellow) {
         highlights.delete(HIGHLIGHT_NAME);
+        highlights.set(ACTIVE_HIGHLIGHT_NAME, new Highlight(...ranges));
+      } else {
+      const hasActiveMatch =
+        activeMatchIndex !== null &&
+        activeMatchIndex >= 0 &&
+        activeMatchIndex < ranges.length;
+        const inactiveRanges = hasActiveMatch
+          ? ranges.filter((_, index) => index !== activeMatchIndex)
+          : ranges;
+
+        if (inactiveRanges.length > 0) {
+          highlights.set(HIGHLIGHT_NAME, new Highlight(...inactiveRanges));
+        } else {
+          highlights.delete(HIGHLIGHT_NAME);
+        }
+
+        if (hasActiveMatch) {
+          highlights.set(
+            ACTIVE_HIGHLIGHT_NAME,
+            new Highlight(ranges[activeMatchIndex]),
+          );
+        } else {
+          highlights.delete(ACTIVE_HIGHLIGHT_NAME);
+        }
+      }
+
+      if (!shouldScrollRef.current) {
+        return;
+      }
+
+      shouldScrollRef.current = false;
+      const hasActiveMatch =
+        !highlightAllMatchesYellow &&
+        activeMatchIndex !== null &&
+        activeMatchIndex >= 0 &&
+        activeMatchIndex < ranges.length;
+      const targetRange = hasActiveMatch ? ranges[activeMatchIndex] : ranges[0];
+      const scrollContainer = currentRoot.closest(
+        "[data-editor-scroll-container]",
+      );
+      if (!scrollContainer) {
+        return;
       }
+
+      const scrollRect = scrollContainer.getBoundingClientRect();
+      const targetRect = targetRange.getBoundingClientRect();
+      const hasVisibleMatch =
+        targetRect.bottom > scrollRect.top && targetRect.top < scrollRect.bottom;
+
+      if (hasVisibleMatch) {
+        return;
+      }
+
+      const scrollTop =
+        scrollContainer.scrollTop +
+        targetRect.top -
+        scrollRect.top -
+        scrollRect.height / 3;
+      scrollContainer.scrollTo({ top: scrollTop, behavior: "instant" });
     };
 
     apply();
     return editor.registerUpdateListener(() => {
       apply();
     });
-  }, [editor, searchWords]);
+  }, [
+    activeMatchIndex,
+    editor,
+    highlightAllMatchesYellow,
+    onMatchCountChange,
+    searchWords,
+  ]);
 
-  useEffect(() => clearHighlight, []);
+  useEffect(() => clearHighlights, []);
 
   return null;
 }
-
-type CSSWithHighlights = typeof CSS & {
-  highlights?: Map<string, Highlight>;
-};