From 0916b5f9900144f456fbc20ef54b0d02c102e981 Mon Sep 17 00:00:00 2001
From: Max Sokolov <i@maxsokolov.net>
Date: Mon, 3 Oct 2016 13:48:59 +0300
Subject: [PATCH] fix actions

---
 Sources/TableRowAction.swift | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/Sources/TableRowAction.swift b/Sources/TableRowAction.swift
index c4bafdb..dac4b49 100644
--- a/Sources/TableRowAction.swift
+++ b/Sources/TableRowAction.swift
@@ -61,25 +61,40 @@ open class TableRowActionData<ItemType, CellType: ConfigurableCell> where CellTy
     }
 }
 
+private enum TableRowActionHandler<ItemType, CellType: ConfigurableCell> where CellType.T == ItemType, CellType: UITableViewCell {
+
+    case voidAction((TableRowActionData<ItemType, CellType>) -> Void)
+    case action((TableRowActionData<ItemType, CellType>) -> Any?)
+
+    func invoke(item: ItemType, cell: UITableViewCell?, path: IndexPath) -> Any? {
+        
+        switch self {
+        case .voidAction(let handler):
+            return handler(TableRowActionData(item: item, cell: cell as? CellType, path: path, userInfo: nil))
+        case .action(let handler):
+            return handler(TableRowActionData(item: item, cell: cell as? CellType, path: path, userInfo: nil))
+        }
+    }
+}
+
 open class TableRowAction<ItemType, CellType: ConfigurableCell> where CellType.T == ItemType, CellType: UITableViewCell {
 
     open let type: TableRowActionType
-    private let handler: ((_ data: TableRowActionData<ItemType, CellType>) -> Any?)?
-
+    private let handler: TableRowActionHandler<ItemType, CellType>
+    
     public init(_ type: TableRowActionType, handler: @escaping (_ data: TableRowActionData<ItemType, CellType>) -> Void) {
 
         self.type = type
-        self.handler = nil
-        //self.handler = handler
+        self.handler = .voidAction(handler)
     }
     
     public init<T>(_ type: TableRowActionType, handler: @escaping (_ data: TableRowActionData<ItemType, CellType>) -> T) {
 
         self.type = type
-        self.handler = handler
+        self.handler = .action(handler)
     }
 
     func invoke(item: ItemType, cell: UITableViewCell?, path: IndexPath) -> Any? {
-        return handler?(TableRowActionData(item: item, cell: cell as? CellType, path: path, userInfo: nil))
+        return handler.invoke(item: item, cell: cell, path: path)
     }
 }