Fix Selection & Cursor Placement, Add Attachment Selection

Author: thecoolwinter

Sources/CodeEditTextView/TextLayoutManager/TextAttachments/TextAttachment.swift

@@ -10,6 +10,7 @@ import AppKit
 /// Represents an attachment type. Attachments take up some set width, and draw their contents in a receiver view.
 public protocol TextAttachment: AnyObject {
     var width: CGFloat { get }
+    var isSelected: Bool { get set }
     func draw(in context: CGContext, rect: NSRect)
 }
 

Sources/CodeEditTextView/TextLayoutManager/TextAttachments/TextAttachmentManager.swift

@@ -15,6 +15,7 @@ import Foundation
 public final class TextAttachmentManager {
     private var orderedAttachments: [AnyTextAttachment] = []
     weak var layoutManager: TextLayoutManager?
+    private var selectionObserver: (any NSObjectProtocol)?
 
     /// Adds a new attachment, keeping `orderedAttachments` sorted by range.location.
     /// If two attachments overlap, the layout phase will later ignore the one with the higher start.
@@ -94,7 +95,7 @@ public final class TextAttachmentManager {
     /// - Returns: An array of `AnyTextAttachment` instances whose ranges intersect `query`.
     public func getAttachmentsOverlapping(_ range: NSRange) -> [AnyTextAttachment] {
         // Find the first attachment whose end is beyond the start of the query.
-        guard let startIdx = firstIndex(where: { $0.range.upperBound > range.location }) else {
+        guard let startIdx = firstIndex(where: { $0.range.upperBound >= range.location }) else {
             return []
         }
 
@@ -107,8 +108,8 @@ public final class TextAttachmentManager {
             if attachment.range.location >= range.upperBound {
                 break
             }
-            if attachment.range.intersection(range)?.length ?? 0 > 0,
-               results.last?.range != attachment.range {
+            if (attachment.range.intersection(range)?.length ?? 0 > 0 || attachment.range.max == range.location)
+                && results.last?.range != attachment.range {
                 results.append(attachment)
             }
             idx += 1
@@ -131,6 +132,43 @@ public final class TextAttachmentManager {
             }
         }
     }
+    
+    /// Set up the attachment manager to listen to selection updates, giving text attachments a chance to respond to
+    /// selection state.
+    ///
+    /// This is specifically not in the initializer to prevent a bit of a chicken-and-the-egg situation where the
+    /// layout manager and selection manager need each other to init.
+    ///
+    /// - Parameter selectionManager: The selection manager to listen to.
+    func setUpSelectionListener(for selectionManager: TextSelectionManager) {
+        if let selectionObserver {
+            NotificationCenter.default.removeObserver(selectionObserver)
+        }
+
+        selectionObserver = NotificationCenter.default.addObserver(
+            forName: TextSelectionManager.selectionChangedNotification,
+            object: selectionManager,
+            queue: .main
+        ) { [weak self] notification in
+            guard let selectionManager = notification.object as? TextSelectionManager else {
+                return
+            }
+            let selectedSet = IndexSet(ranges: selectionManager.textSelections.map({ $0.range }))
+            for attachment in self?.orderedAttachments ?? [] {
+                let isSelected = selectedSet.contains(integersIn: attachment.range)
+                if attachment.attachment.isSelected != isSelected {
+                    self?.layoutManager?.invalidateLayoutForRange(attachment.range)
+                }
+                attachment.attachment.isSelected = isSelected
+            }
+        }
+    }
+
+    deinit {
+        if let selectionObserver {
+            NotificationCenter.default.removeObserver(selectionObserver)
+        }
+    }
 }
 
 private extension TextAttachmentManager {

Sources/CodeEditTextView/TextLayoutManager/TextLayoutManager+Iterator.swift

@@ -196,7 +196,7 @@ public extension TextLayoutManager {
         }
 
         if lastAttachment.range.max > originalPosition.position.range.max,
-           let extendedLinePosition = lineStorage.getLine(atOffset: lastAttachment.range.max) {
+           var extendedLinePosition = lineStorage.getLine(atOffset: lastAttachment.range.max) {
             newPosition = TextLineStorage<TextLine>.TextLinePosition(
                 data: newPosition.data,
                 range: NSRange(start: newPosition.range.location, end: extendedLinePosition.range.max),
@@ -207,6 +207,14 @@ public extension TextLayoutManager {
             maxIndex = max(maxIndex, extendedLinePosition.index)
         }
 
+        if firstAttachment.range.location == newPosition.range.location {
+            minIndex = max(minIndex, 0)
+        }
+
+        if lastAttachment.range.max == newPosition.range.max {
+            maxIndex = min(maxIndex, lineStorage.count - 1)
+        }
+
         // Base case, we haven't updated anything
         if minIndex...maxIndex == originalPosition.indexRange {
             return (newPosition, minIndex...maxIndex)

Sources/CodeEditTextView/TextLayoutManager/TextLayoutManager.swift

@@ -69,7 +69,7 @@ public class TextLayoutManager: NSObject {
     // MARK: - Internal
 
     weak var textStorage: NSTextStorage?
-    var lineStorage: TextLineStorage<TextLine> = TextLineStorage()
+    public var lineStorage: TextLineStorage<TextLine> = TextLineStorage()
     var markedTextManager: MarkedTextManager = MarkedTextManager()
     let viewReuseQueue: ViewReuseQueue<LineFragmentView, LineFragment.ID> = ViewReuseQueue()
     package var visibleLineIds: Set<TextLine.ID> = []

Sources/CodeEditTextView/TextSelectionManager/TextSelectionManager+FillRects.swift

@@ -72,10 +72,15 @@ extension TextSelectionManager {
             }
 
             let maxRect: CGRect
+            let endOfLine = fragmentRange.max <= range.max || range.contains(fragmentRange.max)
+            let endOfDocument = intersectionRange.max == layoutManager.lineStorage.length
+            let emptyLine = linePosition.range.isEmpty
+
             // If the selection is at the end of the line, or contains the end of the fragment, and is not the end
             // of the document, we select the entire line to the right of the selection point.
-            if (fragmentRange.max <= range.max || range.contains(fragmentRange.max))
-                && intersectionRange.max != layoutManager.lineStorage.length {
+            // true, !true = false, false
+            // true, !true = false, true
+            if endOfLine && !(endOfDocument && !emptyLine) {
                 maxRect = CGRect(
                     x: rect.maxX,
                     y: fragmentPosition.yPos + linePosition.yPos,

Sources/CodeEditTextView/TextView/TextView.swift

@@ -346,6 +346,8 @@ public class TextView: NSView, NSTextContent {
         selectionManager = setUpSelectionManager()
         selectionManager.useSystemCursor = useSystemCursor
 
+        layoutManager.attachments.setUpSelectionListener(for: selectionManager)
+
         _undoManager = CEUndoManager(textView: self)
 
         layoutManager.layoutLines()

Tests/CodeEditTextViewTests/LayoutManager/TextLayoutManagerAttachmentsTests.swift

@@ -114,8 +114,8 @@ struct TextLayoutManagerAttachmentsTests {
         let height = try #require(layoutManager.textLineForOffset(0)).height
         layoutManager.attachments.add(DemoTextAttachment(), for: NSRange(start: 0, end: 3))
 
-        // With bug: this the line for offset 3 is > 0 because it wasn't updated for the new attachment.
-        #expect(layoutManager.textLineForOffset(0)?.height == height)
-        #expect(layoutManager.textLineForOffset(3)?.height == 0)
+        // With bug: the line for offset 3 would be the 2nd line (index 1). They should be merged
+        #expect(layoutManager.textLineForOffset(0)?.index == 0)
+        #expect(layoutManager.textLineForOffset(3)?.index == 0)
     }
 }

Tests/CodeEditTextViewTests/TypesetterTests.swift

@@ -3,6 +3,7 @@ import XCTest
 
 final class DemoTextAttachment: TextAttachment {
     var width: CGFloat
+    var isSelected: Bool = false
 
     init(width: CGFloat = 100) {
         self.width = width