From fe0a000fc536ec46efe1a4e49af6ad39e27cfed3 Mon Sep 17 00:00:00 2001
From: Pavel Anokhov <p.anokhov@gmail.com>
Date: Sat, 8 Sep 2018 11:16:42 +0300
Subject: [PATCH] MessagesDataSource & MessagesLayoutDelegate methods for
 MessageType.custom handling.

---
 Sources/Controllers/MessagesViewController.swift |  2 +-
 .../MessagesCollectionViewFlowLayout.swift       |  2 +-
 Sources/Protocols/MessagesDataSource.swift       | 16 ++++++++++++++++
 Sources/Protocols/MessagesLayoutDelegate.swift   | 16 ++++++++++++++++
 4 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/Sources/Controllers/MessagesViewController.swift b/Sources/Controllers/MessagesViewController.swift
index 41ab6b10..c9bf68d3 100644
--- a/Sources/Controllers/MessagesViewController.swift
+++ b/Sources/Controllers/MessagesViewController.swift
@@ -206,7 +206,7 @@ UICollectionViewDelegateFlowLayout, UICollectionViewDataSource {
             cell.configure(with: message, at: indexPath, and: messagesCollectionView)
             return cell
         case .custom:
-            fatalError(MessageKitError.customDataUnresolvedCell)
+            return messagesDataSource.customCell(for: message, at: indexPath, in: messagesCollectionView)
         }
     }
 
diff --git a/Sources/Layout/MessagesCollectionViewFlowLayout.swift b/Sources/Layout/MessagesCollectionViewFlowLayout.swift
index ff123360..26fb385a 100644
--- a/Sources/Layout/MessagesCollectionViewFlowLayout.swift
+++ b/Sources/Layout/MessagesCollectionViewFlowLayout.swift
@@ -151,7 +151,7 @@ open class MessagesCollectionViewFlowLayout: UICollectionViewFlowLayout {
         case .location:
             return locationMessageSizeCalculator
         case .custom:
-            fatalError("Must return a CellSizeCalculator for MessageKind.custom(Any?)")
+            return messagesLayoutDelegate.cellSizeCalculator(for: message, at: indexPath, in: messagesCollectionView)
         }
     }
 
diff --git a/Sources/Protocols/MessagesDataSource.swift b/Sources/Protocols/MessagesDataSource.swift
index 1763a8c8..0d35e592 100644
--- a/Sources/Protocols/MessagesDataSource.swift
+++ b/Sources/Protocols/MessagesDataSource.swift
@@ -92,6 +92,18 @@ public protocol MessagesDataSource: AnyObject {
     ///
     /// The default value returned by this method is `nil`.
     func messageBottomLabelAttributedText(for message: MessageType, at indexPath: IndexPath) -> NSAttributedString?
+    
+    
+    /// Custom collectionView cell for message with `custom` message type.
+    ///
+    /// - Parameters:
+    ///   - message: The `custom` message type
+    ///   - indexPath: The `IndexPath` of the cell.
+    ///   - messagesCollectionView: The `MessagesCollectionView` in which this cell will be displayed.
+    ///
+    /// - Note:
+    ///   This method will call fatalError() on default. You must override this method if you are using MessageType.custom messages.
+    func customCell(for message: MessageType, at indexPath: IndexPath, in messagesCollectionView: MessagesCollectionView) -> UICollectionViewCell
 }
 
 public extension MessagesDataSource {
@@ -115,4 +127,8 @@ public extension MessagesDataSource {
     func messageBottomLabelAttributedText(for message: MessageType, at indexPath: IndexPath) -> NSAttributedString? {
         return nil
     }
+    
+    func customCell(for message: MessageType, at indexPath: IndexPath, in messagesCollectionView: MessagesCollectionView) -> UICollectionViewCell {
+        fatalError(MessageKitError.customDataUnresolvedCell)
+    }
 }
diff --git a/Sources/Protocols/MessagesLayoutDelegate.swift b/Sources/Protocols/MessagesLayoutDelegate.swift
index 24a13c7a..f00955a7 100644
--- a/Sources/Protocols/MessagesLayoutDelegate.swift
+++ b/Sources/Protocols/MessagesLayoutDelegate.swift
@@ -80,6 +80,18 @@ public protocol MessagesLayoutDelegate: AnyObject {
     /// - Note:
     ///   The default value returned by this method is zero.
     func messageBottomLabelHeight(for message: MessageType, at indexPath: IndexPath, in messagesCollectionView: MessagesCollectionView) -> CGFloat
+    
+    
+    /// Cell size calculator for messages with MessageType.custom.
+    ///
+    /// - Parameters:
+    ///   - message: The custom message
+    ///   - indexPath: The `IndexPath` of the cell.
+    ///   - messagesCollectionView: The `MessagesCollectionView` in which this cell will be displayed.
+    ///
+    /// - Note:
+    ///   The default implementation will throw fatalError(). You must override this method if you are using messages with MessageType.custom.
+    func cellSizeCalculator(for message: MessageType, at indexPath: IndexPath, in messagesCollectionView: MessagesCollectionView) -> CellSizeCalculator
 }
 
 public extension MessagesLayoutDelegate {
@@ -103,4 +115,8 @@ public extension MessagesLayoutDelegate {
     func messageBottomLabelHeight(for message: MessageType, at indexPath: IndexPath, in messagesCollectionView: MessagesCollectionView) -> CGFloat {
         return 0
     }
+    
+    func cellSizeCalculator(for message: MessageType, at indexPath: IndexPath, in messagesCollectionView: MessagesCollectionView) -> CellSizeCalculator {
+        fatalError("Must return a CellSizeCalculator for MessageKind.custom(Any?)")
+    }
 }