Preserve attachment enclosing new lines

Proton/Sources/ObjC/PRTextStorage.m

@@ -91,13 +91,29 @@ - (void)replaceCharactersInRange:(NSRange)range withAttributedString:(NSAttribut
         // Out of bounds
         return;
     }
-    
+
     NSMutableAttributedString *replacementString = [attrString mutableCopy];
+    NSAttributedString *substring = [self attributedSubstringFromRange:range];
+
+    if (self.preserveNewlineBeforeBlock
+        && range.location > 0
+        && [self attributedStringHasNewline:substring atStart:NO]
+        && [self isCharacterAdjacentToRangeAnAttachment:self range:range checkBefore:NO]) {
+        replacementString = [self appendNewlineToAttributedString:[attrString mutableCopy] atStart:NO];
+    }
+
+    if (self.preserveNewlineAfterBlock
+        && range.location > 0
+        && [self attributedStringHasNewline:substring atStart:YES]
+        && [self isCharacterAdjacentToRangeAnAttachment:self range:range checkBefore:YES]) {
+        replacementString = [self appendNewlineToAttributedString:[attrString mutableCopy] atStart:YES];
+    }
+
     // Fix any missing attribute that is in the location being replaced, but not in the text that
     // is coming in.
-    if (range.length > 0 && attrString.length > 0) {
+    if (range.length > 0 && replacementString.length > 0) {
         NSDictionary<NSAttributedStringKey, id> *outgoingAttrs = [_storage attributesAtIndex:(range.location + range.length - 1) effectiveRange:nil];
-        NSDictionary<NSAttributedStringKey, id> *incomingAttrs = [attrString attributesAtIndex:0 effectiveRange:nil];
+        NSDictionary<NSAttributedStringKey, id> *incomingAttrs = [replacementString attributesAtIndex:0 effectiveRange:nil];
 
         NSMutableDictionary<NSAttributedStringKey, id> *diff = [NSMutableDictionary dictionary];
         for (NSAttributedStringKey outgoingKey in outgoingAttrs) {
@@ -215,6 +231,68 @@ - (void)removeAttribute:(NSAttributedStringKey)name range:(NSRange)range {
 
 #pragma mark - Private
 
+- (NSMutableAttributedString *)appendNewlineToAttributedString:(NSMutableAttributedString *)attributedString atStart:(BOOL)appendAtStart {
+    if (attributedString.length == 0) {
+        return [[NSMutableAttributedString alloc] initWithString:@"\n"]; // Return just a newline if the original string is empty.
+    }
+
+    // Create a new NSAttributedString with the newline character.
+    NSAttributedString *newlineAttributedString = [[NSAttributedString alloc] initWithString:@"\n"];
+
+    // Create a mutable copy of the original attributed string.
+    NSMutableAttributedString *mutableAttributedString = [[NSMutableAttributedString alloc] initWithAttributedString:attributedString];
+
+    if (appendAtStart) {
+        // Append the attributed newline at the start.
+        [mutableAttributedString insertAttributedString:newlineAttributedString atIndex:0];
+    } else {
+        // Append the attributed newline at the end.
+        [mutableAttributedString appendAttributedString:newlineAttributedString];
+    }
+
+    return mutableAttributedString;
+}
+
+- (BOOL) attributedStringHasNewline:(NSAttributedString *) attributedString atStart: (BOOL)atStart {
+    NSString *string = [attributedString string];
+    if (string.length == 0) {
+        return NO;
+    }
+
+    unichar characterToVerify = [string characterAtIndex: 0];
+    if (atStart == NO) {
+        characterToVerify = [string characterAtIndex:string.length - 1];
+    }
+
+    return [[NSCharacterSet newlineCharacterSet] characterIsMember:characterToVerify];
+}
+
+-(BOOL) isCharacterAdjacentToRangeAnAttachment: (NSAttributedString *) attributedString range: (NSRange) range checkBefore: (BOOL) checkBefore {
+    NSUInteger positionToCheck;
+
+    if (checkBefore) {
+        if (range.location == 0) {
+            return NO; // No character before the start of the string
+        }
+        positionToCheck = range.location - 1;
+    } else {
+        positionToCheck = NSMaxRange(range);
+        if (positionToCheck >= attributedString.length) {
+            return NO; // No character after the end of the string
+        }
+    }
+
+    // Retrieve the attributes at the position to check
+    NSDictionary *attributes = [attributedString attributesAtIndex:positionToCheck effectiveRange:NULL];
+
+    // Check if these attributes contain the NSAttachmentAttributeName
+    if ([attributes objectForKey:@"_isBlockAttachment"] != nil) {
+        return YES; // There is an attachment
+    }
+
+    return NO; // No attachment found
+}
+
 - (void)fixMissingAttributesForDeletedAttributes:(NSArray<NSAttributedStringKey> *)attrs range:(NSRange)range {
     if ((range.location + range.length) > _storage.length) {
         // Out of bounds

Proton/Sources/ObjC/include/PRTextStorage.h

@@ -47,6 +47,9 @@ NS_SWIFT_NAME(TextStorageDelegate)
 @property (weak, nullable) id<PRDefaultTextFormattingProviding> defaultTextFormattingProvider;
 @property (weak, nullable) id<PRTextStorageDelegate> textStorageDelegate;
 
+@property (nonatomic, assign) BOOL preserveNewlineBeforeBlock;
+@property (nonatomic, assign) BOOL preserveNewlineAfterBlock;
+
 @property (nonatomic, readonly) UIFont *defaultFont;
 @property (nonatomic, readonly) NSParagraphStyle *defaultParagraphStyle;
 @property (nonatomic, readonly) UIColor *defaultTextColor;

Proton/Sources/Swift/Core/RichTextView.swift

@@ -35,6 +35,21 @@ class RichTextView: AutogrowingTextView {
 
     private var delegateOverrides = [GestureRecognizerDelegateOverride]()
 
+    var preserveBlockAttachmentNewline: PreserveBlockAttachmentNewline = .none {
+        didSet {
+            richTextStorage.preserveNewlineBeforeBlock = false
+            richTextStorage.preserveNewlineAfterBlock = false
+
+            if preserveBlockAttachmentNewline.contains(.before) {
+                richTextStorage.preserveNewlineBeforeBlock = true
+            }
+
+            if preserveBlockAttachmentNewline.contains(.after) {
+                richTextStorage.preserveNewlineAfterBlock = true
+            }
+        }
+    }
+
     weak var defaultTextFormattingProvider: DefaultTextFormattingProviding?
     {
         get { richTextStorage.defaultTextFormattingProvider }

Proton/Sources/Swift/Editor/EditorView.swift

@@ -65,6 +65,20 @@ public struct AttachmentContentIdentifier {
     }
 }
 
+public struct PreserveBlockAttachmentNewline: OptionSet {
+    public let rawValue: Int
+
+    public init(rawValue: Int) {
+        self.rawValue = rawValue
+    }
+
+    public static let before = PreserveBlockAttachmentNewline(rawValue: 1 << 0)
+    public static let after = PreserveBlockAttachmentNewline(rawValue: 1 << 1)
+
+    public static let none: PreserveBlockAttachmentNewline = []
+    public static let both: PreserveBlockAttachmentNewline = [.before, .after]
+}
+
 /// Defines the height for the Editor
 public enum EditorHeight {
     /// Default controlled via autolayout.
@@ -146,6 +160,11 @@ open class EditorView: UIView {
         get { editorViewContext.delegate }
     }
 
+    public var preserveBlockAttachmentNewline: PreserveBlockAttachmentNewline {
+        get { richTextView.preserveBlockAttachmentNewline }
+        set { richTextView.preserveBlockAttachmentNewline = newValue }
+    }
+
     public var scrollView: UIScrollView {
         richTextView as UIScrollView
     }

Proton/Tests/Editor/EditorViewTests.swift

@@ -870,8 +870,70 @@ class EditorViewTests: XCTestCase {
         XCTAssertNil(dummyAttachment1.containerEditorView)
         XCTAssertNil(dummyAttachment2.containerEditorView)
     }
+
+    func testPreservesNewlineBeforeAttachmentOnDelete() {
+        let viewController = EditorTestViewController()
+        let editor = viewController.editor
+        editor.preserveBlockAttachmentNewline = .before
+
+        let testString = NSAttributedString(string: "test string\n")
+        editor.replaceCharacters(in: .zero, with: testString)
+        let attachment = makePanelAttachment()
+        editor.insertAttachment(in: editor.textEndRange, attachment: attachment)
+        XCTAssertEqual(editor.text, "test string\n\n")
+
+        editor.replaceCharacters(in: NSRange(location: 8, length: 4), with: "")
+        XCTAssertEqual(editor.attachmentsInRange(editor.attributedText.fullRange).first?.range, NSRange(location: 9, length: 1))
+
+        XCTAssertEqual(editor.text, "test str\n\n")
+    }
+
+    func testPreservesNewlineAfterAttachmentOnDelete() {
+        let viewController = EditorTestViewController()
+        let editor = viewController.editor
+        editor.preserveBlockAttachmentNewline = .after
+
+        let testString = NSAttributedString(string: "test string\n  After attachment")
+        editor.replaceCharacters(in: .zero, with: testString)
+        let attachment = makePanelAttachment()
+        editor.insertAttachment(in: NSRange(location: 13, length: 0), attachment: attachment)
+
+        editor.replaceCharacters(in: NSRange(location: 14, length: 7), with: "")
+
+        XCTAssertEqual(editor.text, "test string\n \n attachment")
+        XCTAssertEqual(editor.attachmentsInRange(editor.attributedText.fullRange).first?.range, NSRange(location: 13, length: 1))
+    }
+
+    func testDoesNotPreservesNewlineByDefault() {
+        let viewController = EditorTestViewController()
+        let editor = viewController.editor
+
+        let testString = NSAttributedString(string: "test string\n")
+        editor.replaceCharacters(in: .zero, with: testString)
+        let attachment = makePanelAttachment()
+        editor.insertAttachment(in: editor.textEndRange, attachment: attachment)
+        XCTAssertEqual(editor.text, "test string\n\n")
+
+        editor.replaceCharacters(in: NSRange(location: 8, length: 4), with: "")
+
+        XCTAssertEqual(editor.text, "test str\n")
+        XCTAssertEqual(editor.attachmentsInRange(editor.attributedText.fullRange).first?.range, NSRange(location: 8, length: 1))
+    }
 }
 
+
+func makePanelAttachment() -> Attachment {
+    let panel = PanelView()
+    panel.editor.forceApplyAttributedText = true
+    panel.backgroundColor = .cyan
+    panel.layer.borderWidth = 1.0
+    panel.layer.cornerRadius = 4.0
+    panel.layer.borderColor = UIColor.black.cgColor
+
+    return Attachment(panel, size: .fullWidth)
+}
+
+
 class DummyMultiEditorAttachment: Attachment {
     let view: DummyMultiEditorView
     init(numberOfEditors: Int) {