Merge Trailing Line on Attachments, Select Attachments

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)
 }
 
@@ -18,8 +19,8 @@ public protocol TextAttachment: AnyObject {
 /// This type cannot be initialized outside of `CodeEditTextView`, but will be received when interrogating
 /// the ``TextAttachmentManager``.
 public struct AnyTextAttachment: Equatable {
-    var range: NSRange
-    let attachment: any TextAttachment
+    package(set) public var range: NSRange
+    public let attachment: any TextAttachment
 
     var width: CGFloat {
         attachment.width

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.
@@ -23,11 +24,28 @@ public final class TextAttachmentManager {
         let attachment = AnyTextAttachment(range: range, attachment: attachment)
         let insertIndex = findInsertionIndex(for: range.location)
         orderedAttachments.insert(attachment, at: insertIndex)
+
+        // This is ugly, but if our attachment meets the end of the next line, we need to merge that line with this
+        // one.
+        var getNextOne = false
         layoutManager?.lineStorage.linesInRange(range).dropFirst().forEach {
             if $0.height != 0 {
                 layoutManager?.lineStorage.update(atOffset: $0.range.location, delta: 0, deltaHeight: -$0.height)
             }
+
+            // Only do this if it's not the end of the document
+            if range.max == $0.range.max && range.max != layoutManager?.lineStorage.length {
+                getNextOne = true
+            }
+        }
+
+        if getNextOne,
+            let trailingLine = layoutManager?.lineStorage.getLine(atOffset: range.max),
+           trailingLine.height != 0 {
+            // Update the one trailing line.
+            layoutManager?.lineStorage.update(atOffset: range.max, delta: 0, deltaHeight: -trailingLine.height)
         }
+
         layoutManager?.setNeedsLayout()
     }
 
@@ -77,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 []
         }
 
@@ -90,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
@@ -114,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+Layout.swift

@@ -250,7 +250,7 @@ extension TextLayoutManager {
         let view = viewReuseQueue.getOrCreateView(forKey: lineFragment.data.id) {
             renderDelegate?.lineFragmentView(for: lineFragment.data) ?? LineFragmentView()
         }
-        view.translatesAutoresizingMaskIntoConstraints = false
+        view.translatesAutoresizingMaskIntoConstraints = true // Small optimization for lots of subviews
         view.setLineFragment(lineFragment.data, renderer: lineFragmentRenderer)
         view.frame.origin = CGPoint(x: edgeInsets.left, y: yPos)
         layoutView?.addSubview(view, positioned: .below, relativeTo: nil)

Sources/CodeEditTextView/TextLayoutManager/TextLayoutManager+Public.swift

@@ -203,7 +203,7 @@ extension TextLayoutManager {
     ///   - line: The line to calculate rects for.
     /// - Returns: Multiple bounding rects. Will return one rect for each line fragment that overlaps the given range.
     public func rectsFor(range: NSRange) -> [CGRect] {
-        return lineStorage.linesInRange(range).flatMap { self.rectsFor(range: range, in: $0) }
+        return linesInRange(range).flatMap { self.rectsFor(range: range, in: $0) }
     }
 
     /// Calculates all text bounding rects that intersect with a given range, with a given line position.

Sources/CodeEditTextView/TextLayoutManager/TextLayoutManager.swift

@@ -76,7 +76,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()
     let lineFragmentRenderer: LineFragmentRenderer

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

@@ -108,4 +108,14 @@ struct TextLayoutManagerAttachmentsTests {
         // Line "5" is from the trailing newline. That shows up as an empty line in the view.
         #expect(lines.map { $0.index } == [0, 4])
     }
+
+    @Test
+    func addingAttachmentThatMeetsEndOfLineMergesNextLine() throws {
+        let height = try #require(layoutManager.textLineForOffset(0)).height
+        layoutManager.attachments.add(DemoTextAttachment(), for: NSRange(start: 0, end: 3))
+
+        // 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