Add removeProgress() & removeThen()

inamiy · inamiy · commit 16a2e60158e2 · 2015-05-08T10:26:44.000+09:00
diff --git a/SwiftTask.xcodeproj/project.pbxproj b/SwiftTask.xcodeproj/project.pbxproj
@@ -8,6 +8,8 @@
 
 /* Begin PBXBuildFile section */
 		1F20250219ADA8FD00DE0495 /* BasicTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1F20250119ADA8FD00DE0495 /* BasicTests.swift */; };
+		1F218D5B1AFC3FFD00C849FF /* RemoveHandlerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1F218D5A1AFC3FFD00C849FF /* RemoveHandlerTests.swift */; };
+		1F218D5C1AFC3FFD00C849FF /* RemoveHandlerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1F218D5A1AFC3FFD00C849FF /* RemoveHandlerTests.swift */; };
 		1F46DEDA199EDF1000F97868 /* SwiftTask.h in Headers */ = {isa = PBXBuildFile; fileRef = 1F46DED9199EDF1000F97868 /* SwiftTask.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		1F46DEFB199EDF8100F97868 /* SwiftTask.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1F46DEFA199EDF8100F97868 /* SwiftTask.swift */; };
 		1F46DEFD199EE2C200F97868 /* _TestCase.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1F46DEFC199EE2C200F97868 /* _TestCase.swift */; };
@@ -37,6 +39,7 @@
 
 /* Begin PBXFileReference section */
 		1F20250119ADA8FD00DE0495 /* BasicTests.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = BasicTests.swift; sourceTree = "<group>"; };
+		1F218D5A1AFC3FFD00C849FF /* RemoveHandlerTests.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = RemoveHandlerTests.swift; sourceTree = "<group>"; };
 		1F46DED4199EDF1000F97868 /* SwiftTask.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = SwiftTask.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		1F46DED8199EDF1000F97868 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
 		1F46DED9199EDF1000F97868 /* SwiftTask.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = SwiftTask.h; sourceTree = "<group>"; };
@@ -144,6 +147,7 @@
 				1F46DEE3199EDF1000F97868 /* SwiftTaskTests.swift */,
 				48511C5A19C17563002FE03C /* RetainCycleTests.swift */,
 				485C31F01A1D619A00040DA3 /* TypeInferenceTests.swift */,
+				1F218D5A1AFC3FFD00C849FF /* RemoveHandlerTests.swift */,
 				1F5FA35619A374E600975FB9 /* AlamofireTests.swift */,
 				1F46DEE1199EDF1000F97868 /* Supporting Files */,
 			);
@@ -345,6 +349,7 @@
 			files = (
 				1F20250219ADA8FD00DE0495 /* BasicTests.swift in Sources */,
 				1FF52EB41A4C395A00B4BA28 /* _InterruptableTask.swift in Sources */,
+				1F218D5B1AFC3FFD00C849FF /* RemoveHandlerTests.swift in Sources */,
 				1F6A8CA319A4E4F200369A5D /* SwiftTaskTests.swift in Sources */,
 				1F4C76A41AD8CF40004E47C1 /* AlamofireTests.swift in Sources */,
 				485C31F11A1D619A00040DA3 /* TypeInferenceTests.swift in Sources */,
@@ -359,6 +364,7 @@
 			files = (
 				4822F0DE19D00B2300F5F572 /* SwiftTaskTests.swift in Sources */,
 				4822F0DD19D00B2300F5F572 /* BasicTests.swift in Sources */,
+				1F218D5C1AFC3FFD00C849FF /* RemoveHandlerTests.swift in Sources */,
 				1FF52EB51A4C395A00B4BA28 /* _InterruptableTask.swift in Sources */,
 				1F4C76A51AD8CF41004E47C1 /* AlamofireTests.swift in Sources */,
 				485C31F21A1D619A00040DA3 /* TypeInferenceTests.swift in Sources */,
diff --git a/SwiftTask/SwiftTask.swift b/SwiftTask/SwiftTask.swift
@@ -355,7 +355,13 @@ public class Task<Progress, Value, Error>: Printable
     ///
     public func progress(progressClosure: ProgressTuple -> Void) -> Task
     {
-        self._machine.addProgressTupleHandler(progressClosure)
+        var token: HandlerToken? = nil
+        return self.progress(&token, progressClosure)
+    }
+    
+    public func progress(inout token: HandlerToken?, _ progressClosure: ProgressTuple -> Void) -> Task
+    {
+        self._machine.addProgressTupleHandler(&token, progressClosure)
         
         return self
     }
@@ -367,7 +373,13 @@ public class Task<Progress, Value, Error>: Printable
     ///
     public func then<Value2>(thenClosure: (Value?, ErrorInfo?) -> Value2) -> Task<Progress, Value2, Error>
     {
-        return self.then { (value: Value?, errorInfo: ErrorInfo?) -> Task<Progress, Value2, Error> in
+        var token: HandlerToken? = nil
+        return self.then(&token, thenClosure)
+    }
+    
+    public func then<Value2>(inout token: HandlerToken?, _ thenClosure: (Value?, ErrorInfo?) -> Value2) -> Task<Progress, Value2, Error>
+    {
+        return self.then(&token) { (value: Value?, errorInfo: ErrorInfo?) -> Task<Progress, Value2, Error> in
             return Task<Progress, Value2, Error>(value: thenClosure(value, errorInfo))
         }
     }
@@ -378,6 +390,12 @@ public class Task<Progress, Value, Error>: Printable
     /// - e.g. task.then { value, errorInfo -> NextTaskType in ... }
     ///
     public func then<Progress2, Value2>(thenClosure: (Value?, ErrorInfo?) -> Task<Progress2, Value2, Error>) -> Task<Progress2, Value2, Error>
+    {
+        var token: HandlerToken? = nil
+        return self.then(&token, thenClosure)
+    }
+    
+    public func then<Progress2, Value2>(inout token: HandlerToken?, _ thenClosure: (Value?, ErrorInfo?) -> Task<Progress2, Value2, Error>) -> Task<Progress2, Value2, Error>
     {
         return Task<Progress2, Value2, Error> { [unowned self] newMachine, progress, fulfill, _reject, configure in
             
@@ -389,8 +407,8 @@ public class Task<Progress, Value, Error>: Printable
             // This is especially important for ReactKit's `deinitSignal` behavior.
             //
             let selfMachine = self._machine
-
-            self._then {
+            
+            self._then(&token) {
                 let innerTask = thenClosure(selfMachine.value, selfMachine.errorInfo)
                 _bindInnerTask(innerTask, newMachine, progress, fulfill, _reject, configure)
             }
@@ -399,13 +417,13 @@ public class Task<Progress, Value, Error>: Printable
     }
 
     /// invokes `completionHandler` "now" or "in the future"
-    private func _then(completionHandler: Void -> Void)
+    private func _then(inout token: HandlerToken?, _ completionHandler: Void -> Void)
     {
         switch self.state {
             case .Fulfilled, .Rejected, .Cancelled:
                 completionHandler()
             default:
-                self._machine.addCompletionHandler(completionHandler)
+                self._machine.addCompletionHandler(&token, completionHandler)
         }
     }
     
@@ -416,7 +434,13 @@ public class Task<Progress, Value, Error>: Printable
     ///
     public func success<Value2>(successClosure: Value -> Value2) -> Task<Progress, Value2, Error>
     {
-        return self.success { (value: Value) -> Task<Progress, Value2, Error> in
+        var token: HandlerToken? = nil
+        return self.success(&token, successClosure)
+    }
+    
+    public func success<Value2>(inout token: HandlerToken?, _ successClosure: Value -> Value2) -> Task<Progress, Value2, Error>
+    {
+        return self.success(&token) { (value: Value) -> Task<Progress, Value2, Error> in
             return Task<Progress, Value2, Error>(value: successClosure(value))
         }
     }
@@ -427,13 +451,19 @@ public class Task<Progress, Value, Error>: Printable
     /// - e.g. task.success { value -> NextTaskType in ... }
     ///
     public func success<Progress2, Value2>(successClosure: Value -> Task<Progress2, Value2, Error>) -> Task<Progress2, Value2, Error>
+    {
+        var token: HandlerToken? = nil
+        return self.success(&token, successClosure)
+    }
+    
+    public func success<Progress2, Value2>(inout token: HandlerToken?, _ successClosure: Value -> Task<Progress2, Value2, Error>) -> Task<Progress2, Value2, Error>
     {
         return Task<Progress2, Value2, Error> { [unowned self] newMachine, progress, fulfill, _reject, configure in
             
             let selfMachine = self._machine
             
             // NOTE: using `self._then()` + `selfMachine` instead of `self.then()` will reduce Task allocation
-            self._then {
+            self._then(&token) {
                 if let value = selfMachine.value {
                     let innerTask = successClosure(value)
                     _bindInnerTask(innerTask, newMachine, progress, fulfill, _reject, configure)
@@ -454,7 +484,13 @@ public class Task<Progress, Value, Error>: Printable
     ///
     public func failure(failureClosure: ErrorInfo -> Value) -> Task
     {
-        return self.failure { (errorInfo: ErrorInfo) -> Task in
+        var token: HandlerToken? = nil
+        return self.failure(&token, failureClosure)
+    }
+    
+    public func failure(inout token: HandlerToken?, _ failureClosure: ErrorInfo -> Value) -> Task
+    {
+        return self.failure(&token) { (errorInfo: ErrorInfo) -> Task in
             return Task(value: failureClosure(errorInfo))
         }
     }
@@ -466,12 +502,18 @@ public class Task<Progress, Value, Error>: Printable
     /// - e.g. task.failure { error, isCancelled -> NextTaskType in ... }
     ///
     public func failure<Progress2>(failureClosure: ErrorInfo -> Task<Progress2, Value, Error>) -> Task<Progress2, Value, Error>
+    {
+        var token: HandlerToken? = nil
+        return self.failure(&token, failureClosure)
+    }
+    
+    public func failure<Progress2>(inout token: HandlerToken?, _ failureClosure: ErrorInfo -> Task<Progress2, Value, Error>) -> Task<Progress2, Value, Error>
     {
         return Task<Progress2, Value, Error> { [unowned self] newMachine, progress, fulfill, _reject, configure in
             
             let selfMachine = self._machine
             
-            self._then {
+            self._then(&token) {
                 if let value = selfMachine.value {
                     fulfill(value)
                 }
@@ -503,6 +545,20 @@ public class Task<Progress, Value, Error>: Printable
     {
         return self._machine.handleCancel(error: error)
     }
+    
+    // MARK: Remove Handlers
+    
+    public func removeProgress(handlerToken: HandlerToken)
+    {
+        return self._machine.removeProgressTupleHandler(handlerToken)
+    }
+    
+    /// NOTE: `task.removeThen(token)` will force `let task2 = task.then(&token)` to deinit immediately and tries cancellation if it is still running.
+    public func removeThen(handlerToken: HandlerToken)
+    {
+        return self._machine.removeCompletionHandler(handlerToken)
+    }
+    
 }
 
 // MARK: - Helper
diff --git a/SwiftTask/_StateMachine.swift b/SwiftTask/_StateMachine.swift
@@ -30,8 +30,8 @@ internal class _StateMachine<Progress, Value, Error>
     /// and will be set to `nil` afterward
     internal var initResumeClosure: (Void -> Void)?
     
-    internal private(set) lazy var progressTupleHandlers: [ProgressTupleHandler] = []
-    internal private(set) lazy var completionHandlers: [Void -> Void] = []
+    internal private(set) lazy var progressTupleHandlers = _Handlers<ProgressTupleHandler>()
+    internal private(set) lazy var completionHandlers = _Handlers<Void -> Void>()
     
     internal let configuration = TaskConfiguration()
     
@@ -41,20 +41,30 @@ internal class _StateMachine<Progress, Value, Error>
         self.state = paused ? .Paused : .Running
     }
     
-    internal func addProgressTupleHandler(progressTupleHandler: ProgressTupleHandler)
+    internal func addProgressTupleHandler(inout token: HandlerToken?, _ progressTupleHandler: ProgressTupleHandler)
     {
         if self.state == .Running || self.state == .Paused {
-            self.progressTupleHandlers.append(progressTupleHandler)
+            token = self.progressTupleHandlers.append(progressTupleHandler)
         }
     }
     
-    internal func addCompletionHandler(completionHandler: Void -> Void)
+    internal func removeProgressTupleHandler(handlerToken: HandlerToken)
+    {
+        self.progressTupleHandlers.remove(handlerToken)
+    }
+    
+    internal func addCompletionHandler(inout token: HandlerToken?, _ completionHandler: Void -> Void)
     {
         if self.state == .Running || self.state == .Paused {
-            self.completionHandlers.append(completionHandler)
+            token = self.completionHandlers.append(completionHandler)
         }
     }
     
+    internal func removeCompletionHandler(handlerToken: HandlerToken)
+    {
+        self.completionHandlers.remove(handlerToken)
+    }
+    
     internal func handleProgress(progress: Progress)
     {
         if self.state == .Running {
@@ -185,4 +195,43 @@ internal class _StateMachine<Progress, Value, Error>
         self.initResumeClosure = nil
         self.progress = nil
     }
+}
+
+//--------------------------------------------------
+// MARK: - Utility
+//--------------------------------------------------
+
+public struct HandlerToken
+{
+    internal let key: Int
+}
+
+internal struct _Handlers<T>: SequenceType
+{
+    private var currentKey: Int = 0
+    private var elements = [Int : T]()
+    
+    internal mutating func append(value: T) -> HandlerToken
+    {
+        self.currentKey = self.currentKey &+ 1
+        
+        self.elements[self.currentKey] = value
+        
+        return HandlerToken(key: self.currentKey)
+    }
+    
+    internal mutating func remove(token: HandlerToken)
+    {
+        self.elements.removeValueForKey(token.key)
+    }
+    
+    internal mutating func removeAll(keepCapacity: Bool = false)
+    {
+        self.elements.removeAll(keepCapacity: keepCapacity)
+    }
+    
+    internal func generate() -> GeneratorOf<T>
+    {
+        return GeneratorOf(self.elements.values.generate())
+    }
 }
diff --git a/SwiftTaskTests/RemoveHandlerTests.swift b/SwiftTaskTests/RemoveHandlerTests.swift
