Combine AddLocationCell and LocationCell to one

8 files changed, 181 insertions, 353 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index e56b72cf8e..632600e08f 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -829,8 +829,6 @@
 		F02F41A02B9723AF00625A4F /* AddLocationsViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = F02F419A2B9723AE00625A4F /* AddLocationsViewController.swift */; };
 		F02F41A12B9723AF00625A4F /* AddLocationsDataSource.swift in Sources */ = {isa = PBXBuildFile; fileRef = F02F419B2B9723AE00625A4F /* AddLocationsDataSource.swift */; };
 		F02F41A22B9723AF00625A4F /* AddLocationsCoordinator.swift in Sources */ = {isa = PBXBuildFile; fileRef = F02F419C2B9723AF00625A4F /* AddLocationsCoordinator.swift */; };
-		F02F41A42B9723AF00625A4F /* AddLocationCellViewModel.swift in Sources */ = {isa = PBXBuildFile; fileRef = F02F419E2B9723AF00625A4F /* AddLocationCellViewModel.swift */; };
-		F02F41A52B9723AF00625A4F /* AddLocationCell.swift in Sources */ = {isa = PBXBuildFile; fileRef = F02F419F2B9723AF00625A4F /* AddLocationCell.swift */; };
 		F03580252A13842C00E5DAFD /* IncreasedHitButton.swift in Sources */ = {isa = PBXBuildFile; fileRef = F03580242A13842C00E5DAFD /* IncreasedHitButton.swift */; };
 		F04413612BA45CD70018A6EE /* CustomListLocationNodeBuilder.swift in Sources */ = {isa = PBXBuildFile; fileRef = F04413602BA45CD70018A6EE /* CustomListLocationNodeBuilder.swift */; };
 		F04413622BA45CE30018A6EE /* CustomListLocationNodeBuilder.swift in Sources */ = {isa = PBXBuildFile; fileRef = F04413602BA45CD70018A6EE /* CustomListLocationNodeBuilder.swift */; };
@@ -1991,8 +1989,6 @@
 		F02F419A2B9723AE00625A4F /* AddLocationsViewController.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AddLocationsViewController.swift; sourceTree = "<group>"; };
 		F02F419B2B9723AE00625A4F /* AddLocationsDataSource.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AddLocationsDataSource.swift; sourceTree = "<group>"; };
 		F02F419C2B9723AF00625A4F /* AddLocationsCoordinator.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AddLocationsCoordinator.swift; sourceTree = "<group>"; };
-		F02F419E2B9723AF00625A4F /* AddLocationCellViewModel.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AddLocationCellViewModel.swift; sourceTree = "<group>"; };
-		F02F419F2B9723AF00625A4F /* AddLocationCell.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AddLocationCell.swift; sourceTree = "<group>"; };
 		F03580242A13842C00E5DAFD /* IncreasedHitButton.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IncreasedHitButton.swift; sourceTree = "<group>"; };
 		F04413602BA45CD70018A6EE /* CustomListLocationNodeBuilder.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = CustomListLocationNodeBuilder.swift; sourceTree = "<group>"; };
 		F04F95A02B21D24400431E08 /* shadowsocks.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = shadowsocks.h; sourceTree = "<group>"; };
@@ -3520,8 +3516,6 @@
 			isa = PBXGroup;
 			children = (
 				7A6389D72B7E3BD6008E77E1 /* AddCustomListCoordinator.swift */,
-				F02F419F2B9723AF00625A4F /* AddLocationCell.swift */,
-				F02F419E2B9723AF00625A4F /* AddLocationCellViewModel.swift */,
 				F02F419C2B9723AF00625A4F /* AddLocationsCoordinator.swift */,
 				F02F419B2B9723AE00625A4F /* AddLocationsDataSource.swift */,
 				F02F419A2B9723AE00625A4F /* AddLocationsViewController.swift */,
@@ -5197,7 +5191,6 @@
 				582BB1B1229569620055B6EF /* UINavigationBar+Appearance.swift in Sources */,
 				7A9FA1442A2E3FE5000B728D /* CheckableSettingsCell.swift in Sources */,
 				58ACF6492655365700ACE4B7 /* VPNSettingsViewController.swift in Sources */,
-				F02F41A52B9723AF00625A4F /* AddLocationCell.swift in Sources */,
 				7ABE318D2A1CDD4500DF4963 /* UIFont+Weight.swift in Sources */,
 				58C774BE29A7A249003A1A56 /* CustomNavigationController.swift in Sources */,
 				E1FD0DF528AA7CE400299DB4 /* StatusActivityView.swift in Sources */,
@@ -5454,7 +5447,6 @@
 				7A6389DB2B7E3BD6008E77E1 /* CustomListCellConfiguration.swift in Sources */,
 				F09A297C2A9F8A9B00EA3B6F /* VoucherTextField.swift in Sources */,
 				7A5869B72B56B41500640D27 /* IPOverrideTextViewController.swift in Sources */,
-				F02F41A42B9723AF00625A4F /* AddLocationCellViewModel.swift in Sources */,
 				58ACF64B26553C3F00ACE4B7 /* SettingsSwitchCell.swift in Sources */,
 				7AF9BE952A40461100DBFEDB /* RelayFilterView.swift in Sources */,
 				7A09C98129D99215000C2CAC /* String+FuzzyMatch.swift in Sources */,
diff --git a/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
new file mode 100644
index 0000000000..eaa4f1187c
--- /dev/null
+++ b/ios/MullvadVPN.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
@@ -0,0 +1,23 @@
+{
+  "originHash" : "c15149b2d59d9e9c72375f65339c04f41a19943e1117e682df27fc9f943fdc56",
+  "pins" : [
+    {
+      "identity" : "swift-log",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-log.git",
+      "state" : {
+        "revision" : "173f567a2dfec11d74588eea82cecea555bdc0bc",
+        "version" : "1.4.0"
+      }
+    },
+    {
+      "identity" : "wireguard-apple",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/mullvad/wireguard-apple.git",
+      "state" : {
+        "revision" : "11a00c20dc03f2751db47e94f585c0778c7bde82"
+      }
+    }
+  ],
+  "version" : 3
+}
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationCell.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationCell.swift
deleted file mode 100644
index e0094f0a9a..0000000000
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationCell.swift
+++ /dev/null
@@ -1,204 +0,0 @@
-//
-//  AddLocationCell.swift
-//  MullvadVPN
-//
-//  Created by Mojgan on 2024-02-29.
-//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
-//
-
-import UIKit
-
-protocol AddLocationCellDelegate: AnyObject {
-    func toggleExpanding(cell: AddLocationCell)
-    func toggleSelection(cell: AddLocationCell)
-}
-
-class AddLocationCell: UITableViewCell {
-    weak var delegate: AddLocationCellDelegate?
-
-    private let chevronDown = UIImage(resource: .iconChevronDown)
-    private let chevronUp = UIImage(resource: .iconChevronUp)
-
-    private let locationLabel: UILabel = {
-        let label = UILabel()
-        label.font = UIFont.systemFont(ofSize: 17)
-        label.textColor = .white
-        label.numberOfLines = .zero
-        label.lineBreakStrategy = []
-        return label
-    }()
-
-    private let checkboxButton: UIButton = {
-        let button = UIButton()
-        button.setImage(UIImage(systemName: "checkmark.square.fill"), for: .selected)
-        button.setImage(UIImage(systemName: "square"), for: .normal)
-        button.tintColor = .white
-        return button
-    }()
-
-    private let collapseButton: UIButton = {
-        let button = UIButton(type: .custom)
-        button.accessibilityIdentifier = .collapseButton
-        button.isAccessibilityElement = false
-        button.tintColor = .white
-        return button
-    }()
-
-    var isExpanded = false {
-        didSet {
-            updateCollapseImage()
-            updateAccessibilityCustomActions()
-        }
-    }
-
-    var showsCollapseControl = false {
-        didSet {
-            collapseButton.isHidden = !showsCollapseControl
-            updateAccessibilityCustomActions()
-        }
-    }
-
-    override var indentationLevel: Int {
-        didSet {
-            updateBackgroundColor()
-            setLayoutMargins()
-        }
-    }
-
-    override init(style: UITableViewCell.CellStyle, reuseIdentifier: String?) {
-        super.init(style: style, reuseIdentifier: reuseIdentifier)
-        setupCell()
-    }
-
-    required init?(coder: NSCoder) {
-        fatalError("init(coder:) has not been implemented")
-    }
-
-    private func setLayoutMargins() {
-        let indentation = CGFloat(indentationLevel) * indentationWidth
-
-        var contentMargins = UIMetrics.locationCellLayoutMargins
-        contentMargins.leading += indentation
-
-        contentView.directionalLayoutMargins = contentMargins
-    }
-
-    private func setupCell() {
-        indentationWidth = UIMetrics.TableView.cellIndentationWidth
-
-        backgroundColor = .clear
-        contentView.backgroundColor = .clear
-
-        backgroundView = UIView()
-
-        collapseButton.addTarget(self, action: #selector(handleCollapseButton(_:)), for: .touchUpInside)
-        checkboxButton.addTarget(self, action: #selector(toggleCheckboxButton(_:)), for: .touchUpInside)
-
-        contentView.addConstrainedSubviews([checkboxButton, locationLabel, collapseButton]) {
-            checkboxButton.pinEdgeToSuperviewMargin(.leading(.zero))
-            checkboxButton.centerYAnchor.constraint(equalTo: contentView.centerYAnchor)
-            checkboxButton.widthAnchor
-                .constraint(
-                    equalToConstant: 44.0
-                )
-            checkboxButton.heightAnchor.constraint(equalTo: checkboxButton.widthAnchor, multiplier: 1, constant: 0)
-
-            locationLabel.leadingAnchor.constraint(
-                equalTo: checkboxButton.trailingAnchor,
-                constant: 12
-            )
-
-            locationLabel.trailingAnchor.constraint(lessThanOrEqualTo: collapseButton.leadingAnchor)
-                .withPriority(.defaultHigh)
-            locationLabel.pinEdgesToSuperviewMargins(PinnableEdges([.top(.zero), .bottom(.zero)]))
-
-            collapseButton.widthAnchor
-                .constraint(
-                    equalToConstant: UIMetrics.contentLayoutMargins.leading + UIMetrics
-                        .contentLayoutMargins.trailing + 24.0
-                )
-            collapseButton.pinEdgesToSuperviewMargins(.all().excluding(.leading))
-        }
-
-        updateCollapseImage()
-        updateAccessibilityCustomActions()
-        updateBackgroundColor()
-        setLayoutMargins()
-    }
-
-    private func updateBackgroundColor() {
-        backgroundView?.backgroundColor = backgroundColorForIdentationLevel()
-    }
-
-    private func backgroundColorForIdentationLevel() -> UIColor {
-        switch indentationLevel {
-        case 1:
-            return UIColor.Cell.Background.indentationLevelOne
-        case 2:
-            return UIColor.Cell.Background.indentationLevelTwo
-        case 3:
-            return UIColor.Cell.Background.indentationLevelThree
-        default:
-            return UIColor.Cell.Background.normal
-        }
-    }
-
-    private func updateCollapseImage() {
-        let image = isExpanded ? chevronUp : chevronDown
-        collapseButton.setImage(image, for: .normal)
-    }
-
-    private func updateAccessibilityCustomActions() {
-        if showsCollapseControl {
-            let actionName = isExpanded
-                ? NSLocalizedString(
-                    "ADD_LOCATIONS_COLLAPSE_ACCESSIBILITY_ACTION",
-                    tableName: "AddLocationsLocation",
-                    value: "Collapse location",
-                    comment: ""
-                )
-                : NSLocalizedString(
-                    "ADD_LOCATIONS_EXPAND_ACCESSIBILITY_ACTION",
-                    tableName: "AddLocationsLocation",
-                    value: "Expand location",
-                    comment: ""
-                )
-
-            accessibilityCustomActions = [
-                UIAccessibilityCustomAction(
-                    name: actionName,
-                    target: self,
-                    selector: #selector(toggleCollapseAccessibilityAction)
-                ),
-            ]
-        } else {
-            accessibilityCustomActions = nil
-        }
-    }
-
-    // MARK: - actions
-
-    @objc private func handleCollapseButton(_ sender: UIControl) {
-        delegate?.toggleExpanding(cell: self)
-    }
-
-    @objc private func toggleCollapseAccessibilityAction() -> Bool {
-        delegate?.toggleExpanding(cell: self)
-        return true
-    }
-
-    @objc private func toggleCheckboxButton(_ sender: UIControl) {
-        delegate?.toggleSelection(cell: self)
-    }
-}
-
-extension AddLocationCell {
-    func configure(item: AddLocationCellViewModel) {
-        accessibilityIdentifier = item.node.name
-        locationLabel.text = item.node.name
-        showsCollapseControl = !item.node.children.isEmpty
-        isExpanded = item.node.showsChildren
-        checkboxButton.isSelected = item.isSelected
-        checkboxButton.tintColor = item.isSelected ? .successColor : .white
-    }
-}
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationCellViewModel.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationCellViewModel.swift
deleted file mode 100644
index ffc06aea74..0000000000
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationCellViewModel.swift
+++ /dev/null
@@ -1,15 +0,0 @@
-//
-//  AddLocationCellViewModel.swift
-//  MullvadVPN
-//
-//  Created by Mojgan on 2024-03-04.
-//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
-//
-
-import Foundation
-
-struct AddLocationCellViewModel: Hashable {
-    let node: LocationNode
-    var indentationLevel = 0
-    var isSelected: Bool
-}
diff --git a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
index 3b2125811f..18564b8a12 100644
--- a/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
+++ b/ios/MullvadVPN/Coordinators/CustomLists/AddLocationsDataSource.swift
@@ -11,17 +11,7 @@ import MullvadSettings
 import MullvadTypes
 import UIKit
 
-enum AddLocationsSectionIdentifier: String, Hashable, CaseIterable, CellIdentifierProtocol {
-    case `default`
-
-    var cellClass: AnyClass {
-        switch self {
-        case .default: AddLocationCell.self
-        }
-    }
-}
-
-class AddLocationsDataSource: UITableViewDiffableDataSource<AddLocationsSectionIdentifier, AddLocationCellViewModel> {
+class AddLocationsDataSource: UITableViewDiffableDataSource<LocationSection, LocationCellViewModel> {
     private let tableView: UITableView
     private let nodes: [LocationNode]
     private var customListLocationNode: CustomListLocationNode
@@ -42,25 +32,26 @@ class AddLocationsDataSource: UITableViewDiffableDataSource<AddLocationsSectionI
 
         super.init(tableView: tableView) { _, indexPath, itemIdentifier in
             let cell = tableView.dequeueReusableView(
-                withIdentifier: AddLocationsSectionIdentifier.allCases[indexPath.section],
+                withIdentifier: LocationSection.allCases[indexPath.section],
                 for: indexPath
                 // swiftlint:disable:next force_cast
-            ) as! AddLocationCell
-            cell.configure(item: itemIdentifier)
+            ) as! LocationCell
+            cell.configure(item: itemIdentifier, behavior: .add)
             cell.selectionStyle = .none
             return cell
         }
 
         tableView.delegate = self
-        tableView.registerReusableViews(from: AddLocationsSectionIdentifier.self)
+        tableView.registerReusableViews(from: LocationSection.self)
         defaultRowAnimation = .fade
         reloadWithSelectedLocations()
     }
 
     private func reloadWithSelectedLocations() {
-        var locationsList: [AddLocationCellViewModel] = []
+        var locationsList: [LocationCellViewModel] = []
         nodes.forEach { node in
-            let viewModel = AddLocationCellViewModel(
+            let viewModel = LocationCellViewModel(
+                section: .customLists,
                 node: node,
                 isSelected: customListLocationNode.children.contains(node)
             )
@@ -78,40 +69,36 @@ class AddLocationsDataSource: UITableViewDiffableDataSource<AddLocationsSectionI
 
             locationsList.append(contentsOf: recursivelyCreateCellViewModelTree(
                 for: node,
-                in: .default,
+                in: .customLists,
                 indentationLevel: 1
             ))
         }
-        updateDataSnapshot(with: [locationsList])
+        updateDataSnapshot(with: locationsList)
     }
 
     private func updateDataSnapshot(
-        with list: [[AddLocationCellViewModel]],
+        with list: [LocationCellViewModel],
         animated: Bool = false,
         completion: (() -> Void)? = nil
     ) {
-        var snapshot = NSDiffableDataSourceSnapshot<AddLocationsSectionIdentifier, AddLocationCellViewModel>()
-        let sections = AddLocationsSectionIdentifier.allCases
+        var snapshot = NSDiffableDataSourceSnapshot<LocationSection, LocationCellViewModel>()
 
-        snapshot.appendSections(sections)
+        snapshot.appendSections([.customLists])
+        snapshot.appendItems(list, toSection: .customLists)
 
-        for (index, section) in sections.enumerated() {
-            let items = list[index]
-
-            snapshot.appendItems(items, toSection: section)
-        }
         apply(snapshot, animatingDifferences: animated, completion: completion)
     }
 
     private func recursivelyCreateCellViewModelTree(
         for node: LocationNode,
-        in section: AddLocationsSectionIdentifier,
+        in section: LocationSection,
         indentationLevel: Int
-    ) -> [AddLocationCellViewModel] {
-        var viewModels = [AddLocationCellViewModel]()
+    ) -> [LocationCellViewModel] {
+        var viewModels = [LocationCellViewModel]()
         for childNode in node.children {
             viewModels.append(
-                AddLocationCellViewModel(
+                LocationCellViewModel(
+                    section: .customLists,
                     node: childNode,
                     indentationLevel: indentationLevel,
                     isSelected: customListLocationNode.children.contains(childNode)
@@ -145,7 +132,7 @@ class AddLocationsDataSource: UITableViewDiffableDataSource<AddLocationsSectionI
 
     override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
         // swiftlint:disable:next force_cast
-        let cell = super.tableView(tableView, cellForRowAt: indexPath) as! AddLocationCell
+        let cell = super.tableView(tableView, cellForRowAt: indexPath) as! LocationCell
         cell.delegate = self
         return cell
     }
@@ -157,8 +144,8 @@ extension AddLocationsDataSource: UITableViewDelegate {
     }
 }
 
-extension AddLocationsDataSource: AddLocationCellDelegate {
-    func toggleExpanding(cell: AddLocationCell) {
+extension AddLocationsDataSource: LocationCellDelegate {
+    func toggleExpanding(cell: LocationCell) {
         guard let indexPath = tableView.indexPath(for: cell),
               let item = itemIdentifier(for: indexPath) else { return }
         let isExpanded = item.node.showsChildren
@@ -173,12 +160,12 @@ extension AddLocationsDataSource: AddLocationCellDelegate {
             locationList.removeSubNodes(from: item.node)
         }
 
-        updateDataSnapshot(with: [locationList], animated: true, completion: {
+        updateDataSnapshot(with: locationList, animated: true, completion: {
             self.scroll(to: item, animated: true)
         })
     }
 
-    func toggleSelection(cell: AddLocationCell) {
+    func toggleSelection(cell: LocationCell) {
         guard let index = tableView.indexPath(for: cell)?.row else { return }
 
         var locationList = snapshot().itemIdentifiers
@@ -194,14 +181,14 @@ extension AddLocationsDataSource: AddLocationCellDelegate {
         } else {
             customListLocationNode.remove(selectedLocation: item.node, with: locationList)
         }
-        updateDataSnapshot(with: [locationList], completion: {
+        updateDataSnapshot(with: locationList, completion: {
             self.didUpdateCustomList?(self.customListLocationNode)
         })
     }
 }
 
 extension AddLocationsDataSource {
-    private func scroll(to item: AddLocationCellViewModel, animated: Bool) {
+    private func scroll(to item: LocationCellViewModel, animated: Bool) {
         guard
             let visibleIndexPaths = tableView.indexPathsForVisibleRows,
             let indexPath = indexPath(for: item)
@@ -211,9 +198,9 @@ extension AddLocationsDataSource {
             tableView.scrollToRow(at: indexPath, at: .top, animated: animated)
         } else {
             if let last = item.node.children.last {
-                if let lastInsertedIndexPath = self.indexPath(for: AddLocationCellViewModel(
-                    node: last,
-                    isSelected: false
+                if let lastInsertedIndexPath = self.indexPath(for: LocationCellViewModel(
+                    section: .customLists,
+                    node: last
                 )),
                     let lastVisibleIndexPath = visibleIndexPaths.last,
                     lastInsertedIndexPath >= lastVisibleIndexPath {
@@ -224,34 +211,9 @@ extension AddLocationsDataSource {
     }
 }
 
-// MARK: - Toggle expanding
-
-fileprivate extension [AddLocationCellViewModel] {
-    mutating func addSubNodes(from item: AddLocationCellViewModel, at indexPath: IndexPath) {
-        let row = indexPath.row + 1
-        let locations = item.node.children.map {
-            AddLocationCellViewModel(node: $0, indentationLevel: item.indentationLevel + 1, isSelected: item.isSelected)
-        }
-
-        if row < count {
-            insert(contentsOf: locations, at: row)
-        } else {
-            append(contentsOf: locations)
-        }
-    }
-
-    mutating func removeSubNodes(from node: LocationNode) {
-        for node in node.children {
-            node.showsChildren = false
-            removeAll(where: { node == $0.node })
-            removeSubNodes(from: node)
-        }
-    }
-}
-
 // MARK: - Toggle selection in table view
 
-fileprivate extension [AddLocationCellViewModel] {
+fileprivate extension [LocationCellViewModel] {
     mutating func deselectAncestors(from node: LocationNode?) {
         node?.forEachAncestor { parent in
             guard let index = firstIndex(where: { $0.node == parent }) else {
@@ -274,7 +236,7 @@ fileprivate extension [AddLocationCellViewModel] {
 // MARK: - Update custom list
 
 fileprivate extension CustomListLocationNode {
-    func remove(selectedLocation: LocationNode, with locationList: [AddLocationCellViewModel]) {
+    func remove(selectedLocation: LocationNode, with locationList: [LocationCellViewModel]) {
         if let index = children.firstIndex(of: selectedLocation) {
             children.remove(at: index)
         }
@@ -304,7 +266,7 @@ fileprivate extension CustomListLocationNode {
         }
     }
 
-    private func addSiblings(from locationList: [AddLocationCellViewModel], for node: LocationNode) {
+    private func addSiblings(from locationList: [LocationCellViewModel], for node: LocationNode) {
         guard let parent = node.parent else { return }
         parent.children.forEach { child in
             // adding siblings if they are already selected in snapshot
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
index dfdd791de1..79f847095f 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCell.swift
@@ -8,13 +8,23 @@
 
 import UIKit
 
+//protocol LocationCellDelegate: AnyObject {
+//    func toggle(cell: LocationCell)
+//}
+
 protocol LocationCellDelegate: AnyObject {
-    func toggle(cell: LocationCell)
+    func toggleExpanding(cell: LocationCell)
+    func toggleSelection(cell: LocationCell)
 }
 
 class LocationCell: UITableViewCell {
     weak var delegate: LocationCellDelegate?
 
+    private let container: UIStackView = {
+        let view = UIStackView()
+        return view
+    }()
+
     private let locationLabel: UILabel = {
         let label = UILabel()
         label.font = UIFont.systemFont(ofSize: 16)
@@ -38,6 +48,14 @@ class LocationCell: UITableViewCell {
         return imageView
     }()
 
+    private let checkboxButton: UIButton = {
+        let button = UIButton()
+        button.setImage(UIImage(systemName: "checkmark.square.fill"), for: .selected)
+        button.setImage(UIImage(systemName: "square"), for: .normal)
+        button.tintColor = .white
+        return button
+    }()
+
     private let collapseButton: UIButton = {
         let button = UIButton(type: .custom)
         button.accessibilityIdentifier = .collapseButton
@@ -46,6 +64,8 @@ class LocationCell: UITableViewCell {
         return button
     }()
 
+    private var behavior: LocationCellBehavior = .select
+
     private let chevronDown = UIImage(resource: .iconChevronDown)
     private let chevronUp = UIImage(resource: .iconChevronUp)
 
@@ -106,7 +126,7 @@ class LocationCell: UITableViewCell {
     override func setSelected(_ selected: Bool, animated: Bool) {
         super.setSelected(selected, animated: animated)
 
-        updateTickImage()
+        updateLeadingImage()
         updateStatusIndicatorColor()
     }
 
@@ -122,6 +142,7 @@ class LocationCell: UITableViewCell {
         selectedBackgroundView = UIView()
         selectedBackgroundView?.backgroundColor = UIColor.Cell.Background.selected
 
+        checkboxButton.addTarget(self, action: #selector(toggleCheckboxButton(_:)), for: .touchUpInside)
         collapseButton.addTarget(self, action: #selector(handleCollapseButton(_:)), for: .touchUpInside)
 
         [locationLabel, tickImageView, statusIndicator, collapseButton].forEach { subview in
@@ -131,42 +152,62 @@ class LocationCell: UITableViewCell {
 
         updateCollapseImage()
         updateAccessibilityCustomActions()
-        updateDisabled()
+//        updateDisabled()
         updateBackgroundColor()
         setLayoutMargins()
 
-        NSLayoutConstraint.activate([
-            tickImageView.leadingAnchor.constraint(equalTo: contentView.layoutMarginsGuide.leadingAnchor),
-            tickImageView.centerYAnchor.constraint(equalTo: contentView.centerYAnchor),
+        contentView.addConstrainedSubviews([
+            tickImageView,
+            statusIndicator,
+            locationLabel,
+            collapseButton,
+            checkboxButton
+        ]) {
+            tickImageView.leadingAnchor.constraint(equalTo: contentView.layoutMarginsGuide.leadingAnchor)
+            tickImageView.centerYAnchor.constraint(equalTo: contentView.centerYAnchor)
 
-            statusIndicator.widthAnchor.constraint(equalToConstant: 16),
-            statusIndicator.heightAnchor.constraint(equalTo: statusIndicator.widthAnchor),
-            statusIndicator.centerXAnchor.constraint(equalTo: tickImageView.centerXAnchor),
-            statusIndicator.centerYAnchor.constraint(equalTo: tickImageView.centerYAnchor),
+            statusIndicator.widthAnchor.constraint(equalToConstant: 16)
+            statusIndicator.heightAnchor.constraint(equalTo: statusIndicator.widthAnchor)
+            statusIndicator.centerXAnchor.constraint(equalTo: tickImageView.centerXAnchor)
+            statusIndicator.centerYAnchor.constraint(equalTo: tickImageView.centerYAnchor)
+
+            checkboxButton.pinEdgeToSuperviewMargin(.leading(.zero))
+            checkboxButton.centerYAnchor.constraint(equalTo: contentView.centerYAnchor)
+            checkboxButton.widthAnchor
+                .constraint(
+                    equalToConstant: 44.0
+                )
+            checkboxButton.heightAnchor.constraint(equalTo: checkboxButton.widthAnchor, multiplier: 1, constant: 0)
 
             locationLabel.leadingAnchor.constraint(
                 equalTo: statusIndicator.trailingAnchor,
                 constant: 12
-            ),
+            )
+
             locationLabel.trailingAnchor.constraint(lessThanOrEqualTo: collapseButton.leadingAnchor)
-                .withPriority(.defaultHigh),
-            locationLabel.topAnchor.constraint(equalTo: contentView.layoutMarginsGuide.topAnchor),
-            locationLabel.bottomAnchor.constraint(equalTo: contentView.layoutMarginsGuide.bottomAnchor),
+                .withPriority(.defaultHigh)
+            locationLabel.pinEdgesToSuperviewMargins(PinnableEdges([.top(.zero), .bottom(.zero)]))
 
             collapseButton.widthAnchor
                 .constraint(
                     equalToConstant: UIMetrics.contentLayoutMargins.leading + UIMetrics
-                        .contentLayoutMargins.trailing + 24
-                ),
-            collapseButton.topAnchor.constraint(equalTo: contentView.topAnchor),
-            collapseButton.trailingAnchor.constraint(equalTo: contentView.trailingAnchor),
-            collapseButton.bottomAnchor.constraint(equalTo: contentView.bottomAnchor),
-        ])
+                        .contentLayoutMargins.trailing + 24.0
+                )
+            collapseButton.pinEdgesToSuperviewMargins(.all().excluding(.leading))
+        }
     }
 
-    private func updateTickImage() {
-        statusIndicator.isHidden = isSelected
-        tickImageView.isHidden = !isSelected
+    private func updateLeadingImage() {
+        switch behavior {
+        case .add:
+            checkboxButton.isHidden = false
+            statusIndicator.isHidden = true
+            tickImageView.isHidden = true
+        case .select:
+            checkboxButton.isHidden = true
+            statusIndicator.isHidden = isSelected
+            tickImageView.isHidden = !isSelected
+        }
     }
 
     private func updateStatusIndicatorColor() {
@@ -221,11 +262,11 @@ class LocationCell: UITableViewCell {
     }
 
     @objc private func handleCollapseButton(_ sender: UIControl) {
-        delegate?.toggle(cell: self)
+        delegate?.toggleExpanding(cell: self)
     }
 
     @objc private func toggleCollapseAccessibilityAction() -> Bool {
-        delegate?.toggle(cell: self)
+        delegate?.toggleExpanding(cell: self)
         return true
     }
 
@@ -262,13 +303,31 @@ class LocationCell: UITableViewCell {
             accessibilityCustomActions = nil
         }
     }
+
+    @objc private func toggleCheckboxButton(_ sender: UIControl) {
+        delegate?.toggleSelection(cell: self)
+    }
+}
+
+enum LocationCellBehavior {
+    case add
+    case select
 }
 
 extension LocationCell {
-    func configureCell(item: LocationCellViewModel) {
+    func configure(item: LocationCellViewModel, behavior: LocationCellBehavior) {
         accessibilityIdentifier = item.node.code
         locationLabel.text = item.node.name
         showsCollapseControl = !item.node.children.isEmpty
         isExpanded = item.node.showsChildren
+        checkboxButton.isSelected = item.isSelected
+        checkboxButton.tintColor = item.isSelected ? .successColor : .white
+
+        setBehavior(behavior)
+    }
+
+    private func setBehavior(_ newBehavior: LocationCellBehavior) {
+        self.behavior = newBehavior
+        updateLeadingImage()
     }
 }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
index 2425413fdd..50418e21d8 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationCellViewModel.swift
@@ -12,6 +12,7 @@ struct LocationCellViewModel: Hashable {
     let section: LocationSection
     let node: LocationNode
     var indentationLevel = 0
+    var isSelected: Bool = false
 
     func hash(into hasher: inout Hasher) {
         hasher.combine(section)
@@ -20,6 +21,38 @@ struct LocationCellViewModel: Hashable {
 
     static func == (lhs: Self, rhs: Self) -> Bool {
         lhs.node == rhs.node &&
-            lhs.section == rhs.section
+        lhs.section == rhs.section &&
+        lhs.isSelected == rhs.isSelected
+    }
+}
+
+extension [LocationCellViewModel] {
+    mutating func addSubNodes(from item: LocationCellViewModel, at indexPath: IndexPath) {
+        let section = LocationSection.allCases[indexPath.section]
+        let row = indexPath.row + 1
+
+        let locations = item.node.children.map {
+            LocationCellViewModel(
+                section: section,
+                node: $0,
+                indentationLevel: item.indentationLevel + 1,
+                isSelected: item.isSelected
+            )
+        }
+
+        if row < count {
+            insert(contentsOf: locations, at: row)
+        } else {
+            append(contentsOf: locations)
+        }
+    }
+
+    mutating func removeSubNodes(from node: LocationNode) {
+        for node in node.children {
+            node.showsChildren = false
+            removeAll(where: { node == $0.node })
+
+            removeSubNodes(from: node)
+        }
     }
 }
diff --git a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
index 857037e362..9c1ca8ca06 100644
--- a/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
+++ b/ios/MullvadVPN/View controllers/SelectLocation/LocationDataSource.swift
@@ -39,7 +39,7 @@ final class LocationDataSource: UITableViewDiffableDataSource<LocationSection, L
                 for: indexPath
                 // swiftlint:disable:next force_cast
             ) as! LocationCell
-            cell.configureCell(item: itemIdentifier)
+            cell.configure(item: itemIdentifier, behavior: .select)
             return cell
         }
 
@@ -307,7 +307,11 @@ extension LocationDataSource: UITableViewDelegate {
 }
 
 extension LocationDataSource: LocationCellDelegate {
-    func toggle(cell: LocationCell) {
+    func toggleSelection(cell: LocationCell) {
+        print("Just put a print statement // Marco")
+    }
+
+    func toggleExpanding(cell: LocationCell) {
         guard let indexPath = tableView.indexPath(for: cell),
               let item = itemIdentifier(for: indexPath) else { return }
 
@@ -321,7 +325,7 @@ extension LocationDataSource: LocationCellDelegate {
         if !isExpanded {
             locationList.addSubNodes(from: item, at: indexPath)
         } else {
-            locationList.recursivelyRemoveSubNodes(from: item.node)
+            locationList.removeSubNodes(from: item.node)
         }
 
         let list = sections.enumerated().map { index, section in
@@ -369,29 +373,3 @@ extension LocationDataSource {
         }
     }
 }
-
-private extension [LocationCellViewModel] {
-    mutating func addSubNodes(from item: LocationCellViewModel, at indexPath: IndexPath) {
-        let section = LocationSection.allCases[indexPath.section]
-        let row = indexPath.row + 1
-
-        let locations = item.node.children.map {
-            LocationCellViewModel(section: section, node: $0, indentationLevel: item.indentationLevel + 1)
-        }
-
-        if row < count {
-            insert(contentsOf: locations, at: row)
-        } else {
-            append(contentsOf: locations)
-        }
-    }
-
-    mutating func recursivelyRemoveSubNodes(from node: LocationNode) {
-        for node in node.children {
-            node.showsChildren = false
-            removeAll(where: { node == $0.node })
-
-            recursivelyRemoveSubNodes(from: node)
-        }
-    }
-}