bugfix: Restore notches on external displays after sleep/wake

Sources/Notch/NotchFleet.swift

@@ -114,6 +114,7 @@ final class NotchFleet {
     /// What the fleet settled on, for tests that need to see panels come and
     /// go rather than take our word for it.
     var controllersForTesting: [NotchWindowController] { Array(controllers.values) }
+    var controllersByDisplayKeyForTesting: [NSNumber: NotchWindowController] { controllers }
 
     init(scope: NotchScreenScope, edge: NotchEdge) {
         self.scope = scope
@@ -407,6 +408,10 @@ final class NotchFleet {
         reconcile(screens: NSScreen.screens)
     }
 
+    func reconcileForTesting(screens: [NSScreen]) {
+        reconcile(screens: screens)
+    }
+
     private func reconcile(screens: [NSScreen]) {
         let desired: [NSScreen]
         switch scope {
@@ -433,11 +438,21 @@ final class NotchFleet {
         // instead of tearing a panel down and building another.
         if scope == .mainDisplay, controllers.count == 1,
            let screen = desired.first, let controller = controllers.values.first {
+            controller.displayPreference = displayPreference
             controller.assignedScreen = screen
-            controller.relocate()
+            controller.relocate(on: screen)
             return
         }
         let plan = Self.planReconciliation(current: Set(controllers.keys), desired: keys)
+        let additions = Set(plan.add)
+        let removals = Set(plan.remove)
+        for (screen, id) in zip(desired, keys)
+        where !additions.contains(id) && !removals.contains(id) {
+            guard let controller = controllers[id] else { continue }
+            controller.displayPreference = displayPreference
+            controller.assignedScreen = screen
+            controller.relocate(on: screen)
+        }
         for id in plan.remove {
             controllers[id]?.retire()
             controllers[id] = nil

Sources/Notch/NotchWindowController.swift

@@ -315,15 +315,23 @@ final class NotchWindowController {
     /// been assigned, which is the single-controller case `.mainDisplay`
     /// scope leaves it in.
     func currentScreen() -> NSScreen? {
+        currentScreen(in: NSScreen.screens)
+    }
+
+    func currentScreen(in screens: [NSScreen]) -> NSScreen? {
         if let assigned = assignedScreen,
-           NSScreen.screens.contains(where: { $0 === assigned }) {
+           screens.contains(where: { $0 === assigned }) {
             return assigned
         }
-        return NotchGeometry.preferredScreen(from: NSScreen.screens, preference: displayPreference)
+        return NotchGeometry.preferredScreen(from: screens, preference: displayPreference)
     }
 
     func relocate(cellCount: Int? = nil) {
         guard let screen = currentScreen() else { return }
+        relocate(on: screen, cellCount: cellCount)
+    }
+
+    func relocate(on screen: NSScreen, cellCount: Int? = nil) {
         model.adopt(screen: screen)
         let size = model.panelSize(cellCount: cellCount ?? model.snapshots.count)
         let frame = NotchGeometry.panelFrame(

Tests/NotchLayoutTests.swift

@@ -1056,6 +1056,105 @@ final class NotchFleetScopeTests: XCTestCase {
     }
 }
 
+@MainActor
+private final class SyntheticNSScreen: NSScreen {
+    private let number: Int
+
+    init(number: Int) {
+        self.number = number
+        super.init()
+    }
+
+    override var frame: CGRect {
+        CGRect(x: CGFloat(number) * 1600, y: 0, width: 1440, height: 900)
+    }
+
+    override var visibleFrame: CGRect { frame }
+
+    override var deviceDescription: [NSDeviceDescriptionKey: Any] {
+        [NSDeviceDescriptionKey("NSScreenNumber"): NSNumber(value: number)]
+    }
+}
+
+@MainActor
+final class NotchFleetRebindingTests: XCTestCase {
+    private func screen(_ number: Int) -> SyntheticNSScreen {
+        SyntheticNSScreen(number: number)
+    }
+
+    func testRetainedControllerRebindsToReplacementScreenWithSameKey() {
+        let fleet = NotchFleet(scope: .allDisplays, edge: .right)
+        defer { fleet.stop() }
+
+        let originalScreen = screen(901)
+        let key = NotchFleet.key(for: originalScreen)
+        fleet.reconcileForTesting(screens: [originalScreen])
+        let controller = try! XCTUnwrap(fleet.controllersByDisplayKeyForTesting[key])
+        controller.model.edge = .bottom
+        controller.model.alongOffset = 37
+
+        let replacementScreen = screen(901)
+        fleet.reconcileForTesting(screens: [replacementScreen])
+
+        XCTAssertTrue(fleet.controllersByDisplayKeyForTesting[key] === controller)
+        XCTAssertTrue(controller.assignedScreen === replacementScreen)
+        XCTAssertTrue(controller.currentScreen(in: [replacementScreen]) === replacementScreen)
+        XCTAssertEqual(controller.model.edge, .bottom)
+        XCTAssertEqual(controller.model.alongOffset, 37)
+        XCTAssertEqual(controller.panelFrameForTesting?.minY ?? -1,
+                       replacementScreen.frame.minY, accuracy: 1)
+    }
+
+    func testShrinkingAndRestoringTopologyLeavesOneCorrectControllerPerScreenAndIsIdempotent() {
+        let fleet = NotchFleet(scope: .allDisplays, edge: .right)
+        defer { fleet.stop() }
+        fleet.apply(.alwaysShow)
+
+        let original = [screen(911), screen(912), screen(913)]
+        fleet.reconcileForTesting(screens: original)
+        let originalControllers = fleet.controllersByDisplayKeyForTesting
+        let removedKey = NotchFleet.key(for: original[1])
+        for controller in originalControllers.values {
+            XCTAssertTrue(controller.model.isAlwaysOn)
+            XCTAssertTrue(controller.model.isExpanded)
+        }
+
+        let shrunken = [screen(911), screen(913)]
+        fleet.reconcileForTesting(screens: shrunken)
+
+        XCTAssertEqual(Set(fleet.controllersByDisplayKeyForTesting.keys), Set(shrunken.map(NotchFleet.key)))
+        XCTAssertNil(fleet.controllersByDisplayKeyForTesting[removedKey])
+        XCTAssertTrue(fleet.controllersByDisplayKeyForTesting[NotchFleet.key(for: shrunken[0])] === originalControllers[NotchFleet.key(for: original[0])])
+        XCTAssertTrue(fleet.controllersByDisplayKeyForTesting[NotchFleet.key(for: shrunken[1])] === originalControllers[NotchFleet.key(for: original[2])])
+        XCTAssertTrue(fleet.controllersByDisplayKeyForTesting[NotchFleet.key(for: shrunken[0])]?.model.isAlwaysOn == true)
+        XCTAssertTrue(fleet.controllersByDisplayKeyForTesting[NotchFleet.key(for: shrunken[1])]?.model.isAlwaysOn == true)
+
+        let restored = [screen(911), screen(912), screen(913)]
+        fleet.reconcileForTesting(screens: restored)
+        let finalControllers = fleet.controllersByDisplayKeyForTesting
+
+        XCTAssertEqual(finalControllers.count, restored.count)
+        XCTAssertEqual(Set(finalControllers.keys), Set(restored.map(NotchFleet.key)))
+        for currentScreen in restored {
+            let controller = try! XCTUnwrap(finalControllers[NotchFleet.key(for: currentScreen)])
+            XCTAssertTrue(controller.assignedScreen === currentScreen)
+            XCTAssertTrue(controller.currentScreen(in: restored) === currentScreen)
+            XCTAssertTrue(controller.model.isAlwaysOn)
+            XCTAssertTrue(controller.model.isExpanded)
+        }
+        XCTAssertFalse(finalControllers[removedKey] === originalControllers[removedKey])
+
+        fleet.reconcileForTesting(screens: restored)
+
+        XCTAssertEqual(fleet.controllersByDisplayKeyForTesting.count, restored.count)
+        for currentScreen in restored {
+            let key = NotchFleet.key(for: currentScreen)
+            XCTAssertTrue(fleet.controllersByDisplayKeyForTesting[key] === finalControllers[key])
+            XCTAssertTrue(fleet.controllersByDisplayKeyForTesting[key]?.assignedScreen === currentScreen)
+        }
+    }
+}
+
 /// Renaming the app renames its defaults domain, so every setting moves to a
 /// new empty one unless it is carried across.
 final class RenameMigrationTests: XCTestCase {