Adds Publishers.ConcatMap

Author: ohitsdaniel

Sources/Operators/ConcatMap.swift

@@ -8,6 +8,7 @@
 
 #if canImport(Combine)
 import Combine
+import class Foundation.NSRecursiveLock
 
 @available(OSX 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *)
 public extension Publisher {
@@ -20,8 +21,145 @@ public extension Publisher {
     /// - returns: A publisher emitting the values of all emitted publishers in order.
     func concatMap<T, P>(
         _ transform: @escaping (Self.Output) -> P
-    ) -> Publishers.FlatMap<P, Self> where T == P.Output, P: Publisher, Self.Failure == P.Failure {
-        flatMap(maxPublishers: .max(1), transform)
+    ) -> Publishers.ConcatMap<P, Self> where T == P.Output, P: Publisher, Self.Failure == P.Failure {
+        return Publishers.ConcatMap(upstream: self, transform: transform)
+    }
+}
+
+@available(OSX 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *)
+public extension Publishers {
+    final class ConcatMap<NewPublisher, Upstream>: Publisher where NewPublisher: Publisher, Upstream: Publisher, NewPublisher.Failure == Upstream.Failure  {
+        public typealias Transform = (Upstream.Output) -> NewPublisher
+        public typealias Output = NewPublisher.Output
+        public typealias Failure = Upstream.Failure
+
+        public let upstream: Upstream
+        public let transform: Transform
+
+        public init(
+            upstream: Upstream,
+            transform: @escaping Transform
+        ) {
+            self.upstream = upstream
+            self.transform = transform
+        }
+
+        public func receive<S: Subscriber>(subscriber: S)
+        where Output == S.Input, Failure == S.Failure {
+            subscriber.receive(
+                subscription: Subscription(
+                    upstream: upstream,
+                    transform: transform,
+                    downstream: subscriber
+                )
+            )
+        }
+    }
+}
+
+// MARK: - Subscription
+@available(OSX 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *)
+private extension Publishers.ConcatMap {
+    final class Subscription<Downstream: Subscriber>: Combine.Subscription where Downstream.Input == Output, Downstream.Failure == Failure {
+        private var sink: Sink<Downstream>?
+
+        init(
+            upstream: Upstream,
+            transform: @escaping Transform,
+            downstream: Downstream
+        ) {
+            self.sink = Sink(
+                upstream: upstream,
+                downstream: downstream,
+                transform: { transform($0) }
+            )
+        }
+
+        func request(_ demand: Subscribers.Demand) {
+            sink?.demand(demand)
+        }
+
+        func cancel() {
+            sink = nil
+        }
+    }
+}
+
+// MARK: - Sink
+@available(OSX 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *)
+private extension Publishers.ConcatMap {
+    final class Sink<Downstream: Subscriber>: CombineExt.Sink<Upstream, Downstream>
+    where Downstream.Input == Output, Downstream.Failure == Failure {
+        private let lock = NSRecursiveLock()
+        private let transform: Transform
+        private var activePublisher: NewPublisher?
+        private var bufferedPublishers: [NewPublisher]
+        private var cancellables: Set<AnyCancellable>
+
+        init(
+            upstream: Upstream,
+            downstream: Downstream,
+            transform: @escaping Transform
+        ) {
+            self.transform = transform
+            self.bufferedPublishers = []
+            self.cancellables = []
+            super.init(
+                upstream: upstream,
+                downstream: downstream,
+                transformFailure: { $0 }
+            )
+        }
+
+        override func receive(_ input: Upstream.Output) -> Subscribers.Demand {
+            let mapped = transform(input)
+
+            lock.lock()
+            if activePublisher == nil {
+                lock.unlock()
+                setActivePublisher(mapped)
+            } else {
+                lock.unlock()
+                bufferedPublishers.append(mapped)
+            }
+
+            return .unlimited
+        }
+
+        private func setActivePublisher(_ publisher: NewPublisher) {
+            lock.lock()
+            defer { lock.unlock() }
+            activePublisher = publisher
+
+            publisher.sink(
+                receiveCompletion: { completion in
+                    switch completion {
+                    case .finished:
+                        self.lock.lock()
+                        guard let next = self.bufferedPublishers.first else {
+                            self.lock.unlock()
+                            return
+                        }
+                        self.bufferedPublishers.removeFirst()
+                        self.lock.unlock()
+                        self.setActivePublisher(next)
+                    case .failure(let error):
+                        self.receive(completion: .failure(error))
+                    }
+                },
+                receiveValue: { value in
+                    _ = self.buffer.buffer(value: value)
+                }
+            )
+            .store(in: &cancellables)
+        }
+    }
+}
+
+@available(OSX 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *)
+extension Publishers.ConcatMap.Subscription: CustomStringConvertible {
+    var description: String {
+        return "ConcatMap.Subscription<\(Downstream.Input.self), \(Downstream.Failure.self)>"
     }
 }
 #endif

Tests/ConcatMapTests.swift

@@ -26,10 +26,11 @@ final class ConcatMapTests: XCTestCase {
 
     func test_publishes_values_inOrder() {
         var receivedValues = [Int]()
-        let expectedValues = [1, 2, 3, 4]
+        let expectedValues = [1, 2, 4, 5, 6]
 
         let firstPublisher = P()
         let secondPublisher = P()
+        let thirdPublisher = P()
 
         let sut = PassthroughSubject<P, TestError>()
 
@@ -42,15 +43,20 @@ final class ConcatMapTests: XCTestCase {
 
         sut.send(firstPublisher)
         sut.send(secondPublisher)
+        sut.send(thirdPublisher)
 
         firstPublisher.send(1)
         firstPublisher.send(2)
+        // values sent onto the second publisher will be ignored as long as the first publisher hasn't completed
+        secondPublisher.send(3)
         firstPublisher.send(completion: .finished)
 
-        secondPublisher.send(3)
         secondPublisher.send(4)
+        secondPublisher.send(5)
         secondPublisher.send(completion: .finished)
 
+        thirdPublisher.send(6)
+
         XCTAssertEqual(expectedValues, receivedValues)
     }