PostCategories.diff

diff --git a/WordPress/Classes/ViewRelated/Post/Categories/PostCategoriesViewController.swift b/WordPress/Classes/ViewRelated/Post/Categories/PostCategoriesViewController.swift
index 7883b61e18..a6944041c9 100644
--- a/WordPress/Classes/ViewRelated/Post/Categories/PostCategoriesViewController.swift
+++ b/WordPress/Classes/ViewRelated/Post/Categories/PostCategoriesViewController.swift
@@ -13,13 +13,17 @@ import Foundation
 
 @objc class PostCategoriesViewController: UITableViewController {
     @objc var delegate: PostCategoriesViewControllerDelegate?
-    private var selectionMode: CategoriesSelectionMode
+
     private var blog: Blog
     private var originalSelection: [PostCategory]?
-    private var selectedCategories: [PostCategory]?
-    private var saveButtonItem: UIBarButtonItem?
+    private var selectionMode: CategoriesSelectionMode
+
     private var categories = [PostCategory]()
-    private var categoryIndentationDict = [String: NSNumber]()
+    private var categoryIndentationDict = [Int: Int]()
+    private var selectedCategories = [PostCategory]()
+
+    private var saveButtonItem: UIBarButtonItem?
+
     private var hasSyncedCategories = false
 
     @objc init(blog: Blog, currentSelection: [PostCategory]?, selectionMode: CategoriesSelectionMode) {
@@ -98,58 +102,53 @@ import Foundation
     }
 
     private func reloadCategories() {
-        if selectedCategories == nil {
-            selectedCategories = originalSelection
+        if selectedCategories.isEmpty {
+            selectedCategories = originalSelection ?? []
         }
 
-        var categoryDict = [NSNumber: PostCategory]()
         // Sort categories by parent/child relationship
         let tree = WPCategoryTree(parent: nil)
-        if let sortedCategories = blog.sortedCategories() {
-            tree.getChildrenFromObjects(sortedCategories)
-        }
-
+        tree.getChildrenFromObjects(blog.sortedCategories() ?? [])
         categories = tree.getAllObjects()
+
+        var categoryDict = [Int: PostCategory]()
+        categoryIndentationDict = [:]
+
         categories.forEach { category in
-            categoryDict[category.categoryID] = category
-            let indentationLevel = indentationLevelForCategory(parentID: category.parentID, categoryCollection: categoryDict)
-            categoryIndentationDict[category.categoryID .stringValue] = NSNumber(integerLiteral: indentationLevel)
-        }
+            let categoryID = category.categoryID.intValue
+            let parentID = category.parentID.intValue
 
-        // Remove any previously selected category objects that are no longer available.
-        guard selectedCategories != nil else {
-            return
+            categoryDict[categoryID] = category
+
+            let indentationLevel = indentationLevelForCategory(parentID: parentID, categoryCollection: categoryDict)
+            categoryIndentationDict[categoryID] = indentationLevel
         }
 
-        var didUpdateSelectedCategories = false
-        let updatedSelectedCategories = selectedCategories
-        self.selectedCategories = updatedSelectedCategories?.filter { category in
+        // Remove any previously selected category objects that are no longer available.
+        let selectedCategories = self.selectedCategories
+        self.selectedCategories = selectedCategories.filter { category in
             if let sortedCategories = blog.sortedCategories() as? [PostCategory], sortedCategories.contains(category), !category.isDeleted {
                 return true
-            } else {
-                didUpdateSelectedCategories = true
-                return false
             }
+
+            return false
         }
 
         // Notify the delegate of any changes for selectedCategories.
-        if didUpdateSelectedCategories {
-            delegate?.postCategoriesViewController?(self, didUpdateSelectedCategories: NSSet(array: selectedCategories!))
+        if selectedCategories.count != self.selectedCategories.count {
+            delegate?.postCategoriesViewController?(self, didUpdateSelectedCategories: NSSet(array: self.selectedCategories))
         }
 
         tableView.reloadData()
     }
 
-    private func indentationLevelForCategory(parentID: NSNumber, categoryCollection: [NSNumber: PostCategory]) -> Int {
-        if parentID.intValue == 0 {
+    private func indentationLevelForCategory(parentID: Int, categoryCollection: [Int: PostCategory]) -> Int {
+        guard parentID != 0,
+              let category = categoryCollection[parentID] else {
             return 0
         }
 
-        if let category = categoryCollection[parentID] {
-            return indentationLevelForCategory(parentID: category.parentID, categoryCollection: categoryCollection) + 1
-        } else {
-            return 0
-        }
+        return indentationLevelForCategory(parentID: category.parentID.intValue, categoryCollection: categoryCollection) + 1
     }
 
     //tableView
@@ -159,6 +158,7 @@ import Foundation
     }
 
     override func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
+
         var result = categories.count
         if selectionMode == .parent {
             result = result + 1
@@ -180,8 +180,8 @@ import Foundation
                 WPStyleGuide.configureTableViewDestructiveActionCell(cell)
                 cell.textLabel?.textAlignment = .natural
                 cell.textLabel?.text = NSLocalizedString("No Category", comment: "Text shown (to select no-category) in the parent-category-selection screen when creating a new category.")
-                if selectedCategories == nil {
-                    cell.accessoryType = .checkmark
+                if selectedCategories.isEmpty {
+                    cell.accessoryType = selectedCategories.isEmpty ? .checkmark : .none
                 } else {
                     cell.accessoryType = .none
                 }
@@ -192,17 +192,15 @@ import Foundation
         }
 
         let category = categories[row]
-        let indentationLevel = categoryIndentationDict[category.categoryID .stringValue]?.intValue
+        let indentationLevel = categoryIndentationDict[category.categoryID.intValue]
         cell.indentationLevel = indentationLevel ?? 0
         cell.indentationWidth = Constants.categoryCellIndentation
         cell.textLabel?.text = category.categoryName.stringByDecodingXMLCharacters()
         WPStyleGuide.configureTableViewCell(cell)
-        if let selectedCategories = selectedCategories {
-            if selectedCategories.contains(category) {
-                 cell.accessoryType = .checkmark
-            } else {
-                cell.accessoryType = .none
-            }
+        if selectedCategories.contains(category) {
+             cell.accessoryType = .checkmark
+        } else {
+            cell.accessoryType = .none
         }
 
         return cell
@@ -228,26 +226,27 @@ import Foundation
             }
         case .post:
             category = categories[row]
-            if let category = category, selectedCategories != nil {
-                if selectedCategories!.contains(category) {
-                    selectedCategories!.remove(at: selectedCategories!.firstIndex(of: category)!)
+            if let category = category {
+                if selectedCategories.contains(category),
+                   let index = selectedCategories.firstIndex(of: category) {
+                    selectedCategories.remove(at: index)
                     tableView.cellForRow(at: indexPath)?.accessoryType = .none
                 } else {
-                    selectedCategories?.append(category)
+                    selectedCategories.append(category)
                     tableView.cellForRow(at: indexPath)?.accessoryType = .checkmark
                 }
 
-                delegate?.postCategoriesViewController?(self, didUpdateSelectedCategories: NSSet(array: selectedCategories!))
+                delegate?.postCategoriesViewController?(self, didUpdateSelectedCategories: NSSet(array: selectedCategories))
             }
 
         case .blogDefault:
             category = categories[row]
-            if let category = category, selectedCategories != nil {
-                if selectedCategories!.contains(category) {
+            if let category = category {
+                if selectedCategories.contains(category) {
                     return
                 }
-                selectedCategories!.removeAll()
-                selectedCategories!.append(category)
+                selectedCategories.removeAll()
+                selectedCategories.append(category)
                 tableView.reloadData()
                 delegate?.postCategoriesViewController?(self, didSelectCategory: category)
             }
@@ -266,13 +265,11 @@ extension PostCategoriesViewController: WPAddPostCategoryViewControllerDelegate
     func addPostCategoryViewController(_ controller: WPAddPostCategoryViewController, didAdd category: PostCategory) {
         switch selectionMode {
         case .post, .parent:
-            selectedCategories?.append(category)
-            if let selectedCategories = selectedCategories {
-                delegate?.postCategoriesViewController?(self, didUpdateSelectedCategories: NSSet(array: selectedCategories))
-            }
+            selectedCategories.append(category)
+            delegate?.postCategoriesViewController?(self, didUpdateSelectedCategories: NSSet(array: selectedCategories))
         case .blogDefault:
-            selectedCategories?.removeAll()
-            selectedCategories?.append(category)
+            selectedCategories.removeAll()
+            selectedCategories.append(category)
             delegate?.postCategoriesViewController?(self, didSelectCategory: category)
         }