Add input operation

5 files changed, 281 insertions, 36 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 6933de9e76..b1535c0fa4 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -259,6 +259,14 @@
 		58D67A0A26D7AE3300557C3C /* OSLogHandler.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5823FA4F26CA690600283BF8 /* OSLogHandler.swift */; };
 		58DF28A52417CB4B00E836B0 /* AppStorePaymentManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58DF28A42417CB4B00E836B0 /* AppStorePaymentManager.swift */; };
 		58DF5B7F2852778600E92647 /* OperationSmokeTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58DF5B7E2852778600E92647 /* OperationSmokeTests.swift */; };
+		58DF5B742851FF3F00E92647 /* InputOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58DF5B732851FF3F00E92647 /* InputOperation.swift */; };
+		58DF5B762852108E00E92647 /* InputInjectionBuilder.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58DF5B752852108E00E92647 /* InputInjectionBuilder.swift */; };
+		58DF5B782852178600E92647 /* OperationInputInjectionTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58DF5B772852178600E92647 /* OperationInputInjectionTests.swift */; };
+		58DF5B79285217F300E92647 /* TransformOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58059DDB28465E8F002B1049 /* TransformOperation.swift */; };
+		58DF5B7A285217FA00E92647 /* InputInjectionBuilder.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58DF5B752852108E00E92647 /* InputInjectionBuilder.swift */; };
+		58DF5B7B285217FE00E92647 /* InputOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58DF5B732851FF3F00E92647 /* InputOperation.swift */; };
+		58DF5B7C28521A9F00E92647 /* ResultOperation+Output.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58059DDF2846823E002B1049 /* ResultOperation+Output.swift */; };
+		58DF5B7D28521AAC00E92647 /* OutputOperation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58059DDD28468158002B1049 /* OutputOperation.swift */; };
 		58E0A98827C8F46300FE6BDD /* Tunnel.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58E0A98727C8F46300FE6BDD /* Tunnel.swift */; };
 		58E20771274672CA00DE5D77 /* LaunchViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58E20770274672CA00DE5D77 /* LaunchViewController.swift */; };
 		58E6771F24ADFE7800AA26E7 /* SettingsNavigationController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58E6771E24ADFE7800AA26E7 /* SettingsNavigationController.swift */; };
@@ -537,6 +545,9 @@
 		58D0C7A023F1CECF00FE9BA7 /* MullvadVPNScreenshots.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = MullvadVPNScreenshots.swift; sourceTree = "<group>"; };
 		58DF28A42417CB4B00E836B0 /* AppStorePaymentManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppStorePaymentManager.swift; sourceTree = "<group>"; };
 		58DF5B7E2852778600E92647 /* OperationSmokeTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OperationSmokeTests.swift; sourceTree = "<group>"; };
+		58DF5B732851FF3F00E92647 /* InputOperation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = InputOperation.swift; sourceTree = "<group>"; };
+		58DF5B752852108E00E92647 /* InputInjectionBuilder.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = InputInjectionBuilder.swift; sourceTree = "<group>"; };
+		58DF5B772852178600E92647 /* OperationInputInjectionTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OperationInputInjectionTests.swift; sourceTree = "<group>"; };
 		58E0A98727C8F46300FE6BDD /* Tunnel.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Tunnel.swift; sourceTree = "<group>"; };
 		58E20770274672CA00DE5D77 /* LaunchViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LaunchViewController.swift; sourceTree = "<group>"; };
 		58E6771E24ADFE7800AA26E7 /* SettingsNavigationController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsNavigationController.swift; sourceTree = "<group>"; };
@@ -640,14 +651,16 @@
 				589D28782846250500F9A7B3 /* AsyncOperationQueue.swift */,
 				589D287F28462CB000F9A7B3 /* BackgroundObserver.swift */,
 				589D28812846306C00F9A7B3 /* GroupOperation.swift */,
+				58DF5B752852108E00E92647 /* InputInjectionBuilder.swift */,
+				58DF5B732851FF3F00E92647 /* InputOperation.swift */,
 				5840BE34279EDB16002836BA /* OperationCompletion.swift */,
 				589D28772846250500F9A7B3 /* OperationCondition.swift */,
 				589D28792846250500F9A7B3 /* OperationObserver.swift */,
 				58059DDD28468158002B1049 /* OutputOperation.swift */,
 				5820675D26E6839900655B05 /* PresentAlertOperation.swift */,
 				5846226426E0D9630035F7C2 /* ProductsRequestOperation.swift */,
-				58F7D26427EB50A300E4D821 /* ResultOperation.swift */,
 				5842102D282D3FC200F24E46 /* ResultBlockOperation.swift */,
+				58F7D26427EB50A300E4D821 /* ResultOperation.swift */,
 				58059DE128468255002B1049 /* ResultOperation+Fallible.swift */,
 				58059DDF2846823E002B1049 /* ResultOperation+Output.swift */,
 				58059DDB28465E8F002B1049 /* TransformOperation.swift */,
@@ -802,6 +815,7 @@
 			isa = PBXGroup;
 			children = (
 				58DF5B7E2852778600E92647 /* OperationSmokeTests.swift */,
+				58DF5B772852178600E92647 /* OperationInputInjectionTests.swift */,
 				583E1E292848DF67004838B3 /* OperationObserverTests.swift */,
 				580CBFB72848D503007878F0 /* OperationConditionTests.swift */,
 				582AE3112440CA0D00E6733A /* AccountTokenInputTests.swift */,
@@ -1226,6 +1240,7 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				58DF5B7A285217FA00E92647 /* InputInjectionBuilder.swift in Sources */,
 				582AE3132440CA2700E6733A /* AccountTokenInput.swift in Sources */,
 				58CAF4EF26025954007C5886 /* SimulatorTunnelProvider.swift in Sources */,
 				583E1E232848DE1C004838B3 /* OperationCompletion.swift in Sources */,
@@ -1248,10 +1263,12 @@
 				5896AE88246D7FAF005B36CB /* CustomDateComponentsFormatting.swift in Sources */,
 				5857F23824C8446700CF6F47 /* AsyncBlockOperation.swift in Sources */,
 				582AE3122440CA0D00E6733A /* AccountTokenInputTests.swift in Sources */,
+				58DF5B782852178600E92647 /* OperationInputInjectionTests.swift in Sources */,
 				585DA8A526B14EE000B8C587 /* PacketTunnelStatus.swift in Sources */,
 				58B0A2A9238EE6A100BC001D /* RelayConstraints.swift in Sources */,
 				583E1E282848DE1C004838B3 /* BackgroundObserver.swift in Sources */,
 				5807E2C2243203D000F5FF30 /* StringTests.swift in Sources */,
+				58DF5B7B285217FE00E92647 /* InputOperation.swift in Sources */,
 				5819C2142726CC8D00D6EC38 /* DataSourceSnapshotTests.swift in Sources */,
 				585DA8A326B14E0D00B8C587 /* ServerRelaysResponse.swift in Sources */,
 				583E1E1E2848DE1C004838B3 /* OperationCondition.swift in Sources */,
@@ -1259,11 +1276,14 @@
 				5820676226E75D8500655B05 /* REST.swift in Sources */,
 				58A8055E2716EA6700681642 /* AnyIPAddress.swift in Sources */,
 				583E1E1B2848DE1C004838B3 /* ResultOperation+Fallible.swift in Sources */,
+				58DF5B7D28521AAC00E92647 /* OutputOperation.swift in Sources */,
+				58DF5B79285217F300E92647 /* TransformOperation.swift in Sources */,
 				5857F23024C843ED00CF6F47 /* ChainedError.swift in Sources */,
 				58A8BE81239FBE62006B74AC /* IPEndpoint.swift in Sources */,
 				58871D2325D535D2002297FA /* IPAddressRange+Codable.swift in Sources */,
 				580CBFB82848D503007878F0 /* OperationConditionTests.swift in Sources */,
 				583E1E202848DE1C004838B3 /* OperationObserver.swift in Sources */,
+				58DF5B7C28521A9F00E92647 /* ResultOperation+Output.swift in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -1305,6 +1325,7 @@
 				58F1311527E0B2AB007AC5BC /* Result+Extensions.swift in Sources */,
 				585DA88426B0270700B8C587 /* ServerRelaysResponse.swift in Sources */,
 				5875960726F36B3A00BF6711 /* TunnelIPCError.swift in Sources */,
+				58DF5B742851FF3F00E92647 /* InputOperation.swift in Sources */,
 				58F8AC0E25D3F8CE002BE0ED /* ProblemReportReviewViewController.swift in Sources */,
 				58CCA010224249A1004F3011 /* ConnectViewController.swift in Sources */,
 				58059DE02846823E002B1049 /* ResultOperation+Output.swift in Sources */,
@@ -1333,6 +1354,7 @@
 				58FEEB58260B662E00A621A8 /* AutomaticKeyboardResponder.swift in Sources */,
 				5846227326E22A160035F7C2 /* AppStorePaymentObserver.swift in Sources */,
 				58F2E146276A2C9900A79513 /* StopTunnelOperation.swift in Sources */,
+				58DF5B762852108E00E92647 /* InputInjectionBuilder.swift in Sources */,
 				585DA87A26B024F900B8C587 /* RelayCacheError.swift in Sources */,
 				5856D13727450A8A00DFD627 /* UIImage+TintColor.swift in Sources */,
 				58CB0EE024B86751001EF0D8 /* RESTAPIProxy.swift in Sources */,
diff --git a/ios/MullvadVPN/Operations/InputInjectionBuilder.swift b/ios/MullvadVPN/Operations/InputInjectionBuilder.swift
new file mode 100644
index 0000000000..14f2ed9ce3
--- /dev/null
+++ b/ios/MullvadVPN/Operations/InputInjectionBuilder.swift
@@ -0,0 +1,77 @@
+//
+//  InputInjectionBuilder.swift
+//  MullvadVPN
+//
+//  Created by pronebird on 09/06/2022.
+//  Copyright © 2022 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+protocol OperationInputContext {
+    associatedtype Input
+
+    func reduce() -> Input?
+}
+
+class InputInjectionBuilder<OperationType, Context> where OperationType: InputOperation {
+    typealias InputBlock = (inout Context) -> Void
+
+    private let operation: OperationType
+    private var context: Context
+    private var inputBlocks: [InputBlock] = []
+
+    init(operation: OperationType, context: Context) {
+        self.operation = operation
+        self.context = context
+    }
+
+    func inject<T>(
+        from dependency: T,
+        assignOutputTo keyPath: WritableKeyPath<Context, T.Output?>
+    ) -> Self
+        where T: OutputOperation
+    {
+        return inject(from: dependency) { context, output in
+            context[keyPath: keyPath] = output
+        }
+    }
+
+    func inject<T>(
+        from dependency: T,
+        via block: @escaping (inout Context, T.Output) -> Void
+    ) -> Self
+        where T: OutputOperation
+    {
+        inputBlocks.append { context in
+            if let output = dependency.output {
+                block(&context, output)
+            }
+        }
+
+        operation.addDependency(dependency)
+
+        return self
+    }
+
+    func reduce(_ reduceBlock: @escaping (Context) -> OperationType.Input?) {
+        operation.setInputBlock {
+            for inputBlock in self.inputBlocks {
+                inputBlock(&self.context)
+            }
+
+            return reduceBlock(self.context)
+        }
+    }
+}
+
+extension InputInjectionBuilder
+    where Context: OperationInputContext,
+          Context.Input == OperationType.Input
+{
+    func reduce() {
+        reduce { context in
+            return context.reduce()
+        }
+    }
+}
diff --git a/ios/MullvadVPN/Operations/InputOperation.swift b/ios/MullvadVPN/Operations/InputOperation.swift
new file mode 100644
index 0000000000..7f351d5ba4
--- /dev/null
+++ b/ios/MullvadVPN/Operations/InputOperation.swift
@@ -0,0 +1,47 @@
+//
+//  InputOperation.swift
+//  MullvadVPN
+//
+//  Created by pronebird on 09/06/2022.
+//  Copyright © 2022 Mullvad VPN AB. All rights reserved.
+//
+
+import Foundation
+
+protocol InputOperation: Operation {
+    associatedtype Input
+
+    var input: Input? { get }
+
+    func setInputBlock(_ block: @escaping () -> Input?)
+
+    func inject<T>(from dependency: T)
+        where T: OutputOperation, T.Output == Input
+
+    func inject<T>(from dependency: T, via block: @escaping (T.Output) -> Input)
+        where T: OutputOperation
+}
+
+extension InputOperation {
+    func inject<T>(from dependency: T) where T: OutputOperation, T.Output == Input {
+        inject(from: dependency, via: { $0 })
+    }
+
+    func inject<T>(from dependency: T, via block: @escaping (T.Output) -> Input)
+        where T: OutputOperation
+    {
+        setInputBlock {
+            return dependency.output.map { value in
+                return block(value)
+            }
+        }
+        addDependency(dependency)
+    }
+
+    func injectMany<Context>(context: Context) -> InputInjectionBuilder<Self, Context> {
+        return InputInjectionBuilder(
+            operation: self,
+            context: context
+        )
+    }
+}
diff --git a/ios/MullvadVPN/Operations/TransformOperation.swift b/ios/MullvadVPN/Operations/TransformOperation.swift
index b09b5780c1..2a2f5e2b5a 100644
--- a/ios/MullvadVPN/Operations/TransformOperation.swift
+++ b/ios/MullvadVPN/Operations/TransformOperation.swift
@@ -7,13 +7,34 @@
 
 import Foundation
 
-final class TransformOperation<Input, Output, Failure: Error>: ResultOperation<Output, Failure> {
-    typealias ExecutionBlock = ((Input, TransformOperation<Input, Output, Failure>) -> Void)
-    typealias ThrowingExecutionBlock = ((Input) throws -> Output)
+final class TransformOperation<Input, Output, Failure: Error>:
+    ResultOperation<Output, Failure>,
+    InputOperation
+{
+    typealias ExecutionBlock = (Input, TransformOperation<Input, Output, Failure>) -> Void
+    typealias ThrowingExecutionBlock = (Input) throws -> Output
+
+    typealias InputBlock = () -> Input?
+
+    private let nslock = NSLock()
+
+    private(set) var input: Input? {
+        get {
+            nslock.lock()
+            defer { nslock.unlock() }
+            return _input
+        }
+        set {
+            nslock.lock()
+            _input = newValue
+            nslock.unlock()
+        }
+    }
+
+    private var _input: Input?
+    private var inputBlock: InputBlock?
 
-    private var input: Input?
     private var executionBlock: ExecutionBlock?
-    private var configurationBlocks: [() -> Void] = []
     private var cancellationBlocks: [() -> Void] = []
 
     init(
@@ -22,38 +43,35 @@ final class TransformOperation<Input, Output, Failure: Error>: ResultOperation<O
         block: ExecutionBlock? = nil
     )
     {
-        self.input = input
-        self.executionBlock = block
+        _input = input
+        executionBlock = block
 
         super.init(dispatchQueue: dispatchQueue)
     }
 
-    convenience init(
-        dispatchQueue: DispatchQueue?,
+    init(
+        dispatchQueue: DispatchQueue? = nil,
         input: Input? = nil,
-        block: @escaping ThrowingExecutionBlock
+        throwingBlock: @escaping ThrowingExecutionBlock
     )
     {
-        self.init(
-            dispatchQueue: dispatchQueue,
-            input: input,
-            block: Self.wrapThrowingBlock(block)
-        )
+        _input = input
+        executionBlock = Self.wrapThrowingBlock(throwingBlock)
+
+        super.init(dispatchQueue: dispatchQueue)
     }
 
     override func main() {
-        for configurationBlock in configurationBlocks {
-            configurationBlock()
-        }
+        let inputValue = inputBlock?()
 
-        configurationBlocks.removeAll()
+        input = inputValue
 
-        guard let input = input, let executionBlock = executionBlock else {
+        guard let inputValue = inputValue, let executionBlock = executionBlock else {
             finish(completion: .cancelled)
             return
         }
 
-        executionBlock(input, self)
+        executionBlock(inputValue, self)
     }
 
     override func operationDidCancel() {
@@ -70,6 +88,8 @@ final class TransformOperation<Input, Output, Failure: Error>: ResultOperation<O
         executionBlock = nil
     }
 
+    // MARK: - Block handlers
+
     func setExecutionBlock(_ block: @escaping ExecutionBlock) {
         dispatchQueue.async {
             assert(!self.isExecuting && !self.isFinished)
@@ -91,22 +111,12 @@ final class TransformOperation<Input, Output, Failure: Error>: ResultOperation<O
         }
     }
 
-    func inject<T>(from dependency: T) where T: OutputOperation, T.Output == Input {
-        inject(from: dependency, via: { $0 })
-    }
+    // MARK: - Input injection
 
-    func inject<T>(from dependency: T, via block: @escaping (T.Output) -> Input) where T: OutputOperation {
+    func setInputBlock(_ block: @escaping () -> Input?) {
         dispatchQueue.async {
-            self.configurationBlocks.append { [weak self] in
-                guard let self = self else { return }
-
-                if let output = dependency.output {
-                    self.input = block(output)
-                }
-            }
-
+            self.inputBlock = block
         }
-        addDependency(dependency)
     }
 
     private class func wrapThrowingBlock(_ executionBlock: @escaping ThrowingExecutionBlock) -> ExecutionBlock {
@@ -122,6 +132,5 @@ final class TransformOperation<Input, Output, Failure: Error>: ResultOperation<O
             }
         }
     }
-
 }
 
diff --git a/ios/MullvadVPNTests/OperationInputInjectionTests.swift b/ios/MullvadVPNTests/OperationInputInjectionTests.swift
new file mode 100644
index 0000000000..d62214af42
--- /dev/null
+++ b/ios/MullvadVPNTests/OperationInputInjectionTests.swift
@@ -0,0 +1,90 @@
+//
+//  OperationInputInjectionTests.swift
+//  MullvadVPNTests
+//
+//  Created by pronebird on 09/06/2022.
+//  Copyright © 2022 Mullvad VPN AB. All rights reserved.
+//
+
+import XCTest
+
+class OperationInputInjectionTests: XCTestCase {
+
+    func testInject() throws {
+        let provider = ResultBlockOperation<Int, Error> {
+            return 1
+        }
+
+        let consumer = TransformOperation<Int, Int, Error> { input in
+            return input + 1
+        }
+
+        consumer.inject(from: provider)
+
+        let operationQueue = AsyncOperationQueue()
+
+        operationQueue.addOperations([provider, consumer], waitUntilFinished: true)
+
+        XCTAssertEqual(consumer.output, 2)
+    }
+
+    func testInjectVia() throws {
+        let provider = ResultBlockOperation<Int, Error> {
+            return 1
+        }
+
+        let consumer = TransformOperation<String, Int, Error> { input in
+            return Int(input)!
+        }
+
+        consumer.inject(from: provider) { output in
+            return "\(output)"
+        }
+
+        let operationQueue = AsyncOperationQueue()
+
+        operationQueue.addOperations([provider, consumer], waitUntilFinished: true)
+
+        XCTAssertEqual(consumer.output, 1)
+    }
+
+    func testInjectMany() throws {
+        struct Context: OperationInputContext {
+            var a: Int?
+            var b: Int?
+
+            func reduce() -> Int? {
+                guard let a = a, let b = b else { return nil }
+
+                return a + b
+            }
+        }
+
+        let operationQueue = AsyncOperationQueue()
+
+        let providerA = ResultBlockOperation<Int, Error> {
+            return 1
+        }
+
+        let providerB = ResultBlockOperation<Int, Error> {
+            return 2
+        }
+
+        let consumer = TransformOperation<Int, String, Error> { input in
+            return "\(input)"
+        }
+
+        consumer.injectMany(context: Context())
+            .inject(from: providerA, assignOutputTo: \.a)
+            .inject(from: providerB, assignOutputTo: \.b)
+            .reduce()
+
+        operationQueue.addOperations(
+            [providerA, providerB, consumer],
+            waitUntilFinished: true
+        )
+
+        XCTAssertEqual(consumer.output, "3")
+    }
+
+}