From f2e21ee179b9bafd026f0e2bdcfaa38eea6fbb2c Mon Sep 17 00:00:00 2001
From: Jared Pereira <jared@awarm.space>
Date: Tue, 2 Jul 2024 17:23:46 -0400
Subject: [PATCH] fix a lot of little text-selection and toolbar issues

Namely switch to iterating over every position in a slection and
checking if there is a mark, and then updating the mark if neccessary
instead of toggling it.

Also add a specific case for handling links since they aren't inclusive
of the last position
---
 components/TextBlock/index.tsx              |  15 ++-
 components/Toolbar/HighlightButton.tsx      |  18 ++-
 components/Toolbar/LinkButton.tsx           |  17 ++-
 components/Toolbar/TextDecorationButton.tsx |  13 +-
 src/utils/prosemirror/rangeHasMark.ts       |  30 +++++
 src/utils/prosemirror/setMark.ts            | 138 ++++++++++++++++++++
 6 files changed, 210 insertions(+), 21 deletions(-)
 create mode 100644 src/utils/prosemirror/rangeHasMark.ts
 create mode 100644 src/utils/prosemirror/setMark.ts

diff --git a/components/TextBlock/index.tsx b/components/TextBlock/index.tsx
index 695fd041..f1098d51 100644
--- a/components/TextBlock/index.tsx
+++ b/components/TextBlock/index.tsx
@@ -17,7 +17,7 @@ import {
 } from "src/replicache";
 import { isVisible } from "src/utils/isVisible";
 
-import { EditorState } from "prosemirror-state";
+import { EditorState, TextSelection } from "prosemirror-state";
 import { ySyncPlugin } from "y-prosemirror";
 import { Replicache } from "replicache";
 import { generateKeyBetween } from "fractional-indexing";
@@ -35,6 +35,8 @@ import { BlockOptions } from "components/BlockOptions";
 import { setEditorState, useEditorStates } from "src/state/useEditorState";
 import { isIOS } from "@react-aria/utils";
 import { useIsMobile } from "src/hooks/isMobile";
+import { setMark } from "src/utils/prosemirror/setMark";
+import { rangeHasMark } from "src/utils/prosemirror/rangeHasMark";
 
 export function TextBlock(props: BlockProps & { className: string }) {
   let initialized = useInitialPageLoad();
@@ -429,7 +431,16 @@ export function BaseTextBlock(props: BlockProps & { className: string }) {
 function CommandHandler(props: { entityID: string }) {
   let cb = useEditorEventCallback(
     (view, args: { mark: MarkType; attrs?: any }) => {
-      toggleMark(args.mark, args.attrs)(view.state, view.dispatch);
+      let { to, from, $cursor, $to, $from } = view.state
+        .selection as TextSelection;
+      let mark = rangeHasMark(view.state, args.mark, from, to);
+      if (
+        mark &&
+        (!args.attrs ||
+          JSON.stringify(args.attrs) === JSON.stringify(mark.attrs))
+      ) {
+        toggleMark(args.mark, args.attrs)(view.state, view.dispatch);
+      } else setMark(args.mark, args.attrs)(view.state, view.dispatch);
     },
   );
   useAppEventListener(props.entityID, "toggleMark", cb, []);
diff --git a/components/Toolbar/HighlightButton.tsx b/components/Toolbar/HighlightButton.tsx
index f54b6609..38568219 100644
--- a/components/Toolbar/HighlightButton.tsx
+++ b/components/Toolbar/HighlightButton.tsx
@@ -28,6 +28,7 @@ import { useEntity, useReplicache } from "src/replicache";
 import { useMemo, useState } from "react";
 import { useColorAttribute } from "components/ThemeManager/useColorAttribute";
 import { useParams } from "next/navigation";
+import { rangeHasMark } from "src/utils/prosemirror/rangeHasMark";
 
 export const HighlightColorButton = (props: {
   color: "1" | "2" | "3";
@@ -40,16 +41,20 @@ export const HighlightColorButton = (props: {
   let hasMark: boolean = false;
   if (focusedEditor) {
     let { to, from, $cursor } = focusedEditor.editor.selection as TextSelection;
+
+    let mark = rangeHasMark(
+      focusedEditor.editor,
+      schema.marks.highlight,
+      from,
+      to,
+    );
     if ($cursor)
       hasMark = !!schema.marks.highlight.isInSet(
         focusedEditor.editor.storedMarks || $cursor.marks(),
       );
-    else
-      hasMark = focusedEditor.editor.doc.rangeHasMark(
-        from,
-        to,
-        schema.marks.highlight,
-      );
+    else {
+      hasMark = !!mark;
+    }
   }
   return (
     <button
@@ -58,7 +63,6 @@ export const HighlightColorButton = (props: {
         toggleMarkInFocusedBlock(schema.marks.highlight, {
           color: props.color,
         });
-        schema.marks.highlight.create({ color: props.color });
         props.setLastUsedHightlight(props.color);
       }}
     >
diff --git a/components/Toolbar/LinkButton.tsx b/components/Toolbar/LinkButton.tsx
index 47cfb660..f16311ac 100644
--- a/components/Toolbar/LinkButton.tsx
+++ b/components/Toolbar/LinkButton.tsx
@@ -7,6 +7,7 @@ import { useState } from "react";
 import { Separator } from "components/Layout";
 import { MarkType } from "prosemirror-model";
 import { setEditorState, useEditorStates } from "src/state/useEditorState";
+import { rangeHasMark } from "src/utils/prosemirror/rangeHasMark";
 
 export function LinkButton(props: { setToolBarState: (s: "link") => void }) {
   let focusedBlock = useUIState((s) => s.focusedBlock);
@@ -17,15 +18,13 @@ export function LinkButton(props: { setToolBarState: (s: "link") => void }) {
   if (focusedEditor) {
     let { to, from, $cursor } = focusedEditor.editor.selection as TextSelection;
     if ($cursor) isLink = !!schema.marks.link.isInSet($cursor.marks());
-    else {
-      isLink = true;
-      for (let pos = from + 1; pos < to - 1; pos++) {
-        const $pos = focusedEditor.editor.doc.resolve(pos);
-        if (!$pos.marks().find((mark) => mark.type === schema.marks.link)) {
-          isLink = false;
-        }
-      }
-    }
+    if (to !== from)
+      isLink = !!rangeHasMark(
+        focusedEditor.editor,
+        schema.marks.link,
+        from,
+        to,
+      );
   }
   return (
     <ToolbarButton
diff --git a/components/Toolbar/TextDecorationButton.tsx b/components/Toolbar/TextDecorationButton.tsx
index 6e3b705d..c4e4ac89 100644
--- a/components/Toolbar/TextDecorationButton.tsx
+++ b/components/Toolbar/TextDecorationButton.tsx
@@ -1,10 +1,11 @@
-import { MarkType } from "prosemirror-model";
+import { Mark, MarkType } from "prosemirror-model";
 import { useUIState } from "src/useUIState";
 import { ToolbarButton } from ".";
 import { toggleMark } from "prosemirror-commands";
 import { TextSelection } from "prosemirror-state";
 import { publishAppEvent } from "src/eventBus";
 import { useEditorStates } from "src/state/useEditorState";
+import { rangeHasMark } from "src/utils/prosemirror/rangeHasMark";
 
 export function TextDecorationButton(props: {
   mark: MarkType;
@@ -15,13 +16,19 @@ export function TextDecorationButton(props: {
     focusedBlock ? s.editorStates[focusedBlock.entityID] : null,
   );
   let hasMark: boolean = false;
+  let mark: Mark | null = null;
   if (focusedEditor) {
-    let { to, from, $cursor } = focusedEditor.editor.selection as TextSelection;
+    let { to, from, $cursor, $to, $from } = focusedEditor.editor
+      .selection as TextSelection;
+
+    mark = rangeHasMark(focusedEditor.editor, props.mark, from, to);
     if ($cursor)
       hasMark = !!props.mark.isInSet(
         focusedEditor.editor.storedMarks || $cursor.marks(),
       );
-    else hasMark = focusedEditor.editor.doc.rangeHasMark(from, to, props.mark);
+    else {
+      hasMark = !!mark;
+    }
   }
 
   return (
diff --git a/src/utils/prosemirror/rangeHasMark.ts b/src/utils/prosemirror/rangeHasMark.ts
new file mode 100644
index 00000000..4eec5ba9
--- /dev/null
+++ b/src/utils/prosemirror/rangeHasMark.ts
@@ -0,0 +1,30 @@
+import { schema } from "components/TextBlock/schema";
+import { Mark, MarkType } from "prosemirror-model";
+import { EditorState } from "prosemirror-state";
+
+export function rangeHasMark(
+  state: EditorState,
+  type: MarkType,
+  start: number,
+  end: number,
+) {
+  let mark: Mark | null = null;
+  let length = end - start;
+  if (!type.spec.inclusive) length = length - 1;
+  for (let i = 1; i <= length; i++) {
+    let pos = state.doc.resolve(start + i);
+    let markAtPos = pos.marks().find((f) => f.type === type);
+    if (!mark) {
+      if (!markAtPos) {
+        return null;
+      }
+      mark = markAtPos;
+    } else {
+      if (!markAtPos) return null;
+      if (mark !== markAtPos) {
+        return null;
+      }
+    }
+  }
+  return mark;
+}
diff --git a/src/utils/prosemirror/setMark.ts b/src/utils/prosemirror/setMark.ts
new file mode 100644
index 00000000..a6fcc70d
--- /dev/null
+++ b/src/utils/prosemirror/setMark.ts
@@ -0,0 +1,138 @@
+// Taken from https://github.com/ueberdosis/tiptap/blob/dfacb3b987b57b3ab518bae87bc3d263ebfb60d0/packages/core/src/commands/setMark.ts#L66
+import { MarkType, ResolvedPos } from "@tiptap/pm/model";
+import { EditorState, Transaction } from "@tiptap/pm/state";
+
+import { TextSelection } from "prosemirror-state";
+import { Mark, Schema } from "prosemirror-model";
+
+function canSetMark(
+  state: EditorState,
+  tr: Transaction,
+  newMarkType: MarkType,
+) {
+  const { selection } = tr;
+  let cursor: ResolvedPos | null = null;
+
+  if (selection instanceof TextSelection) {
+    cursor = selection.$cursor;
+  }
+
+  if (cursor) {
+    const currentMarks = state.storedMarks ?? cursor.marks();
+
+    // There can be no current marks that exclude the new mark
+    return (
+      !!newMarkType.isInSet(currentMarks) ||
+      !currentMarks.some((mark) => mark.type.excludes(newMarkType))
+    );
+  }
+
+  const { ranges } = selection;
+
+  return ranges.some(({ $from, $to }) => {
+    let someNodeSupportsMark =
+      $from.depth === 0
+        ? state.doc.inlineContent && state.doc.type.allowsMarkType(newMarkType)
+        : false;
+
+    state.doc.nodesBetween($from.pos, $to.pos, (node, _pos, parent) => {
+      // If we already found a mark that we can enable, return false to bypass the remaining search
+      if (someNodeSupportsMark) {
+        return false;
+      }
+
+      if (node.isInline) {
+        const parentAllowsMarkType =
+          !parent || parent.type.allowsMarkType(newMarkType);
+        const currentMarksAllowMarkType =
+          !!newMarkType.isInSet(node.marks) ||
+          !node.marks.some((otherMark) => otherMark.type.excludes(newMarkType));
+
+        someNodeSupportsMark =
+          parentAllowsMarkType && currentMarksAllowMarkType;
+      }
+      return !someNodeSupportsMark;
+    });
+
+    return someNodeSupportsMark;
+  });
+}
+export const setMark =
+  (type: MarkType, attributes = {}) =>
+  (state: EditorState, dispatch: (tr: Transaction) => void) => {
+    const { selection } = state;
+    const { empty, ranges } = selection;
+
+    let tr = state.tr;
+    if (empty) {
+      const oldAttributes = getMarkAttributes(state, type);
+
+      tr.addStoredMark(
+        type.create({
+          ...oldAttributes,
+          ...attributes,
+        }),
+      );
+    } else {
+      ranges.forEach((range) => {
+        const from = range.$from.pos;
+        const to = range.$to.pos;
+
+        state.doc.nodesBetween(from, to, (node, pos) => {
+          const trimmedFrom = Math.max(pos, from);
+          const trimmedTo = Math.min(pos + node.nodeSize, to);
+          const someHasMark = node.marks.find((mark) => mark.type === type);
+
+          // if there is already a mark of this type
+          // we know that we have to merge its attributes
+          // otherwise we add a fresh new mark
+          if (someHasMark) {
+            node.marks.forEach((mark) => {
+              if (type === mark.type) {
+                tr.addMark(
+                  trimmedFrom,
+                  trimmedTo,
+                  type.create({
+                    ...mark.attrs,
+                    ...attributes,
+                  }),
+                );
+              }
+            });
+          } else {
+            tr.addMark(trimmedFrom, trimmedTo, type.create(attributes));
+          }
+        });
+      });
+    }
+
+    dispatch(tr);
+  };
+
+function getMarkAttributes(
+  state: EditorState,
+  type: MarkType,
+): Record<string, any> {
+  const { from, to, empty } = state.selection;
+  const marks: Mark[] = [];
+
+  if (empty) {
+    if (state.storedMarks) {
+      marks.push(...state.storedMarks);
+    }
+
+    marks.push(...state.selection.$head.marks());
+  } else {
+    state.doc.nodesBetween(from, to, (node) => {
+      marks.push(...node.marks);
+    });
+  }
+
+  const mark = marks.find((markItem) => markItem.type.name === type.name);
+
+  if (!mark) {
+    return {};
+  }
+
+  return { ...mark.attrs };
+}
-- 
2.54.0