Preserve the start/end of the selection for non-text points when crea…

…ting link nodes

packages/lexical-link/src/index.ts

@@ -14,6 +14,7 @@ import type {
   LexicalCommand,
   LexicalNode,
   NodeKey,
+  Point,
   RangeSelection,
   SerializedElementNode,
 } from 'lexical';
@@ -28,10 +29,13 @@ import {
   $getSelection,
   $isElementNode,
   $isRangeSelection,
+  $normalizeSelection__EXPERIMENTAL,
+  $setSelection,
   createCommand,
   ElementNode,
   Spread,
 } from 'lexical';
+import invariant from 'shared/invariant';
 
 export type LinkAttributes = {
   rel?: null | string;
@@ -477,6 +481,66 @@ export const TOGGLE_LINK_COMMAND: LexicalCommand<
   string | ({url: string} & LinkAttributes) | null
 > = createCommand('TOGGLE_LINK_COMMAND');
 
+function $getPointNode(point: Point, offset: number): LexicalNode | null {
+  if (point.type === 'element') {
+    const node = point.getNode();
+    invariant(
+      $isElementNode(node),
+      '$getPointNode: element point is not an ElementNode',
+    );
+    const childNode = node.getChildren()[point.offset + offset];
+    return childNode || null;
+  }
+  return null;
+}
+
+/**
+ * Preserve the logical start/end of a RangeSelection in situations where
+ * the point is an element that may be reparented in the callback.
+ *
+ * @param $fn The function to run
+ * @returns The result of the callback
+ */
+function $withSelectedNodes<T>($fn: () => T): T {
+  const initialSelection = $getSelection();
+  if (!$isRangeSelection(initialSelection)) {
+    return $fn();
+  }
+  const normalized = $normalizeSelection__EXPERIMENTAL(initialSelection);
+  const isBackwards = normalized.isBackward();
+  const anchorNode = $getPointNode(normalized.anchor, isBackwards ? -1 : 0);
+  const focusNode = $getPointNode(normalized.focus, isBackwards ? 0 : -1);
+  const rval = $fn();
+  if (anchorNode || focusNode) {
+    const updatedSelection = $getSelection();
+    if ($isRangeSelection(updatedSelection)) {
+      const finalSelection = updatedSelection.clone();
+      if (anchorNode) {
+        const anchorParent = anchorNode.getParent();
+        if (anchorParent) {
+          finalSelection.anchor.set(
+            anchorParent.getKey(),
+            anchorNode.getIndexWithinParent() + (isBackwards ? 1 : 0),
+            'element',
+          );
+        }
+      }
+      if (focusNode) {
+        const focusParent = focusNode.getParent();
+        if (focusParent) {
+          finalSelection.focus.set(
+            focusParent.getKey(),
+            focusNode.getIndexWithinParent() + (isBackwards ? 0 : 1),
+            'element',
+          );
+        }
+      }
+      $setSelection($normalizeSelection__EXPERIMENTAL(finalSelection));
+    }
+  }
+  return rval;
+}
+
 /**
  * Generates or updates a LinkNode. It can also delete a LinkNode if the URL is null,
  * but saves any children and brings them up to the parent node.
@@ -515,93 +579,82 @@ export function $toggleLink(
         parentLink.remove();
       }
     });
-  } else {
-    // Add or merge LinkNodes
-    if (nodes.length === 1) {
-      const firstNode = nodes[0];
-      // if the first node is a LinkNode or if its
-      // parent is a LinkNode, we update the URL, target and rel.
-      const linkNode = $getAncestor(firstNode, $isLinkNode);
-      if (linkNode !== null) {
-        linkNode.setURL(url);
-        if (target !== undefined) {
-          linkNode.setTarget(target);
-        }
-        if (rel !== null) {
-          linkNode.setRel(rel);
-        }
-        if (title !== undefined) {
-          linkNode.setTitle(title);
-        }
-        return;
-      }
+    return;
+  }
+  const updatedNodes = new Set<NodeKey>();
+  const updateLinkNode = (linkNode: LinkNode) => {
+    if (updatedNodes.has(linkNode.getKey())) {
+      return;
+    }
+    updatedNodes.add(linkNode.getKey());
+    linkNode.setURL(url);
+    if (target !== undefined) {
+      linkNode.setTarget(target);
+    }
+    if (rel !== undefined) {
+      linkNode.setRel(rel);
     }
+    if (title !== undefined) {
+      linkNode.setTitle(title);
+    }
+  };
+  // Add or merge LinkNodes
+  if (nodes.length === 1) {
+    const firstNode = nodes[0];
+    // if the first node is a LinkNode or if its
+    // parent is a LinkNode, we update the URL, target and rel.
+    const linkNode = $getAncestor(firstNode, $isLinkNode);
+    if (linkNode !== null) {
+      return updateLinkNode(linkNode);
+    }
+  }
 
-    let prevParent: ElementNode | LinkNode | null = null;
+  $withSelectedNodes(() => {
     let linkNode: LinkNode | null = null;
-
-    nodes.forEach((node) => {
-      const parent = node.getParent();
-
-      if (
-        parent === linkNode ||
-        parent === null ||
-        ($isElementNode(node) && !node.isInline())
-      ) {
-        return;
+    for (const node of nodes) {
+      if (!node.isAttached()) {
+        continue;
       }
-
-      if ($isLinkNode(parent)) {
-        linkNode = parent;
-        parent.setURL(url);
-        if (target !== undefined) {
-          parent.setTarget(target);
-        }
-        if (rel !== null) {
-          linkNode.setRel(rel);
-        }
-        if (title !== undefined) {
-          linkNode.setTitle(title);
-        }
-        return;
+      const parentLinkNode = $getAncestor(node, $isLinkNode);
+      if (parentLinkNode) {
+        updateLinkNode(parentLinkNode);
+        continue;
       }
-
-      if (!parent.is(prevParent)) {
-        prevParent = parent;
-        linkNode = $createLinkNode(url, {rel, target, title});
-
-        if ($isLinkNode(parent)) {
-          if (node.getPreviousSibling() === null) {
-            parent.insertBefore(linkNode);
-          } else {
-            parent.insertAfter(linkNode);
-          }
-        } else {
-          node.insertBefore(linkNode);
-        }
-      }
-
-      if ($isLinkNode(node)) {
-        if (node.is(linkNode)) {
-          return;
+      if ($isElementNode(node)) {
+        if (!node.isInline()) {
+          // Ignore block nodes, if there are any children we will see them
+          // later and wrap in a new LinkNode
+          continue;
         }
-        if (linkNode !== null) {
-          const children = node.getChildren();
-
-          for (let i = 0; i < children.length; i++) {
-            linkNode.append(children[i]);
+        if ($isLinkNode(node)) {
+          // If it's not an autolink node and we don't already have a LinkNode
+          // in this block then we can update it and re-use it
+          if (
+            !$isAutoLinkNode(node) &&
+            (linkNode === null || !linkNode.getParentOrThrow().isParentOf(node))
+          ) {
+            updateLinkNode(node);
+            linkNode = node;
+            continue;
+          }
+          // Unwrap LinkNode, we already have one or it's an AutoLinkNode
+          for (const child of node.getChildren()) {
+            node.insertBefore(child);
           }
+          node.remove();
+          continue;
         }
-
-        node.remove();
-        return;
       }
-
-      if (linkNode !== null) {
-        linkNode.append(node);
+      const prevLinkNode = node.getPreviousSibling();
+      if ($isLinkNode(prevLinkNode) && prevLinkNode.is(linkNode)) {
+        prevLinkNode.append(node);
+        continue;
       }
-    });
-  }
+      linkNode = $createLinkNode(url, {rel, target, title});
+      node.insertAfter(linkNode);
+      linkNode.append(node);
+    }
+  });
 }
 /** @deprecated renamed to {@link $toggleLink} by @lexical/eslint-plugin rules-of-lexical */
 export const toggleLink = $toggleLink;