From 52d9ee79a36fc23b8af3f325dd739c1bb1edfd86 Mon Sep 17 00:00:00 2001
From: Alf Eaton <alf.eaton@overleaf.com>
Date: Tue, 4 Jul 2023 09:15:02 +0100
Subject: [PATCH] [visual] Handle selections adjacent to lists and section
 headings (#13581)

GitOrigin-RevId: 35b289102110f88587679740eeed575e16f6788b
---
 .../source-editor/extensions/toolbar/lists.ts |  9 ++--
 .../extensions/toolbar/sections.ts            | 41 +++++++++++++++----
 .../utils/tree-operations/ancestors.ts        | 22 ++++++++++
 3 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/services/web/frontend/js/features/source-editor/extensions/toolbar/lists.ts b/services/web/frontend/js/features/source-editor/extensions/toolbar/lists.ts
index cc6fb39a3a..b2aac520ca 100644
--- a/services/web/frontend/js/features/source-editor/extensions/toolbar/lists.ts
+++ b/services/web/frontend/js/features/source-editor/extensions/toolbar/lists.ts
@@ -17,6 +17,7 @@ import {
   ancestorOfNodeWithType,
   ancestorWithType,
   descendantsOfNodeWithType,
+  wrappedNodeOfType,
 } from '../../utils/tree-operations/ancestors'
 import { getEnvironmentName } from '../../utils/tree-operations/environments'
 import { ListEnvironment } from '../../lezer-latex/latex.terms.mjs'
@@ -195,11 +196,9 @@ const toggleListForRange = (
   range: SelectionRange,
   environment: string
 ) => {
-  const ancestorNode = ancestorNodeOfType(
-    view.state,
-    range.head,
-    ListEnvironment
-  )
+  const ancestorNode =
+    ancestorNodeOfType(view.state, range.head, ListEnvironment) ??
+    wrappedNodeOfType(view.state, range, ListEnvironment)
 
   if (ancestorNode) {
     const beginEnvNode = ancestorNode.getChild('BeginEnv')
diff --git a/services/web/frontend/js/features/source-editor/extensions/toolbar/sections.ts b/services/web/frontend/js/features/source-editor/extensions/toolbar/sections.ts
index 140869c1e7..f6f45a8fe4 100644
--- a/services/web/frontend/js/features/source-editor/extensions/toolbar/sections.ts
+++ b/services/web/frontend/js/features/source-editor/extensions/toolbar/sections.ts
@@ -1,7 +1,7 @@
 import { EditorSelection, EditorState, SelectionRange } from '@codemirror/state'
 import { EditorView } from '@codemirror/view'
 import { syntaxTree } from '@codemirror/language'
-import { ancestorsOfNodeWithType } from '../../utils/tree-operations/ancestors'
+import { ancestorOfNodeWithType } from '../../utils/tree-operations/ancestors'
 import { SyntaxNode } from '@lezer/common'
 
 export const findCurrentSectionHeadingLevel = (state: EditorState) => {
@@ -25,13 +25,38 @@ export const rangeInfo = (
   range: SelectionRange
 ): RangeInfo => {
   const tree = syntaxTree(state)
-  const node = tree.resolveInner(range.anchor)
-  const command = ancestorsOfNodeWithType(node, 'SectioningCommand').next()
-    .value
-  const ctrlSeq = command?.firstChild
-  const level = ctrlSeq
-    ? state.sliceDoc(ctrlSeq.from + 1, ctrlSeq.to).trim()
-    : 'text'
+
+  const fromNode = tree.resolveInner(range.from, 1)
+  const fromAncestor = ancestorOfNodeWithType(fromNode, 'SectioningCommand')
+
+  const toNode = tree.resolveInner(range.to, -1)
+  const toAncestor = ancestorOfNodeWithType(toNode, 'SectioningCommand')
+
+  const command = fromAncestor ?? toAncestor
+
+  // from and to are both outside section heading
+  if (!command) {
+    return { range, level: 'text' }
+  }
+
+  if (fromAncestor && toAncestor) {
+    // from and to are inside different section headings
+    if (fromAncestor !== toAncestor) {
+      return { range, level: 'text' }
+    }
+  } else {
+    // the range isn't empty and only one end is inside a section heading
+    if (!range.empty) {
+      return { range, level: 'text' }
+    }
+  }
+
+  const ctrlSeq = command.firstChild
+  if (!ctrlSeq) {
+    return { range, level: 'text' }
+  }
+
+  const level = state.sliceDoc(ctrlSeq.from + 1, ctrlSeq.to).trim()
 
   return { command, ctrlSeq, level, range }
 }
diff --git a/services/web/frontend/js/features/source-editor/utils/tree-operations/ancestors.ts b/services/web/frontend/js/features/source-editor/utils/tree-operations/ancestors.ts
index b698f0c056..67236f6434 100644
--- a/services/web/frontend/js/features/source-editor/utils/tree-operations/ancestors.ts
+++ b/services/web/frontend/js/features/source-editor/utils/tree-operations/ancestors.ts
@@ -64,6 +64,28 @@ export function getAncestorStack(
   return stack.reverse()
 }
 
+export const wrappedNodeOfType = (
+  state: EditorState,
+  range: SelectionRange,
+  type: string | number
+): SyntaxNode | null => {
+  if (range.empty) {
+    return null
+  }
+
+  const ancestorNode = ancestorNodeOfType(state, range.from, type, 1)
+
+  if (
+    ancestorNode &&
+    ancestorNode.from === range.from &&
+    ancestorNode.to === range.to
+  ) {
+    return ancestorNode
+  }
+
+  return null
+}
+
 export const ancestorNodeOfType = (
   state: EditorState,
   pos: number,