Implement CellIdentifierProtocol for Select Location

2 files changed, 9 insertions, 31 deletions

diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index d237ef1d15..857037e362 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -28,23 +28,24 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
         customLists: LocationDataSourceProtocol
     ) {
         self.tableView = tableView
-
         #if DEBUG
         self.dataSources.append(customLists)
         #endif
         self.dataSources.append(allLocations)
 
         super.init(tableView: tableView) { _, indexPath, itemIdentifier in
-            let reuseIdentifier = LocationSection.Cell.locationCell.reuseIdentifier
-            // swiftlint:disable:next force_cast
-            let cell = tableView.dequeueReusableCell(withIdentifier: reuseIdentifier, for: indexPath) as! LocationCell
+            let cell = tableView.dequeueReusableView(
+                withIdentifier: LocationSection.allCases[indexPath.section],
+                for: indexPath
+                // swiftlint:disable:next force_cast
+            ) as! LocationCell
             cell.configureCell(item: itemIdentifier)
             return cell
         }
 
         tableView.delegate = self
+        tableView.registerReusableViews(from: LocationSection.self)
         defaultRowAnimation = .fade
-        registerClasses()
     }
 
     func setRelays(_ response: REST.ServerRelaysResponse, selectedRelays: UserSelectedRelays?, filter: RelayFilter) {
@@ -137,12 +138,6 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
         }
     }
 
-    private func registerClasses() {
-        LocationSection.allCases.forEach {
-            tableView.register($0.cell.reusableViewClass, forCellReuseIdentifier: $0.cell.reuseIdentifier)
-        }
-    }
-
     private func mapSelectedItem(from selectedRelays: UserSelectedRelays?) {
         let allLocationsDataSource =
             dataSources.first(where: { $0 is AllLocationDataSource }) as? AllLocationDataSource
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationSection.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationSection.swift
index 6ebf676adb..51c4d00305 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationSection.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationSection.swift
@@ -7,7 +7,7 @@
 //
 
 import Foundation
-enum LocationSection: Int, Hashable, CustomStringConvertible, CaseIterable {
+enum LocationSection: String, Hashable, CustomStringConvertible, CaseIterable, CellIdentifierProtocol {
     case customLists
     case allLocations
 
@@ -28,8 +28,8 @@ enum LocationSection: Int, Hashable, CustomStringConvertible, CaseIterable {
         }
     }
 
-    var cell: Cell {
-        .locationCell
+    var cellClass: AnyClass {
+        LocationCell.self
     }
 
     static var allCases: [LocationSection] {
@@ -40,20 +40,3 @@ enum LocationSection: Int, Hashable, CustomStringConvertible, CaseIterable {
         #endif
     }
 }
-
-extension LocationSection {
-    enum Cell: String, CaseIterable {
-        case locationCell
-
-        var reusableViewClass: AnyClass {
-            switch self {
-            case .locationCell:
-                return LocationCell.self
-            }
-        }
-
-        var reuseIdentifier: String {
-            self.rawValue
-        }
-    }
-}