How to avoid flickering when resetting selection or limit the optional range of selection

[1111mp]

This is what I have achieved so far:

Untitled.mov

What I ultimately want to achieve is that when the cursor starts selecting from a CodeNode, only nodes inside the CodeNode can be selected, and nodes outside of the CodeNode are not allowed to be selected. I want to limit the selection range.

This is my current code:

function getPreviousTextNodeAfterCode(codeNode: HTMLElement): Text | null {
  let next = codeNode.previousSibling;
  while (next) {
    const textNode = next.firstChild;
    if (textNode && textNode.nodeType === Node.TEXT_NODE)
      return textNode as Text;
    next = textNode;
  }
  return null;
}

function getNextTextNodeAfterCode(codeNode: HTMLElement): Text | null {
  let next = codeNode.nextSibling;
  while (next) {
    const textNode = next.firstChild;
    if (textNode && textNode.nodeType === Node.TEXT_NODE)
      return textNode as Text;
    next = textNode;
  }
  return null;
}

function RestrictCodeSelectionPlugin(): JSX.Element | null {
  const [editor] = useLexicalComposerContext();

  useEffect(() => {
    return editor.registerCommand(
      SELECTION_CHANGE_COMMAND,
      () => {
        const selection = $getSelection();
        if (!$isRangeSelection(selection)) return false;

        const root = $getRoot();
        if (root.getChildrenSize() <= 1) return false;
        if (isSelectedAll(root, selection)) return false;

        const anchorNode = selection.anchor.getNode();
        const focusNode = selection.focus.getNode();

        function getParentUntilQuote(node: any) {
          let current = node.getTopLevelElement();
          console.log('current', current);
          if ($isCodePlusNode(current)) {
            return current;
          }
          return null;
        }

        const anchorQuote = getParentUntilQuote(anchorNode);
        const focusQuote = getParentUntilQuote(focusNode);
        console.log('selection', selection);
        console.log('anchorQuote', anchorQuote);
        console.log('focusQuote', focusQuote);

        if (anchorQuote && !focusQuote) {
          // const nativeSelection = getDOMSelection(editor._window);
          // console.log('nativeSelection', nativeSelection);
          // if (!nativeSelection) return false;
          // nativeSelection.removeAllRanges();
          // const domRange = document.createRange();
          // const lastDesc = anchorQuote.getLastDescendant()!;
          // const dom = editor.getElementByKey(lastDesc.getKey())!;
          // domRange.setStart(
          //   editor.getElementByKey(anchorNode.getKey())!.firstChild!,
          //   selection.anchor.offset,
          // );
          // domRange.setEnd(dom.firstChild!, lastDesc.getTextContentSize());
          // nativeSelection.removeAllRanges();
          // nativeSelection.addRange(domRange);

          // editor.update(() => {
          // $addUpdateTag(SKIP_SELECTION_FOCUS_TAG);
          const newSelection = $createRangeSelection();
          newSelection.anchor.set(
            anchorNode.getKey(),
            selection.anchor.offset,
            'text',
          );
          const lastNode = anchorQuote.getLastChild()!;
          newSelection.focus.set(
            lastNode.getKey(),
            lastNode.getTextContentSize(),
            'text',
          );
          $setSelection(newSelection);
          // });
          // const prevEditable = editor.isEditable();
          // editor.setEditable(false);
          // editor.update(
          //   () => {
          //     //
          //   },
          //   {
          //     onUpdate: () => {
          //       editor.setEditable(prevEditable);
          //     },
          //   },
          // );
          return true;
        }
        return true;
      },
      COMMAND_PRIORITY_CRITICAL,
    );
  }, [editor]);

  useEffect(() => {
    const root = editor.getRootElement();
    if (!root) return;

    const handleMouseDown = (e: MouseEvent) => {
      if (!e.shiftKey) return;

      const nativeSelection = getDOMSelection(editor._window);
      const target = e.target as HTMLElement | null;
      if (!target || !nativeSelection) return;

      editor.read(() => {
        const selection = $getSelection();

        if ($isRangeSelection(selection)) {
          const anchorNode = selection.anchor.getNode();
          const topNode = anchorNode.getTopLevelElement();
          if (
            $isCodePlusNode(topNode) &&
            !editor.getElementByKey(topNode.getKey())?.contains(target)
          ) {
            e.preventDefault();
            e.stopPropagation();

            if (target.parentElement?.tagName.toLowerCase() === 'code') {
              const domRange = document.createRange();
              const firstDesc = topNode.getFirstDescendant()!;
              domRange.setStart(
                editor.getElementByKey(firstDesc.getKey())!.firstChild!,
                0,
              );
              domRange.setEnd(
                editor.getElementByKey(selection.anchor.key)!.firstChild!,
                selection.anchor.offset,
              );
              nativeSelection.removeAllRanges();
              nativeSelection.addRange(domRange);
              return;
            }

            const anchorDom = editor.getElementByKey(selection.anchor.key)!;
            const domRange = document.createRange();
            if (
              anchorDom.compareDocumentPosition(target.parentElement!) &
              Node.DOCUMENT_POSITION_FOLLOWING
            ) {
              const lastDesc = topNode.getLastDescendant();
              const dom = editor.getElementByKey(lastDesc!.getKey())!;
              domRange.setStart(
                editor.getElementByKey(selection.anchor.key)?.firstChild || dom,
                selection.anchor.offset,
              );
              domRange.setEnd(
                dom.firstChild || dom,
                lastDesc!.getTextContentSize(),
              );
            } else {
              const firstDesc = topNode.getFirstDescendant()!;
              const dom = editor.getElementByKey(firstDesc!.getKey())!;
              domRange.setStart(dom.firstChild!, 0);
              domRange.setEnd(anchorDom.firstChild!, selection.anchor.offset);
            }

            nativeSelection.removeAllRanges();
            nativeSelection.addRange(domRange);
          }

          if (
            !$isCodePlusNode(topNode) &&
            target.parentElement?.tagName.toLowerCase() === 'code'
          ) {
            e.preventDefault();
            e.stopPropagation();

            const domRange = document.createRange();
            const anchorDom = editor.getElementByKey(selection.anchor.key)!;
            if (
              anchorDom.compareDocumentPosition(target.parentElement) &
              Node.DOCUMENT_POSITION_FOLLOWING
            ) {
              const textNode = getPreviousTextNodeAfterCode(
                target.parentElement,
              );
              if (!textNode) return;

              const anchorDomText = anchorDom.firstChild!;
              domRange.setStart(anchorDomText, selection.anchor.offset);
              domRange.setEnd(textNode, textNode.length);
              nativeSelection.removeAllRanges();
              nativeSelection.addRange(domRange);
            } else {
              const textNode = getNextTextNodeAfterCode(target.parentElement);
              if (!textNode) return;

              domRange.setStart(textNode, 0);
              const anchorDomText = anchorDom.firstChild!;
              domRange.setEnd(anchorDomText, selection.anchor.offset);
            }

            nativeSelection.removeAllRanges();
            nativeSelection.addRange(domRange);
          }
        }
      });
    };

    root.addEventListener('mousedown', handleMouseDown, true);

    return () => {
      root.removeEventListener('mousedown', handleMouseDown, true);
    };
  }, [editor]);

  return null;
}

export { RestrictCodeSelectionPlugin };

I implemented the effect of limiting the selection area by Shift+Click:

Screen.Recording.2025-10-06.at.15.09.59.mov

I have tried many ways but can't achieve it perfectly. Is there a better way?

Thanks in advance. 🙏

[etrepum]

There isn't currently a hook at the right time for doing this in a nice way, when SELECTION_CHANGE_COMMAND is dispatched the native selection and DOM are already updated so when you update it again it's basically going to schedule for the next frame. The TableSelection code basically handles this kind of thing but it's going about it another way entirely and it's a real mess. You might be able to manage by listening directly to document.addEventListener('selectionchange', …) and modifying only the native selection but there still might be some flicker (haven't tried recently).

There's also some stuff for handling this sort of thing with shadow roots, but you'd have to have a wrapper node around the code node for that since the children of the code node are not blocks themselves.