     }
 }
 
-// MARK: - MessageCellLayoutContext
+// MARK: - Helpers
 
 extension MessagesCollectionViewFlowLayout {
 
-    internal func cellLayoutContext(for message: MessageType, at indexPath: IndexPath) -> MessageCellLayoutContext {
-        guard let cachedContext = layoutContextCache.object(forKey: message.messageId as NSString) else {
-            let newContext = newCellLayoutContext(for: message, at: indexPath)
-
-            layoutContextCache.setObject(newContext, forKey: message.messageId as NSString)
-
-            return newContext
-        }
-        return cachedContext
-    }
-
-    internal func newCellLayoutContext(for message: MessageType, at indexPath: IndexPath) -> MessageCellLayoutContext {
-        currentLayoutContext = MessageCellLayoutContext()
-        currentLayoutContext.avatarPosition = avatarPosition(for: message, at: indexPath)
-        currentLayoutContext.avatarSize = avatarSize(for: message, at: indexPath)
-        currentLayoutContext.messageContainerPadding = messageContainerPadding(for: message, at: indexPath)
-        currentLayoutContext.messageLabelInsets = messageLabelInsets(for: message, at: indexPath)
-        currentLayoutContext.messageContainerMaxWidth = messageContainerMaxWidth(for: message, at: indexPath)
-        currentLayoutContext.messageContainerSize = messageContainerSize(for: message, at: indexPath)
-        currentLayoutContext.topLabelAlignment = cellTopLabelAlignment(for: message, at: indexPath)
-        currentLayoutContext.topLabelMaxWidth = cellTopLabelMaxWidth(for: message, at: indexPath)
-        currentLayoutContext.topLabelSize = cellTopLabelSize(for: message, at: indexPath)
-        currentLayoutContext.bottomLabelAlignment = cellBottomLabelAlignment(for: message, at: indexPath)
-        currentLayoutContext.bottomLabelMaxWidth = cellBottomLabelMaxWidth(for: message, at: indexPath)
-        currentLayoutContext.bottomLabelSize = cellBottomLabelSize(for: message, at: indexPath)
-        currentLayoutContext.itemHeight = cellContentHeight(for: message, at: indexPath)
-        return currentLayoutContext
-    }
-}
-
-// MARK: - Helpers
-
-extension MessagesCollectionViewFlowLayout {
-
-    /// Returns the size required fit a NSAttributedString considering a constrained max width.
-    ///
-    /// - Parameters:
-    ///   - attributedText: The `NSAttributedString` used to calculate a size that fits.
-    ///   - maxWidth: The max width available for the label.
     internal func labelSize(for attributedText: NSAttributedString, considering maxWidth: CGFloat) -> CGSize {
 
-        let estimatedHeight = attributedText.height(considering: maxWidth)
-        let estimatedWidth = attributedText.width(considering: estimatedHeight)
-
-        let finalHeight = estimatedHeight.rounded(.up)
-        let finalWidth = estimatedWidth > maxWidth ? maxWidth : estimatedWidth.rounded(.up)
-
-        return CGSize(width: finalWidth, height: finalHeight)
-    }
-
-    /// Returns the size required to fit a String considering a constrained max width.
-    ///
-    /// - Parameters:
-    ///   - text: The `String` used to calculate a size that fits.
-    ///   - maxWidth: The max width available for the label.
-    internal func labelSize(for text: String, considering maxWidth: CGFloat, and font: UIFont) -> CGSize {
-
-        let estimatedHeight = text.height(considering: maxWidth, and: font)
-        let estimatedWidth = text.width(considering: estimatedHeight, and: font)
-
-        let finalHeight = estimatedHeight.rounded(.up)
-        let finalWidth = estimatedWidth > maxWidth ? maxWidth : estimatedWidth.rounded(.up)
-
-        return CGSize(width: finalWidth, height: finalHeight)
+        guard let cachedSize = attributedTextSizeCache.object(forKey: attributedText.hash as NSNumber)?.cgSizeValue else {
+            let estimatedHeight = attributedText.height(considering: maxWidth)
+            let estimatedWidth = attributedText.width(considering: estimatedHeight)
+
+            let finalHeight = estimatedHeight.rounded(.up)
+            let finalWidth = estimatedWidth > maxWidth ? maxWidth : estimatedWidth.rounded(.up)
+
+            return CGSize(width: finalWidth, height: finalHeight)
+        }
+
+        return cachedSize
     }
 
     /// Convenience property for accessing the layout object's `MessagesCollectionView`.

    ///   - indexPath: The `IndexPath` of the cell.
    open func sizeForItem(at indexPath: IndexPath) -> CGSize {
        let message = messagesDataSource.messageForItem(at: indexPath, in: messagesCollectionView)
        let context = cellLayoutContext(for: message, at: indexPath)
        guard let itemHeight = context.itemHeight else {
            fatalError("Unexpectedly received a nil itemHeight")
        }
        return CGSize(width: itemWidth, height: itemHeight)
    }
}

        switch message.data {
        case .text(let text):
            let messageInsets = messageLabelInsets(for: message, at: indexPath)
            let attributedText = NSAttributedString(string: text, attributes: [.font: messageLabelFont])
            messageContainerSize = labelSize(for: attributedText, considering: maxWidth)
            messageContainerSize.width += messageInsets.horizontal
            messageContainerSize.height += messageInsets.vertical
        case .attributedText(let text):
