diff --git a/Sources/Layout/MessagesCollectionViewFlowLayout.swift b/Sources/Layout/MessagesCollectionViewFlowLayout.swift
index 4a42d762..5b2aa124 100644
--- a/Sources/Layout/MessagesCollectionViewFlowLayout.swift
+++ b/Sources/Layout/MessagesCollectionViewFlowLayout.swift
@@ -114,7 +114,7 @@ open class MessagesCollectionViewFlowLayout: UICollectionViewFlowLayout {
     open var videoMessageSizeCalculator = MediaMessageSizeCalculator()
     open var locationMessageSizeCalculator = LocationMessageSizeCalculator()
 
-    open func cellSizeCalculatorForItem(at indexPath: IndexPath) -> MessageSizeCalculator {
+    open func cellSizeCalculatorForItem(at indexPath: IndexPath) -> CellSizeCalculator {
         let message = messagesDataSource.messageForItem(at: indexPath, in: messagesCollectionView)
         switch message.data {
         case .text:
diff --git a/Sources/Models/MessageData.swift b/Sources/Models/MessageData.swift
index 0d3f8b91..25f619ad 100644
--- a/Sources/Models/MessageData.swift
+++ b/Sources/Models/MessageData.swift
@@ -54,7 +54,7 @@ public enum MessageData {
     /// A custom message.
     /// - Note: Using this case requires that you override the following methods and handle this case:
     ///   - `collectionView(_:cellForItemAt indexPath: IndexPath) -> UICollectionViewCell`
-    ///   - `cellSizeCalculatorForItem(at indexPath: IndexPath) -> MessageSizeCalculator`
+    ///   - `cellSizeCalculatorForItem(at indexPath: IndexPath) -> CellSizeCalculator`
     case custom(Any?)
 
     // MARK: - Not supported yet