From be84fc5a6bbe9d2c4054c3e61bb45bf9d431f88d Mon Sep 17 00:00:00 2001
From: Maurice Parker <mo@vincode.io>
Date: Wed, 11 Sep 2019 16:53:27 -0500
Subject: [PATCH] Simplify the section and folder expand/collapse logic

---
 iOS/MasterFeed/MasterFeedViewController.swift | 26 +++---
 iOS/SceneCoordinator.swift                    | 88 ++++---------------
 2 files changed, 28 insertions(+), 86 deletions(-)

diff --git a/iOS/MasterFeed/MasterFeedViewController.swift b/iOS/MasterFeed/MasterFeedViewController.swift
index 52216c544..66af80f01 100644
--- a/iOS/MasterFeed/MasterFeedViewController.swift
+++ b/iOS/MasterFeed/MasterFeedViewController.swift
@@ -364,11 +364,11 @@ class MasterFeedViewController: UITableViewController, UndoableCommandRunner {
 		
 		if sectionNode.isExpanded {
 			headerView.disclosureExpanded = false
-			coordinator.collapseSection(sectionIndex)
+			coordinator.collapse(sectionNode)
 			self.applyChanges(animate: true)
 		} else {
 			headerView.disclosureExpanded = true
-			coordinator.expandSection(sectionIndex)
+			coordinator.expand(sectionNode)
 			self.applyChanges(animate: true)
 		}
 		
@@ -408,8 +408,8 @@ class MasterFeedViewController: UITableViewController, UndoableCommandRunner {
 	}
 	
 	@objc func expandSelectedRows(_ sender: Any?) {
-		if let indexPath = coordinator.currentFeedIndexPath {
-			coordinator.expandFolder(indexPath)
+		if let indexPath = coordinator.currentFeedIndexPath, let node = dataSource.itemIdentifier(for: indexPath) {
+			coordinator.expand(node)
 			self.applyChanges(animate: true) {
 				self.reloadAllVisibleCells()
 			}
@@ -417,8 +417,8 @@ class MasterFeedViewController: UITableViewController, UndoableCommandRunner {
 	}
 	
 	@objc func collapseSelectedRows(_ sender: Any?) {
-		if let indexPath = coordinator.currentFeedIndexPath {
-			coordinator.collapseFolder(indexPath)
+		if let indexPath = coordinator.currentFeedIndexPath, let node = dataSource.itemIdentifier(for: indexPath) {
+			coordinator.collapse(node)
 			self.applyChanges(animate: true) {
 				self.reloadAllVisibleCells()
 			}
@@ -478,7 +478,7 @@ class MasterFeedViewController: UITableViewController, UndoableCommandRunner {
 		}
 		
 		if !sectionNode.isExpanded {
-			coordinator.expandSection(sectionIndex)
+			coordinator.expand(sectionNode)
 			self.applyChanges(animate: true) {
 				completion?()
 			}
@@ -502,12 +502,12 @@ class MasterFeedViewController: UITableViewController, UndoableCommandRunner {
 		}
 	
 		// It wasn't already visable, so expand its folder and try again
-		guard let parent = node.parent, let indexPath = dataSource.indexPath(for: parent) else {
+		guard let parent = node.parent else {
 			completion?()
 			return
 		}
 		
-		coordinator.expandFolder(indexPath)
+		coordinator.expand(parent)
 		reloadNode(parent)
 
 		self.applyChanges(animate: true, adjustScroll: true) { [weak self] in
@@ -664,18 +664,18 @@ private extension MasterFeedViewController {
 	}
 
 	func expand(_ cell: MasterFeedTableViewCell) {
-		guard let indexPath = tableView.indexPath(for: cell)  else {
+		guard let indexPath = tableView.indexPath(for: cell), let node = dataSource.itemIdentifier(for: indexPath) else {
 			return
 		}
-		coordinator.expandFolder(indexPath)
+		coordinator.expand(node)
 		self.applyChanges(animate: true)
 	}
 
 	func collapse(_ cell: MasterFeedTableViewCell) {
-		guard let indexPath = tableView.indexPath(for: cell) else {
+		guard let indexPath = tableView.indexPath(for: cell), let node = dataSource.itemIdentifier(for: indexPath) else {
 			return
 		}
-		coordinator.collapseFolder(indexPath)
+		coordinator.collapse(node)
 		self.applyChanges(animate: true)
 	}
 
diff --git a/iOS/SceneCoordinator.swift b/iOS/SceneCoordinator.swift
index 28cca5102..1aac0b2e2 100644
--- a/iOS/SceneCoordinator.swift
+++ b/iOS/SceneCoordinator.swift
@@ -437,103 +437,45 @@ class SceneCoordinator: NSObject, UndoableCommandRunner, UnreadCountProvider {
 		return 0
 	}
 		
-	func expandSection(_ section: Int) {
-		guard let expandNode = treeController.rootNode.childAtIndex(section), !expandNode.isExpanded else {
-			return
-		}
-
-		expandNode.isExpanded = true
-		
+	func expand(_ node: Node) {
+		node.isExpanded = true
 		animatingChanges = true
-		
-		var i = 0
-		
-		func addNode(_ node: Node) {
-			shadowTable[section].insert(node, at: i)
-			i = i + 1
-		}
-		
-		for child in expandNode.childNodes {
-			addNode(child)
-			if child.isExpanded {
-				for gChild in child.childNodes {
-					addNode(gChild)
-				}
-			}
-		}
-		
+		rebuildShadowTable()
 		animatingChanges = false
 	}
 	
 	func expandAllSectionsAndFolders() {
-		for (sectionIndex, sectionNode) in treeController.rootNode.childNodes.enumerated() {
-
-			expandSection(sectionIndex)
-			
+		for sectionNode in treeController.rootNode.childNodes {
+			sectionNode.isExpanded = true
 			for topLevelNode in sectionNode.childNodes {
-				if topLevelNode.representedObject is Folder, let indexPath = indexPathFor(topLevelNode) {
-					expandFolder(indexPath)
+				if topLevelNode.representedObject is Folder {
+					topLevelNode.isExpanded = true
 				}
 			}
-
 		}
-	}
-	
-	func expandFolder(_ indexPath: IndexPath) {
-		guard let expandNode = nodeFor(indexPath), !expandNode.isExpanded && expandNode.representedObject is Folder else {
-			return
-		}
-		
-		expandNode.isExpanded = true
-		
 		animatingChanges = true
-		
-		for i in 0..<expandNode.childNodes.count {
-			if let child = expandNode.childAtIndex(i) {
-				let nextIndex = indexPath.row + i + 1
-				shadowTable[indexPath.section].insert(child, at: nextIndex)
-			}
-		}
-		
+		rebuildShadowTable()
 		animatingChanges = false
 	}
 	
-	func collapseSection(_ section: Int) {
-		guard let collapseNode = treeController.rootNode.childAtIndex(section), collapseNode.isExpanded else {
-			return
-		}
-		
+	func collapse(_ node: Node) {
+		node.isExpanded = false
 		animatingChanges = true
-		collapseNode.isExpanded = false
-		shadowTable[section] = [Node]()
+		rebuildShadowTable()
 		animatingChanges = false
 	}
 	
 	func collapseAllFolders() {
 		for sectionNode in treeController.rootNode.childNodes {
+			sectionNode.isExpanded = true
 			for topLevelNode in sectionNode.childNodes {
-				if topLevelNode.representedObject is Folder, let indexPath = indexPathFor(topLevelNode) {
-					collapseFolder(indexPath)
+				if topLevelNode.representedObject is Folder {
+					topLevelNode.isExpanded = true
 				}
 			}
 		}
-	}
-	
-	func collapseFolder(_ indexPath: IndexPath) {
-		guard let collapseNode = nodeFor(indexPath), collapseNode.isExpanded && collapseNode.representedObject is Folder else {
-			return
-		}
-		
 		animatingChanges = true
-
-		collapseNode.isExpanded = false
-		
-		for child in collapseNode.childNodes {
-			if let index = shadowTable[indexPath.section].firstIndex(of: child) {
-				shadowTable[indexPath.section].remove(at: index)
-			}
-		}
-		
+		rebuildShadowTable()
 		animatingChanges = false
 	}