Swift：如何实现根据行为返回不同类型的chunked方法重载？

现有一个将数组拆分为固定大小块的Array扩展方法：

extension Array {
  func chunked(ofCount chunkSize: Int) -> [[Element]] {
    stride(from: 0, to: count, by: chunkSize).map {
      Array(self[$0 ..< Swift.min($0 + chunkSize, count)])
    }
  }
}

想要新增一个重载版本，支持自动用nil填充最后一块（返回类型变为[[Element?]]），但直接用fillWithNils: Bool作为参数的话，无论传入true还是false，返回类型都会统一为[[Element?]]，不符合传入false时返回[[Element]]的预期。

希望通过协议+关联类型的方式实现类似如下的调用：

let strings = ["a", "b", "c", "d", "e"]
let noNils: [[String]] = strings.chunked(ofCount: 4, with: NoNils.self)
let yesNils: [[String?]] = strings.chunked(ofCount: 4, with: YesNils.self)

// 动态控制行为
let behavior: AnyChunkedBehavior<String> = ...
let dynamicResult = strings.chunked(ofCount: 4, with: behavior)

但最初的协议实现报错type 'NoNils' does not conform to protocol 'ChunkedBehavior'，原代码如下：

protocol ChunkedBehavior {
  associatedtype Element
  associatedtype ResultElement
}

enum NoNils: ChunkedBehavior { // 报错：不遵守协议
  typealias ResultElement = Element
}

enum YesNils: ChunkedBehavior {
  typealias ResultElement = Optional<Element>
}

extension Array {
  func chunked<C: ChunkedBehavior>(ofCount: Int, behavior: C) -> [[C.ResultElement]] where C.Element == Element {
    return [[]]
  }
}

问题根源 #

原协议的Element关联类型没有被具体类型绑定，枚举NoNils/YesNils无法确定Element的具体类型，导致协议一致性检查失败。我们需要调整协议设计，让行为类型作为策略，明确如何根据输入的数组元素类型生成对应的结果类型，并处理分块逻辑。

正确实现代码 #

1. 重新定义ChunkedBehavior协议

让协议关联输入的InputElement和输出的OutputElement，并提供处理单个分块的方法：

protocol ChunkedBehavior {
    associatedtype InputElement
    associatedtype OutputElement
    
    /// 处理单个分块，返回符合行为要求的结果
    static func processChunk(_ chunk: [InputElement], chunkSize: Int) -> [OutputElement]
}

2. 实现两种行为类型（用空结构体作为标记）

// 不填充Nil的行为：直接返回原块，输出类型和输入一致
struct NoNils<Element>: ChunkedBehavior {
    static func processChunk(_ chunk: [Element], chunkSize: Int) -> [Element] {
        chunk
    }
}

// 填充Nil的行为：将不足chunkSize的块用Nil补全，输出类型为可选型
struct YesNils<Element>: ChunkedBehavior {
    static func processChunk(_ chunk: [Element], chunkSize: Int) -> [Element?] {
        var filledChunk = chunk.map(Optional.some)
        while filledChunk.count < chunkSize {
            filledChunk.append(nil)
        }
        return filledChunk
    }
}

3. 扩展Array实现通用的chunked方法

extension Array {
    func chunked<C: ChunkedBehavior>(ofCount chunkSize: Int, with behavior: C.Type) -> [[C.OutputElement]] where C.InputElement == Element {
        guard chunkSize > 0 else { return [] }
        
        return stride(from: 0, to: count, by: chunkSize).map { startIndex in
            let endIndex = Swift.min(startIndex + chunkSize, count)
            let chunk = Array(self[startIndex..<endIndex])
            return C.processChunk(chunk, chunkSize: chunkSize)
        }
    }
}

4. 支持动态行为的类型擦除（可选）

如果需要用变量动态控制行为，由于带关联类型的协议不能直接作为变量类型，需要实现类型擦除：

struct AnyChunkedBehavior<InputElement, OutputElement>: ChunkedBehavior {
    private let _processChunk: ([InputElement], Int) -> [OutputElement]
    
    init<B: ChunkedBehavior>(_ behavior: B.Type) where B.InputElement == InputElement, B.OutputElement == OutputElement {
        self._processChunk = B.processChunk
    }
    
    static func processChunk(_ chunk: [InputElement], chunkSize: Int) -> [OutputElement] {
        fatalError("不会被调用，实际使用的是存储的闭包")
    }
    
    func processChunk(_ chunk: [InputElement], chunkSize: Int) -> [OutputElement] {
        _processChunk(chunk, chunkSize: chunkSize)
    }
}

// 扩展Array支持动态行为的调用
extension Array {
    func chunked<Output>(ofCount chunkSize: Int, with behavior: AnyChunkedBehavior<Element, Output>) -> [[Output]] {
        guard chunkSize > 0 else { return [] }
        
        return stride(from: 0, to: count, by: chunkSize).map { startIndex in
            let endIndex = Swift.min(startIndex + chunkSize, count)
            let chunk = Array(self[startIndex..<endIndex])
            return behavior.processChunk(chunk, chunkSize: chunkSize)
        }
    }
}

调用示例 #

let strings = ["a", "b", "c", "d", "e"]

// 不填充Nil：返回[[String]]
let noNilsResult = strings.chunked(ofCount: 4, with: NoNils.self)
// 输出：[["a", "b", "c", "d"], ["e"]]

// 填充Nil：返回[[String?]]
let yesNilsResult = strings.chunked(ofCount: 4, with: YesNils.self)
// 输出：[["a", "b", "c", "d"], ["e", nil, nil, nil]]

// 动态控制行为
let dynamicBehavior: AnyChunkedBehavior<String, String?> = AnyChunkedBehavior(YesNils.self)
let dynamicResult = strings.chunked(ofCount: 4, with: dynamicBehavior)
// 输出和yesNilsResult一致

关键说明 #

• 用空结构体代替枚举作为行为标记，因为结构体可以通过泛型参数绑定Element类型，解决原枚举无法关联类型的问题。
• 协议中定义processChunk静态方法，让每种行为自己实现分块的处理逻辑，职责更清晰。
• 类型擦除的AnyChunkedBehavior解决了带关联类型的协议不能作为变量类型的问题，支持动态切换行为。

内容的提问来源于stack exchange，提问作者Senseful