Parallel map for arrays in Swift

I needed a parallel map for parallel computation of glyphs in VimR. I use the following (I don’t need any throwing transformations)

extension RandomAccessCollection where Index == Int {

  func parallelMap<T>(
    chunkSize: Int = 1,
    _ transform: @escaping (Element) -> T
  ) -> [T] {
    let count = self.count
    guard count > chunkSize else { return self.map(transform) }

    var result = Array<T?>(repeating: nil, count: count)

    result.withUnsafeMutableBufferPointer { pointer in
      if chunkSize == 1 {
        DispatchQueue.concurrentPerform(iterations: count) { i in
          pointer[i] = transform(self[i])
        }
      } else {
        let chunkCount = Int(ceil(Double(self.count) / Double(chunkSize)))
        DispatchQueue.concurrentPerform(iterations: chunkCount) { chunkIndex in
          let start = Swift.min(chunkIndex * chunkSize, count)
          let end = Swift.min(start + chunkSize, count)

          (start..<end).forEach { i in pointer[i] = transform(self[i]) }
        }
      }
    }

    return result.map { $0! }
  }
}

Using withUnsafeMutableBufferPointer(_:), you avoid synchronization. I guess you cannot avoid the array of Optionals and the map step at the end until the Accessing an Array’s Uninitialized Buffer proposal is implemented. However, if you can provide a default value for the type T, then you can omit the map step by initializing the result array with the default value

var result = Array<T>(repeating: someDefaultValue, count: count)
// [...]
return result
comments powered by Disqus