如何在JSQMessageViewController Swift 4中添加用户镜像

import UIKit
import JSQMessagesViewController

class MyChatViewController: JSQMessagesViewController {

    // 假设你有一个用户头像的字典
    var userAvatars: [String: JSQMessagesAvatarImage] = [:]

    override func viewDidLoad() {
        super.viewDidLoad()
        // 初始化用户头像
        initializeUserAvatars()
    }

    // 初始化用户头像
    func initializeUserAvatars() {
        // 假设你有一些用户ID和对应的头像图片名称
        let userIDs = ["user1", "user2"]
        let avatarNames = ["avatar1", "avatar2"]

        for (index, userID) in userIDs.enumerated() {
            if let avatarImage = UIImage(named: avatarNames[index]) {
                let avatar = JSQMessagesAvatarImageFactory.avatarImage(with: avatarImage, diameter: 64)
                userAvatars[userID] = avatar
            }
        }
    }

    // 提供每个消息单元格的用户头像
    override func collectionView(_ collectionView: JSQMessagesCollectionView, avatarImageDataForItemAt indexPath: IndexPath) -> JSQMessageAvatarImageDataSource? {
        let message = messages[indexPath.item]
        if let senderId = message.senderId, let avatar = userAvatars[senderId] {
            return avatar
        }
        return nil
    }

    // 其他必要的方法...
}