AsyncMergeSort.swift

//
//  AsyncMergeSort.swift
//  BitRemote
//
//  Created by 荒木辰造 on R 5/03/08.
//

import Foundation

extension Array {
    func asyncSorted(using comparators: [KeyPathComparator<Element>], threadsCount: Int = 8) async -> Self {
        await withTaskGroup(of: Self.self) { group in
            let arrays = chunked(into: threadsCount)
            for array in arrays {
                group.addTask(operation: { await array.threadedSorted(using: comparators) })
            }

            var results = [Self]()
            results.reserveCapacity(count)
            for await array in group {
                results.append(array)
            }
            return await Self.asyncMerge(arrays: results, using: comparators)
        }
    }

    static func asyncMerge(arrays: [Self], using comparators: [KeyPathComparator<Element>]) async -> Self {
        await withTaskGroup(of: Self.self) { group in
            for (index, array) in arrays.enumerated() where index % 2 == 0 {
                if index + 1 < arrays.count {
                    group.addTask(operation: { await array.asyncMerged(with: arrays[index + 1], using: comparators) })
                } else {
                    group.addTask(operation: { array })
                }
            }

            var results = [Self]()
            for await array in group {
                results.append(array)
            }
            if results.count <= 1 {
                return results.first ?? .init()
            } else {
                return await asyncMerge(arrays: results, using: comparators)
            }
        }
    }

    func asyncMerged(with array: Self, using comparators: [KeyPathComparator<Element>]) async -> Self {
        let task = Task {
            var leftIndex = 0
            var rightIndex = 0
            var mergedArray = Self()

            while leftIndex < count && rightIndex < array.count {
                let leftElement = self[leftIndex]
                let rightElement = array[rightIndex]

                switch comparators.compare(leftElement, rightElement) {
                case .orderedAscending:
                    mergedArray.append(leftElement)
                    leftIndex += 1

                case .orderedSame:
                    mergedArray.append(leftElement)
                    leftIndex += 1
                    mergedArray.append(rightElement)
                    rightIndex += 1

                case .orderedDescending:
                    mergedArray.append(rightElement)
                    rightIndex += 1
                }
            }

            while leftIndex < count {
                mergedArray.append(self[leftIndex])
                leftIndex += 1
            }
            while rightIndex < array.count {
                mergedArray.append(array[rightIndex])
                rightIndex += 1
            }

            return mergedArray
        }
        return await withTaskCancellationHandler(operation: { await task.value }, onCancel: { task.cancel() })
    }

    private func threadedSorted(using comparators: [KeyPathComparator<Element>]) async -> Self {
        let task = Task { sorted(using: comparators) }
        return await withTaskCancellationHandler(operation: { await task.value }, onCancel: { task.cancel() })
    }

    private func chunked(into count: Int) -> [Self] {
        let size = self.count / count + 1
        return stride(from: 0, to: self.count, by: size).map {
            .init(self[$0 ..< Swift.min($0 + size, self.count)])
        }
    }
}