Add LocationDataSource

1 files changed, 129 insertions, 386 deletions

diff --git a/ios/MullvadVPN/SelectLocationViewController.swift b/ios/MullvadVPN/SelectLocationViewController.swift
index 47a6dab663..39464bce5c 100644
--- a/ios/MullvadVPN/SelectLocationViewController.swift
+++ b/ios/MullvadVPN/SelectLocationViewController.swift
@@ -6,35 +6,49 @@
 //  Copyright © 2019 Mullvad VPN AB. All rights reserved.
 //
 
-import DiffableDataSources
 import UIKit
 import Logging
 
-private let kCellIdentifier = "Cell"
+class SelectLocationViewController: UIViewController, RelayCacheObserver, UITableViewDelegate {
 
-class SelectLocationViewController: UITableViewController, RelayCacheObserver {
+    private enum ReuseIdentifiers: String {
+        case cell
+        case header
+    }
 
-    private enum Error: ChainedError {
-        case loadRelayList(RelayCacheError)
-        case getRelayConstraints(TunnelManager.Error)
+    private lazy var tableView: UITableView = {
+        let tableView = UITableView(frame: view.bounds, style: .plain)
+        tableView.translatesAutoresizingMaskIntoConstraints = true
+        tableView.autoresizingMask = [.flexibleWidth, .flexibleHeight]
+        tableView.backgroundColor = .clear
+        tableView.separatorColor = .secondaryColor
+        tableView.separatorInset = .zero
+        tableView.estimatedRowHeight = 53
+        tableView.estimatedSectionHeaderHeight = 109
+        tableView.indicatorStyle = .white
 
-        var errorDescription: String? {
-            switch self {
-            case .loadRelayList:
-                return "Failure to load a relay list"
-            case .getRelayConstraints:
-                return "Failure to get relay constraints"
-            }
-        }
-    }
+        tableView.register(SelectLocationHeaderView.self, forHeaderFooterViewReuseIdentifier: ReuseIdentifiers.header.rawValue)
+        tableView.register(SelectLocationCell.self, forCellReuseIdentifier: ReuseIdentifiers.cell.rawValue)
+
+        return tableView
+    }()
 
     private let logger = Logger(label: "SelectLocationController")
-    private var cachedRelays: CachedRelays?
-    private var relayConstraints: RelayConstraints?
-    private var expandedItems = [RelayLocation]()
-    private var dataSource: DataSource?
+    private var dataSource: LocationDataSource?
+    private var setCachedRelaysOnViewDidLoad: CachedRelays?
+    private var setRelayLocationOnViewDidLoad: RelayLocation?
+    private var isViewAppeared = false
 
-    var didSelectLocationHandler: ((RelayLocation) -> Void)?
+    var didSelectRelayLocation: ((SelectLocationViewController, RelayLocation) -> Void)?
+    var scrollToSelectedRelayOnViewWillAppear = true
+
+    init() {
+        super.init(nibName: nil, bundle: nil)
+    }
+
+    required init?(coder: NSCoder) {
+        fatalError("init(coder:) has not been implemented")
+    }
 
     // MARK: - View lifecycle
 
@@ -42,440 +56,169 @@ class SelectLocationViewController: UITableViewController, RelayCacheObserver {
         super.viewDidLoad()
 
         view.backgroundColor = .secondaryColor
-        tableView.tableHeaderView = SelectLocationHeaderView(frame: CGRect(x: 0, y: 0, width: 50, height: 50))
-        tableView.register(SelectLocationCell.self, forCellReuseIdentifier: kCellIdentifier)
-        tableView.separatorColor = .secondaryColor
-        tableView.separatorInset = .zero
+        view.addSubview(tableView)
 
-        dataSource = DataSource(
+        dataSource = LocationDataSource(
             tableView: self.tableView,
             cellProvider: { [weak self] (tableView, indexPath, item) -> UITableViewCell? in
                 guard let self = self else { return nil }
 
                 let cell = tableView.dequeueReusableCell(
-                    withIdentifier: kCellIdentifier, for: indexPath) as! SelectLocationCell
+                    withIdentifier: ReuseIdentifiers.cell.rawValue, for: indexPath) as! SelectLocationCell
 
-                cell.accessibilityIdentifier = item.relayLocation.stringRepresentation
-                cell.isDisabled = !item.hasActiveRelays()
-                cell.locationLabel.text = item.displayName()
-                cell.statusIndicator.isActive = item.hasActiveRelays()
-                cell.showsCollapseControl = item.isCollapsibleLevel()
-                cell.isExpanded = self.expandedItems.contains(item.relayLocation)
+                cell.accessibilityIdentifier = item.location.stringRepresentation
+                cell.isDisabled = !item.isActive
+                cell.locationLabel.text = item.displayName
+                cell.statusIndicator.isActive = item.isActive
+                cell.showsCollapseControl = item.isCollapsible
+                cell.isExpanded = item.showsChildren
                 cell.didCollapseHandler = { [weak self] (cell) in
                     self?.collapseCell(cell)
                 }
 
                 return cell
-        })
+            })
 
-        dataSource?.defaultRowAnimation = .top
+        tableView.delegate = self
         tableView.dataSource = dataSource
 
-        RelayCache.shared.addObserver(self)
-
-        updateDataSource(animateDifferences: false) {
-            self.updateTableViewSelection(scroll: true, animated: false)
+        if let setCachedRelaysOnViewDidLoad = self.setCachedRelaysOnViewDidLoad {
+            dataSource?.setRelays(setCachedRelaysOnViewDidLoad.relays)
         }
-    }
-
-    override func viewDidLayoutSubviews() {
-        super.viewDidLayoutSubviews()
-
-        updateTableHeaderViewSizeIfNeeded()
-    }
 
-    // MARK: - UITableViewDelegate
-
-    override func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
-        return dataSource?.itemIdentifier(for: indexPath)?.hasActiveRelays() ?? false
-    }
-
-    override func tableView(_ tableView: UITableView, indentationLevelForRowAt indexPath: IndexPath) -> Int {
-        return dataSource?.itemIdentifier(for: indexPath)?.indentationLevel() ?? 0
-    }
-
-    override func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
-        guard let item = dataSource?.itemIdentifier(for: indexPath) else { return }
-
-        // Disable interaction with the controller after selection
-        tableView.isUserInteractionEnabled = false
-
-        DispatchQueue.main.asyncAfter(deadline: .now() + .milliseconds(250)) {
-            self.didSelectLocationHandler?(item.relayLocation)
+        if let setRelayLocationOnViewDidLoad = setRelayLocationOnViewDidLoad {
+            dataSource?.setSelectedRelayLocation(
+                setRelayLocationOnViewDidLoad,
+                showHiddenParents: true,
+                animated: false,
+                scrollPosition: .none
+            )
         }
-    }
 
-    // MARK: - RelayCacheObserver
-
-    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelays cachedRelays: CachedRelays) {
-        self.didReceiveCachedRelays(cachedRelays) { (result) in
-            DispatchQueue.main.async {
-                switch result {
-                case .success(let (cachedRelays, relayConstraints)):
-                    self.didReceiveCachedRelays(cachedRelays, relayConstraints: relayConstraints)
-
-                case .failure(let error):
-                    self.logger.error(chainedError: error)
-                }
-            }
-        }
-    }
-
-    // MARK: - Public
-
-    func prefetchData(completionHandler: @escaping () -> Void) {
-        fetchRelays { (result) in
-            DispatchQueue.main.async {
-                switch result {
-                case .success(let (cachedRelays, relayConstraints)):
-                    self.didReceiveCachedRelays(cachedRelays, relayConstraints: relayConstraints)
-
-                case .failure(let error):
-                    self.logger.error(chainedError: error)
-                }
-
-                completionHandler()
-            }
-        }
-    }
-
-    // MARK: - Relay list handling
-
-    private func fetchRelays(completionHandler: @escaping (Result<(CachedRelays, RelayConstraints), Error>) -> Void) {
-        RelayCache.shared.read { (result) in
-            switch result {
-            case .success(let cachedRelays):
-                self.didReceiveCachedRelays(cachedRelays, completionHandler: completionHandler)
-
-            case .failure(let error):
-                completionHandler(.failure(.loadRelayList(error)))
-            }
-        }
+        RelayCache.shared.addObserver(self)
     }
 
-    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays, completionHandler: @escaping (Result<(CachedRelays, RelayConstraints), Error>) -> Void) {
-        TunnelManager.shared.getRelayConstraints { (result) in
-            let result = result
-                .map { (cachedRelays, $0) }
-                .mapError { Error.getRelayConstraints($0) }
+    override func viewWillAppear(_ animated: Bool) {
+        super.viewWillAppear(animated)
 
-            completionHandler(result)
+        if let indexPath = dataSource?.indexPathForSelectedRelay(), scrollToSelectedRelayOnViewWillAppear, !isViewAppeared {
+            self.tableView.scrollToRow(at: indexPath, at: .middle, animated: false)
         }
     }
 
-    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays, relayConstraints: RelayConstraints) {
-        self.cachedRelays = cachedRelays
-        self.relayConstraints = relayConstraints
-
-        let relayLocation = relayConstraints.location.value
-        expandedItems = relayLocation?.ascendants ?? []
+    override func viewDidAppear(_ animated: Bool) {
+        super.viewDidAppear(animated)
 
-        updateDataSource(animateDifferences: false)
-        tableView.reloadData()
+        isViewAppeared = true
 
-        updateTableViewSelection(scroll: true, animated: false)
+        tableView.flashScrollIndicators()
     }
 
-    private func computeIndexPathForSelectedLocation(relayLocation: RelayLocation) -> IndexPath? {
-        guard let row = dataSource?.snapshot()
-            .itemIdentifiers
-            .firstIndex(where: { $0.relayLocation == relayLocation }) else {
-                return nil
-        }
+    override func viewDidDisappear(_ animated: Bool) {
+        super.viewDidDisappear(animated)
 
-        return IndexPath(row: row, section: 0)
+        isViewAppeared = false
     }
 
-    // MARK: - Collapsible cells
-
-    private func updateTableViewSelection(scroll: Bool, animated: Bool) {
-        guard let relayLocation = relayConstraints?.location.value else { return }
-
-        let indexPath = computeIndexPathForSelectedLocation(relayLocation: relayLocation)
+    // MARK: - UITableViewDelegate
 
-        let scrollPosition: UITableView.ScrollPosition = scroll ? .middle : .none
-        tableView.selectRow(at: indexPath, animated: animated, scrollPosition: scrollPosition)
+    func tableView(_ tableView: UITableView, shouldHighlightRowAt indexPath: IndexPath) -> Bool {
+        return dataSource?.item(for: indexPath)?.isActive ?? false
     }
 
-    private func updateDataSource(animateDifferences: Bool, completion: (() -> Void)? = nil) {
-        let items = self.cachedRelays.map { (cachedRelays) -> [DataSourceItem] in
-            return cachedRelays.relays.makeDataSource { (item) -> Bool in
-                return expandedItems.contains(item.relayLocation)
-            }
-        } ?? []
-
-        var snapshot = DataSourceSnapshot()
-        snapshot.appendSections([.locations])
-        snapshot.appendItems(items, toSection: .locations)
-
-        dataSource?.apply(
-            snapshot,
-            animatingDifferences: animateDifferences,
-            completion: completion
-        )
+    func tableView(_ tableView: UITableView, indentationLevelForRowAt indexPath: IndexPath) -> Int {
+        return dataSource?.item(for: indexPath)?.indentationLevel ?? 0
     }
 
-    private func collapseCell(_ cell: SelectLocationCell) {
-        guard let cellIndexPath = tableView.indexPath(for: cell),
-            let item = dataSource?.itemIdentifier(for: cellIndexPath) else {
-                return
-        }
-
-        let itemLocation = item.relayLocation
-
-        if let index = expandedItems.firstIndex(of: itemLocation) {
-            expandedItems.remove(at: index)
-            cell.isExpanded = false
-        } else {
-            expandedItems.append(itemLocation)
-            cell.isExpanded = true
-        }
-
-        updateDataSource(animateDifferences: true) {
-            self.updateTableViewSelection(scroll: false, animated: true)
+    func tableView(_ tableView: UITableView, willDisplay cell: UITableViewCell, forRowAt indexPath: IndexPath) {
+        if let item = dataSource?.item(for: indexPath), item.location == dataSource?.selectedRelayLocation {
+            cell.setSelected(true, animated: false)
         }
     }
 
-    // MARK: - UITableView header
-
-    private func updateTableHeaderViewSizeIfNeeded() {
-        guard let header = tableView.tableHeaderView else { return }
+    func tableView(_ tableView: UITableView, didSelectRowAt indexPath: IndexPath) {
+        guard let item = dataSource?.item(for: indexPath) else { return }
 
-        // measure the view size
-        let sizeConstraint = CGSize(
-            width: tableView.bounds.width,
-            height: UIView.layoutFittingCompressedSize.height
+        dataSource?.setSelectedRelayLocation(
+            item.location,
+            showHiddenParents: false,
+            animated: false,
+            scrollPosition: .none
         )
-
-        let newSize = header.systemLayoutSizeFitting(sizeConstraint)
-        let oldSize = header.frame.size
-
-        if oldSize.height != newSize.height {
-            header.frame.size.height = newSize.height
-
-            // reset the header view to force UITableView layout pass
-            tableView.tableHeaderView = header
-        }
+        didSelectRelayLocation?(self, item.location)
     }
-}
-
-private extension RelayLocation {
 
-    /// A list of `RelayLocation` items preceding the given one in the relay tree
-    var ascendants: [RelayLocation] {
-        switch self {
-        case .hostname(let country, let city, _):
-            return [.country(country), .city(country, city)]
+    func tableView(_ tableView: UITableView, viewForHeaderInSection section: Int) -> UIView? {
+        assert(section == 0)
 
-        case .city(let country, _):
-            return [.country(country)]
+        let view = tableView.dequeueReusableHeaderFooterView(withIdentifier: ReuseIdentifiers.header.rawValue) as! SelectLocationHeaderView
 
-        case .country:
-            return []
+        // When contained within the navigation controller, we want the distance between the navigation title
+        // and the table header label to be exactly 24pt.
+        if let navigationBar = navigationController?.navigationBar as? CustomNavigationBar {
+            view.topLayoutMarginAdjustmentForNavigationBarTitle = navigationBar.titleLabelBottomInset
         }
-    }
 
-}
-
-/// Enum describing the table view sections
-private enum DataSourceSection {
-    case locations
-}
-
-/// Data source type
-private typealias DataSource = TableViewDiffableDataSource<DataSourceSection, DataSourceItem>
-
-/// Data source snapshot type
-private typealias DataSourceSnapshot = DiffableDataSourceSnapshot<DataSourceSection, DataSourceItem>
-
-/// A wrapper type for RelayList to be able to represent it as a flat list
-private enum DataSourceItem: Hashable {
-
-    struct Country {
-        let location: String
-        let name: String
-        let hasActiveRelays: Bool
-    }
-
-    struct City {
-        let location: String
-        let name: String
-        let hasActiveRelays: Bool
+        return view
     }
 
-    struct Hostname {
-        let location: String
-        let hostname: String
-        let active: Bool
-    }
-
-    case country(Country)
-    case city(City)
-    case hostname(Hostname)
+    // MARK: - RelayCacheObserver
 
-    var relayLocation: RelayLocation {
-        switch self {
-        case .country(let country):
-            return .country(country.location)
-        case .city(let city):
-            let split = city.location.split(separator: "-", maxSplits: 2).map(String.init)
-            return .city(split[0], split[1])
-        case .hostname(let host):
-            let split = host.location.split(separator: "-", maxSplits: 2).map(String.init)
-            return .hostname(split[0], split[1], host.hostname)
+    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelays cachedRelays: CachedRelays) {
+        DispatchQueue.main.async {
+            self.didReceiveCachedRelays(cachedRelays)
         }
     }
 
-    static func == (lhs: DataSourceItem, rhs: DataSourceItem) -> Bool {
-        lhs.relayLocation == rhs.relayLocation
-    }
-
-    func hash(into hasher: inout Hasher) {
-        hasher.combine(relayLocation)
-    }
-
-    func indentationLevel() -> Int {
-        switch self {
-        case .country:
-            return 0
-        case .city:
-            return 1
-        case .hostname:
-            return 2
-        }
-    }
+    // MARK: - Public
 
-    func displayName() -> String {
-        switch self {
-        case .country(let country):
-            return country.name
-        case .city(let city):
-            return city.name
-        case .hostname(let relay):
-            return relay.hostname
-        }
-    }
+    func prefetchData(completionHandler: @escaping (RelayCacheError?) -> Void) {
+        RelayCache.shared.read { (result) in
+            DispatchQueue.main.async {
+                switch result {
+                case .success(let cachedRelays):
+                    self.didReceiveCachedRelays(cachedRelays)
+                    completionHandler(nil)
 
-    func hasActiveRelays() -> Bool {
-        switch self {
-        case .country(let country):
-            return country.hasActiveRelays
-        case .city(let city):
-            return city.hasActiveRelays
-        case .hostname(let host):
-            return host.active
+                case .failure(let error):
+                    completionHandler(error)
+                }
+            }
         }
     }
 
-    func isCollapsibleLevel() -> Bool {
-        switch self {
-        case .country, .city:
-            return self.hasActiveRelays()
-        case .hostname:
-            return false
+    func setSelectedRelayLocation(_ relayLocation: RelayLocation?, animated: Bool, scrollPosition: UITableView.ScrollPosition) {
+        guard isViewLoaded else {
+            self.setRelayLocationOnViewDidLoad = relayLocation
+            return
         }
-    }
-
-}
-
-extension ServerRelaysResponse {
-    fileprivate static func lexicalSortComparator(_ a: String, _ b: String) -> Bool {
-        return a.localizedCaseInsensitiveCompare(b) == .orderedAscending
-    }
 
-    fileprivate static func fileSortComparator(_ a: String, _ b: String) -> Bool {
-        return a.localizedStandardCompare(b) == .orderedAscending
+        self.dataSource?.setSelectedRelayLocation(
+            relayLocation,
+            showHiddenParents: true,
+            animated: animated,
+            scrollPosition: scrollPosition
+        )
     }
 
-    fileprivate func makeDataSource(evaluator: (DataSourceItem) -> Bool) -> [DataSourceItem] {
-        let relaysByCountry = Dictionary(grouping: wireguard.relays) { (relay) -> String in
-            return relay.location.split(separator: "-").first.flatMap(String.init)!
-        }
-
-        var items = [DataSourceItem]()
-
-        var countryItems = [DataSourceItem.Country]()
-        var cityItems = [String: [DataSourceItem.City]]()
-        var relayItems = [String: [DataSourceItem.Hostname]]()
-
-        for (countryCode, relays) in relaysByCountry {
-            let relaysByCity = Dictionary(grouping: relays) { (relay) -> String in
-                return relay.location
-            }
-
-            if let (cityCode, relays) = relaysByCity.first {
-                guard let location = locations[cityCode] else {
-                    continue
-                }
-
-                let country = DataSourceItem.Country(
-                    location: countryCode,
-                    name: location.country,
-                    hasActiveRelays: relays.contains(where: { (serverRelay) -> Bool in
-                        return serverRelay.active
-                    }))
-
-                countryItems.append(country)
-                if !evaluator(.country(country)) {
-                    continue
-                }
-            }
-
-            for (cityCode, relays) in relaysByCity {
-                guard let location = locations[cityCode] else {
-                    // TODO: log to file?
-                    print("Location not found: \(cityCode)")
-                    continue
-                }
-
-                let city = DataSourceItem.City(
-                    location: cityCode,
-                    name: location.city,
-                    hasActiveRelays: relays.contains(where: { (serverRelay) -> Bool in
-                        return serverRelay.active
-                    }))
-
-                if var cities = cityItems[countryCode] {
-                    cities.append(city)
-                    cityItems[countryCode] = cities
-                } else {
-                    cityItems[countryCode] = [city]
-                }
-
-                if !evaluator(.city(city)) {
-                    continue
-                }
-
-                relayItems[cityCode] = relays.map { (relay) -> DataSourceItem.Hostname in
-                    return DataSourceItem.Hostname(location: relay.location, hostname: relay.hostname, active: relay.active)
-                }
-            }
-        }
+    // MARK: - Relay list handling
 
-        countryItems.sort { (a, b) -> Bool in
-            return Self.lexicalSortComparator(a.name, b.name)
+    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays) {
+        guard isViewLoaded else {
+            self.setCachedRelaysOnViewDidLoad = cachedRelays
+            return
         }
+        self.dataSource?.setRelays(cachedRelays.relays)
+    }
 
-        for country in countryItems {
-            items.append(.country(country))
-
-            if var cities = cityItems[country.location] {
-                cities.sort { (a, b) -> Bool in
-                    return Self.lexicalSortComparator(a.name, b.name)
-                }
-                for city in cities {
-                    items.append(.city(city))
+    // MARK: - Collapsible cells
 
-                    if var relays = relayItems[city.location] {
-                        relays.sort { (a, b) -> Bool in
-                            return Self.fileSortComparator(a.hostname, b.hostname)
-                        }
-                        items.append(contentsOf: relays.map { DataSourceItem.hostname($0) })
-                    }
-                }
-            }
+    private func collapseCell(_ cell: SelectLocationCell) {
+        guard let cellIndexPath = tableView.indexPath(for: cell),
+              let dataSource = dataSource, let location = dataSource.relayLocation(for: cellIndexPath) else {
+            return
         }
 
-        return items
+        dataSource.toggleChildren(location, animated: true)
     }
 }