Attachment selection reset

Proton/Sources/ObjC/PRTextStorage.m

@@ -23,6 +23,7 @@
 
 @interface PRTextStorage ()
 @property (nonatomic) NSTextStorage *storage;
+@property (nonatomic, assign) BOOL isEditingAttributedText;
 @end
 
 @implementation PRTextStorage
@@ -74,6 +75,8 @@ - (void)edited:(NSTextStorageEditActions)editedMask range:(NSRange)editedRange c
 }
 
 - (void)replaceCharactersInRange:(NSRange)range withAttributedString:(NSAttributedString *)attrString {
+    _isEditingAttributedText = YES;
+    [self beginEditing];
     // TODO: Add undo behaviour
 
     // Handles the crash when nested list receives enter key in quick succession that unindents the list item.
@@ -103,23 +106,39 @@ - (void)replaceCharactersInRange:(NSRange)range withAttributedString:(NSAttribut
 
     NSAttributedString *deletedText = [_storage attributedSubstringFromRange:range];
     [_textStorageDelegate textStorage:self will:deletedText insertText:replacementString in:range];
+    NSArray<NSTextAttachment *> *attachmentsToDelete = [self attachmentsForRange:range];
     [super replaceCharactersInRange:range withAttributedString:replacementString];
+
+    [self endEditing];
+    _isEditingAttributedText = NO;
+    // Deleting of Attachment needs to happen after editing has ended. If invoked while textStorage editing is
+    // taking place, this may sometimes result in a crash(_fillLayoutHoleForCharacterRange).
+    [self deleteAttachments:attachmentsToDelete];
 }
 
 - (void)replaceCharactersInRange:(NSRange)range withString:(NSString *)str {
-    [self beginEditing];
-    NSInteger delta = str.length - range.length;
-
     NSArray<NSTextAttachment *> *attachmentsToDelete = [self attachmentsForRange:range];
-    for (NSTextAttachment *attachment in attachmentsToDelete) {
-        [_textStorageDelegate textStorage:self didDelete:attachment];
+    if(_isEditingAttributedText == NO) {
+        [self beginEditing];
     }
-
+    NSInteger delta = str.length - range.length;
     [_storage replaceCharactersInRange:range withString:str];
     [_storage fixAttributesInRange:NSMakeRange(0, _storage.length)];
     [self edited:NSTextStorageEditedCharacters & NSTextStorageEditedAttributes range:range changeInLength:delta];
+    if(_isEditingAttributedText == NO) {
+        [self endEditing];
+        // Deleting of Attachment needs to happen after editing has ended. If invoked while textStorage editing is
+        // taking place, this may sometimes result in a crash(_fillLayoutHoleForCharacterRange).
+        [self deleteAttachments:attachmentsToDelete];
+    }
+}
 
-    [self endEditing];
+-(void)deleteAttachments:(NSArray<NSTextAttachment *>*) attachments {
+    // Deleting of Attachment needs to happen after editing has ended. If invoked while textStorage editing is
+    // taking place, this may sometimes result in a crash(_fillLayoutHoleForCharacterRange).
+    for (NSTextAttachment *attachment in attachments) {
+        [_textStorageDelegate textStorage:self didDelete:attachment];
+    }
 }
 
 - (void)setAttributes:(NSDictionary<NSString *, id> *)attrs range:(NSRange)range {

Proton/Sources/Swift/Attachment/Attachment.swift

@@ -322,6 +322,11 @@ open class Attachment: NSTextAttachment, BoundsObserving {
             self.isAsyncRendered = true
             self.containerEditorView?.asyncAttachmentRenderingDelegate?.didRenderAttachment(self, in: containerEditorView)
         }
+        selectionView.onRemoveFromSuperview = { [weak self] in
+            guard let self,
+                  let range = rangeInContainer()?.nextPosition else { return }
+            self.containerEditorView?.selectedRange = range
+        }
     }
 
     private func setup(image: AttachmentImage) {