Update Text Container Width Correctly

Sources/CodeEditTextView/CEScrollView.swift

@@ -10,6 +10,12 @@ import STTextView
 
 class CEScrollView: NSScrollView {
 
+    override open var contentSize: NSSize {
+        var proposedSize = super.contentSize
+        proposedSize.width -= verticalRulerView?.requiredThickness ?? 0.0
+        return proposedSize
+    }
+
     override func mouseDown(with event: NSEvent) {
 
         if let textView = self.documentView as? STTextView,

Sources/CodeEditTextView/Controller/STTextViewController+Highlighter.swift

@@ -0,0 +1,42 @@
+//
+//  STTextViewController+Highlighter.swift
+//
+//
+//  Created by Khan Winter on 4/21/23.
+//
+
+import AppKit
+import SwiftTreeSitter
+
+extension STTextViewController {
+    /// Configures the `Highlighter` object
+    internal func setUpHighlighter() {
+        self.highlighter = Highlighter(
+            textView: textView,
+            highlightProvider: highlightProvider,
+            theme: theme,
+            attributeProvider: self,
+            language: language
+        )
+    }
+
+    /// Sets the highlight provider and re-highlights all text. This method should be used sparingly.
+    internal func setHighlightProvider(_ highlightProvider: HighlightProviding? = nil) {
+        var provider: HighlightProviding?
+
+        if let highlightProvider = highlightProvider {
+            provider = highlightProvider
+        } else {
+            let textProvider: ResolvingQueryCursor.TextProvider = { [weak self] range, _ -> String? in
+                return self?.textView.textContentStorage.textStorage?.mutableString.substring(with: range)
+            }
+
+            provider = TreeSitterClient(codeLanguage: language, textProvider: textProvider)
+        }
+
+        if let provider = provider {
+            self.highlightProvider = provider
+            highlighter?.setHighlightProvider(provider)
+        }
+    }
+}

Sources/CodeEditTextView/Controller/STTextViewController+TextContainer.swift

@@ -0,0 +1,39 @@
+//
+//  STTextViewController+TextContainer.swift
+//  
+//
+//  Created by Khan Winter on 4/21/23.
+//
+
+import AppKit
+import STTextView
+
+extension STTextViewController {
+    /// Update the text view's text container if needed.
+    ///
+    /// Effectively updates the container to reflect the `wrapLines` setting, and to reflect any updates to the ruler,
+    /// scroll view, or window frames.
+    internal func updateTextContainerWidthIfNeeded() {
+        let previousTrackingSetting = textView.widthTracksTextView
+        textView.widthTracksTextView = wrapLines
+        if wrapLines {
+            var proposedSize = ((view as? NSScrollView)?.contentSize ?? .zero)
+            proposedSize.height = .greatestFiniteMagnitude
+
+            if textView.textContainer.size != proposedSize || textView.frame.size != proposedSize {
+                textView.textContainer.size = proposedSize
+                textView.setFrameSize(proposedSize)
+            }
+        } else {
+            var proposedSize = textView.frame.size
+            proposedSize.width = ((view as? NSScrollView)?.contentSize ?? .zero).width
+            if previousTrackingSetting != wrapLines {
+                textView.textContainer.size = CGSize(
+                    width: CGFloat.greatestFiniteMagnitude, height: CGFloat.greatestFiniteMagnitude
+                )
+                textView.setFrameSize(proposedSize)
+                textView.textLayoutManager.textViewportLayoutController.layoutViewport()
+            }
+        }
+    }
+}

Sources/CodeEditTextView/Controller/STTextViewController.swift

@@ -9,7 +9,6 @@ import AppKit
 import SwiftUI
 import Combine
 import STTextView
-import SwiftTreeSitter
 import CodeEditLanguages
 import TextFormation
 
@@ -97,7 +96,7 @@ public class STTextViewController: NSViewController, STTextViewDelegate, ThemeAt
     internal var highlighter: Highlighter?
 
     /// The provided highlight provider.
-    private var highlightProvider: HighlightProviding?
+    internal var highlightProvider: HighlightProviding?
 
     // MARK: Init
 
@@ -183,7 +182,6 @@ public class STTextViewController: NSViewController, STTextViewDelegate, ThemeAt
         textView.selectedLineHighlightColor = theme.lineHighlight
         textView.string = self.text.wrappedValue
         textView.isEditable = self.isEditable
-        textView.widthTracksTextView = self.wrapLines
         textView.highlightSelectedLine = true
         textView.allowsUndo = true
         textView.setupMenus()
@@ -223,6 +221,7 @@ public class STTextViewController: NSViewController, STTextViewDelegate, ThemeAt
                                                queue: .main) { [weak self] _ in
             guard let self = self else { return }
             (self.view as? NSScrollView)?.contentView.contentInsets.bottom = self.bottomContentInsets
+            self.updateTextContainerWidthIfNeeded()
         }
 
         NotificationCenter.default.addObserver(
@@ -232,10 +231,19 @@ public class STTextViewController: NSViewController, STTextViewDelegate, ThemeAt
         ) { [weak self] _ in
             self?.updateCursorPosition()
         }
+
+        NotificationCenter.default.addObserver(
+            forName: NSView.frameDidChangeNotification,
+            object: (self.view as? NSScrollView)?.verticalRulerView,
+            queue: .main
+        ) { [weak self] _ in
+            self?.updateTextContainerWidthIfNeeded()
+        }
     }
 
-    public override func viewDidAppear() {
-        super.viewDidAppear()
+    public override func viewWillAppear() {
+        super.viewWillAppear()
+        updateTextContainerWidthIfNeeded()
     }
 
     public func textViewDidChangeText(_ notification: Notification) {
@@ -299,6 +307,7 @@ public class STTextViewController: NSViewController, STTextViewDelegate, ThemeAt
         }
 
         highlighter?.invalidate()
+        updateTextContainerWidthIfNeeded()
     }
 
     /// Gets all attributes for the given capture including the line height, background color, and text color.
@@ -324,37 +333,6 @@ public class STTextViewController: NSViewController, STTextViewDelegate, ThemeAt
         ((self.lineHeight) - font.lineHeight) / 2
     }
 
-    // MARK: - Highlighting
-
-    /// Configures the `Highlighter` object
-    private func setUpHighlighter() {
-        self.highlighter = Highlighter(textView: textView,
-                                       highlightProvider: highlightProvider,
-                                       theme: theme,
-                                       attributeProvider: self,
-                                       language: language)
-    }
-
-    /// Sets the highlight provider and re-highlights all text. This method should be used sparingly.
-    public func setHighlightProvider(_ highlightProvider: HighlightProviding? = nil) {
-        var provider: HighlightProviding?
-
-        if let highlightProvider = highlightProvider {
-            provider = highlightProvider
-        } else {
-            let textProvider: ResolvingQueryCursor.TextProvider = { [weak self] range, _ -> String? in
-                return self?.textView.textContentStorage.textStorage?.mutableString.substring(with: range)
-            }
-
-            provider = TreeSitterClient(codeLanguage: language, textProvider: textProvider)
-        }
-
-        if let provider = provider {
-            self.highlightProvider = provider
-            highlighter?.setHighlightProvider(provider)
-        }
-    }
-
     // MARK: Selectors
 
     override public func keyDown(with event: NSEvent) {