Fix Drawing Ordering, Use Attachment Ranges

Author: thecoolwinter

Sources/CodeEditSourceEditor/LineFolding/Model/LineFoldCalculator.swift

@@ -57,7 +57,7 @@ actor LineFoldCalculator {
         // Depth: Open range
         var openFolds: [Int: LineFoldStorage.RawFold] = [:]
         var currentDepth: Int = 0
-        var iterator = await controller.textView.layoutManager.linesInRange(documentRange)
+        var iterator = await controller.textView.layoutManager.lineStorage.makeIterator()
 
         var lines = await self.getMoreLines(
             controller: controller,
@@ -109,25 +109,27 @@ actor LineFoldCalculator {
 
         let attachments = await controller.textView.layoutManager.attachments
             .getAttachmentsOverlapping(documentRange)
-            .compactMap { $0.attachment as? LineFoldPlaceholder }
-            .map {
-                LineFoldStorage.DepthStartPair(depth: $0.fold.depth, start: $0.fold.range.lowerBound)
+            .compactMap { attachmentBox -> LineFoldStorage.DepthStartPair? in
+                guard let attachment = attachmentBox.attachment as? LineFoldPlaceholder else {
+                    return nil
+                }
+                return LineFoldStorage.DepthStartPair(depth: attachment.fold.depth, start: attachmentBox.range.location)
             }
 
         let storage = LineFoldStorage(
             documentLength: foldCache.max(
                 by: { $0.range.upperBound < $1.range.upperBound }
             )?.range.upperBound ?? documentRange.length,
             folds: foldCache.sorted(by: { $0.range.lowerBound < $1.range.lowerBound }),
-            collapsedProvider: { Set(attachments) }
+            collapsedRanges: Set(attachments)
         )
         valueStreamContinuation.yield(storage)
     }
 
     @MainActor
     private func getMoreLines(
         controller: TextViewController,
-        iterator: inout TextLayoutManager.RangeIterator,
+        iterator: inout TextLineStorage<TextLine>.TextLineStorageIterator,
         previousDepth: Int,
         foldProvider: LineFoldProvider
     ) -> [LineFoldProviderLineInfo]? {

Sources/CodeEditSourceEditor/LineFolding/Model/LineFoldStorage.swift

@@ -17,14 +17,6 @@ struct FoldRange: Sendable, Equatable {
     var isCollapsed: Bool
 }
 
-/// Represents a single fold run with stable identifier and collapse state
-struct FoldRun: Sendable {
-    let id: FoldRange.FoldIdentifier
-    let depth: Int
-    let range: Range<Int>
-    let isCollapsed: Bool
-}
-
 /// Sendable data model for code folding using RangeStore
 struct LineFoldStorage: Sendable {
     /// A temporary fold representation without stable ID
@@ -51,9 +43,9 @@ struct LineFoldStorage: Sendable {
     private var foldRanges: [FoldRange.FoldIdentifier: FoldRange] = [:]
 
     /// Initialize with the full document length
-    init(documentLength: Int, folds: [RawFold] = [], collapsedProvider: () -> Set<DepthStartPair> = { [] }) {
+    init(documentLength: Int, folds: [RawFold] = [], collapsedRanges: Set<DepthStartPair> = []) {
         self.store = RangeStore<FoldStoreElement>(documentLength: documentLength)
-        self.updateFolds(from: folds, collapsedProvider: collapsedProvider)
+        self.updateFolds(from: folds, collapsedRanges: collapsedRanges)
     }
 
     private mutating func nextFoldId() -> FoldRange.FoldIdentifier {
@@ -63,17 +55,14 @@ struct LineFoldStorage: Sendable {
 
     /// Replace all fold data from raw folds, preserving collapse state via callback
     /// - Parameter rawFolds: newly computed folds (depth + range)
-    /// - Parameter collapsedProvider: callback returning current collapsed ranges/depths
-    mutating func updateFolds(from rawFolds: [RawFold], collapsedProvider: () -> Set<DepthStartPair>) {
+    /// - Parameter collapsedRanges: Current collapsed ranges/depths
+    mutating func updateFolds(from rawFolds: [RawFold], collapsedRanges: Set<DepthStartPair>) {
         // Build reuse map by start+depth, carry over collapse state
         var reuseMap: [DepthStartPair: FoldRange] = [:]
         for region in foldRanges.values {
             reuseMap[DepthStartPair(depth: region.depth, start: region.range.lowerBound)] = region
         }
 
-        // Determine which ranges are currently collapsed
-        let collapsedSet = collapsedProvider()
-
         // Build new regions
         foldRanges.removeAll(keepingCapacity: true)
         store = RangeStore<FoldStoreElement>(documentLength: store.length)
@@ -85,7 +74,7 @@ struct LineFoldStorage: Sendable {
             let id = prior?.id ?? nextFoldId()
             let wasCollapsed = prior?.isCollapsed ?? false
             // override collapse if provider says so
-            let isCollapsed = collapsedSet.contains(key) || wasCollapsed
+            let isCollapsed = collapsedRanges.contains(key) || wasCollapsed
             let fold = FoldRange(id: id, depth: raw.depth, range: raw.range, isCollapsed: isCollapsed)
 
             foldRanges[id] = fold
@@ -99,6 +88,12 @@ struct LineFoldStorage: Sendable {
         store.storageUpdated(editedRange: editedRange, changeInLength: delta)
     }
 
+    mutating func toggleCollapse(forFold fold: FoldRange) {
+        guard var existingRange = foldRanges[fold.id] else { return }
+        existingRange.isCollapsed.toggle()
+        foldRanges[fold.id] = existingRange
+    }
+
     /// Query a document subrange and return all folds as an ordered list by start position
     func folds(in queryRange: Range<Int>) -> [FoldRange] {
         let runs = store.runs(in: queryRange.clamped(to: 0..<store.length))
@@ -120,20 +115,6 @@ struct LineFoldStorage: Sendable {
         }
 
         return result.sorted { $0.range.lowerBound < $1.range.lowerBound }
-//        let runs = store.runs(in: queryRange.clamped(to: 0..<store.length))
-//        var currentLocation = queryRange.lowerBound
-//        return runs.compactMap { run in
-//            defer {
-//                currentLocation += run.length
-//            }
-//            guard let value = run.value else { return nil }
-//            return FoldRun(
-//                id: value.id,
-//                depth: value.depth,
-//                range: currentLocation..<(currentLocation + run.length),
-//                isCollapsed: foldRanges[value.id]?.isCollapsed ?? false
-//            )
-//        }
     }
 
     /// Given a depth and a location, return the full original fold region

Sources/CodeEditSourceEditor/LineFolding/Model/LineFoldingModel.swift

@@ -77,11 +77,19 @@ class LineFoldingModel: NSObject, NSTextStorageDelegate {
     /// Finds the deepest cached fold and depth of the fold for a line number.
     /// - Parameter lineNumber: The line number to query, zero-indexed.
     /// - Returns: The deepest cached fold and depth of the fold if it was found.
-    func getCachedFoldAt(lineNumber: Int) -> (range: FoldRange, depth: Int)? {
+    func getCachedFoldAt(lineNumber: Int) -> FoldRange? {
         guard let lineRange = controller?.textView.layoutManager.textLineForIndex(lineNumber)?.range else { return nil }
-        guard let deepestFold = foldCache.folds(in: lineRange.intRange).max(by: { $0.depth < $1.depth }) else {
+        guard let deepestFold = foldCache.folds(in: lineRange.intRange).max(by: {
+            if $0.isCollapsed != $1.isCollapsed {
+                $1.isCollapsed // Collapsed folds take precedence.
+            } else if $0.isCollapsed {
+                $0.depth > $1.depth
+            } else {
+                $0.depth < $1.depth
+            }
+        }) else {
             return nil
         }
-        return (deepestFold, deepestFold.depth)
+        return deepestFold
     }
 }

Sources/CodeEditSourceEditor/LineFolding/View/FoldingRibbonView+Draw.swift

@@ -40,7 +40,15 @@ extension FoldingRibbonView {
         let textRange = rangeStart.range.location..<rangeEnd.range.upperBound
 
         let folds = getDrawingFolds(forTextRange: textRange)
-        for fold in folds {
+        for fold in folds.filter({ !$0.isCollapsed }) {
+            drawFoldMarker(
+                fold,
+                in: context,
+                using: layoutManager
+            )
+        }
+
+        for fold in folds.filter({ $0.isCollapsed }) {
             drawFoldMarker(
                 fold,
                 in: context,
@@ -141,6 +149,7 @@ extension FoldingRibbonView {
         context.setLineJoin(.round)
         context.setLineWidth(1.3)
 
+        context.setFillColor(NSColor.tertiaryLabelColor.cgColor)
         context.fill(fillRect)
         context.addPath(chevron)
         context.strokePath()

Sources/CodeEditSourceEditor/LineFolding/View/FoldingRibbonView.swift

@@ -130,28 +130,31 @@ class FoldingRibbonView: NSView {
         guard let layoutManager = model?.controller?.textView.layoutManager,
               event.type == .leftMouseDown,
               let lineNumber = layoutManager.textLineForPosition(clickPoint.y)?.index,
-              let fold = model?.getCachedFoldAt(lineNumber: lineNumber) else {
+              let fold = model?.getCachedFoldAt(lineNumber: lineNumber),
+              let firstLineInFold = layoutManager.textLineForOffset(fold.range.lowerBound) else {
             super.mouseDown(with: event)
             return
         }
         if let attachment = model?.controller?.textView?.layoutManager.attachments
-            .getAttachmentsStartingIn(NSRange(fold.range.range))
-            .filter({ $0.attachment is LineFoldPlaceholder })
-            .first {
+            .getAttachmentsStartingIn(NSRange(fold.range))
+            .filter({ $0.attachment is LineFoldPlaceholder && firstLineInFold.range.contains($0.range.location) }).first {
             layoutManager.attachments.remove(atOffset: attachment.range.location)
-//            fold.range.isCollapsed = false
             attachments.removeAll(where: { $0 === attachment.attachment })
         } else {
-            let placeholder = LineFoldPlaceholder(fold: fold.range)
-            layoutManager.attachments.add(placeholder, for: NSRange(fold.range.range))
+            let placeholder = LineFoldPlaceholder(fold: fold)
+            layoutManager.attachments.add(placeholder, for: NSRange(fold.range))
             attachments.append(placeholder)
-//            fold.range.collapsed = true
         }
 
+        model?.foldCache.toggleCollapse(forFold: fold)
         model?.controller?.textView.needsLayout = true
     }
 
     override func mouseMoved(with event: NSEvent) {
+        defer {
+            super.mouseMoved(with: event)
+        }
+
         let pointInView = convert(event.locationInWindow, from: nil)
         guard let lineNumber = model?.controller?.textView.layoutManager.textLineForPosition(pointInView.y)?.index,
               let fold = model?.getCachedFoldAt(lineNumber: lineNumber) else {
@@ -160,15 +163,15 @@ class FoldingRibbonView: NSView {
             return
         }
 
-        guard fold.range.range != hoveringFold?.range else {
+        guard fold.range != hoveringFold?.range else {
             return
         }
         hoverAnimationTimer?.invalidate()
         // We only animate the first hovered fold. If the user moves the mouse vertically into other folds we just
         // show it immediately.
         if hoveringFold == nil {
             hoverAnimationProgress = 0.0
-            hoveringFold = fold.range
+            hoveringFold = fold
 
             let duration: TimeInterval = 0.2
             let startTime = CACurrentMediaTime()
@@ -186,7 +189,7 @@ class FoldingRibbonView: NSView {
 
         // Don't animate these
         hoverAnimationProgress = 1.0
-        hoveringFold = fold.range
+        hoveringFold = fold
     }
 
     override func mouseExited(with event: NSEvent) {