From 514dba15b9726866d3ca0f137b4203fe63f44b96 Mon Sep 17 00:00:00 2001
From: Christoffer Winterkvist <christoffer@winterkvist.com>
Date: Wed, 12 Jun 2024 20:41:00 +0200
Subject: [PATCH] refactor: update SystemWindowRelativeFocus for state tracking

Refactor SystemWindowRelativeFocus to a final class to handle state
tracking for consumed windows and previous direction. Add new
instance variables and methods to reset state and insert consumed
windows. Update usages of SystemWindowRelativeFocus' methods to
use instance methods instead of static methods.

Add consumed windows tracking to ensure windows are not reused.
Update SystemCommandRunner to initialize and use the new
SystemWindowRelativeFocus instance.
---
 .../System/SystemWindowRelativeFocus.swift    | 29 ++++++++++++++++---
 .../Core/Runners/SystemCommandRunner.swift    | 22 ++++++++------
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/App/Sources/Core/Runners/System/SystemWindowRelativeFocus.swift b/App/Sources/Core/Runners/System/SystemWindowRelativeFocus.swift
index bc2b85f0..aedc7549 100644
--- a/App/Sources/Core/Runners/System/SystemWindowRelativeFocus.swift
+++ b/App/Sources/Core/Runners/System/SystemWindowRelativeFocus.swift
@@ -4,12 +4,26 @@ import AXEssibility
 import Foundation
 import Windows
 
-enum SystemWindowRelativeFocus {
+final class SystemWindowRelativeFocus {
   enum Direction {
     case up, down, left, right
   }
 
-  static func run(_ direction: Direction, snapshot: UserSpace.Snapshot) throws {
+  var consumedWindows = Set<WindowModel>()
+  var previousDirection: Direction?
+
+  init() {}
+
+  func reset() {
+    consumedWindows.removeAll()
+  }
+
+  func run(_ direction: Direction, snapshot: UserSpace.Snapshot) throws {
+    if direction != previousDirection {
+      previousDirection = direction
+      consumedWindows.removeAll()
+    }
+
     var windows = indexWindowsInStage(getWindows())
     let frontMostApplication = snapshot.frontMostApplication
     let frontMostAppElement = AppAccessibilityElement(frontMostApplication.ref.processIdentifier)
@@ -25,6 +39,7 @@ enum SystemWindowRelativeFocus {
 
       if window.id == focusedWindow.id {
         activeWindow = window
+        consumedWindows.insert(window)
         windows.remove(at: offset)
         break
       }
@@ -35,6 +50,8 @@ enum SystemWindowRelativeFocus {
       windows.removeFirst()
     }
 
+    windows.removeAll(where: { consumedWindows.contains($0) })
+
     guard let activeWindow = activeWindow else { return }
 
     var matchedWindow: WindowModel?
@@ -51,6 +68,8 @@ enum SystemWindowRelativeFocus {
 
     guard let matchedWindow else { return }
 
+    consumedWindows.insert(matchedWindow)
+
     let processIdentifier = pid_t(matchedWindow.ownerPid.rawValue)
     guard let runningApplication = NSRunningApplication(processIdentifier: processIdentifier) else { return }
     let appElement = AppAccessibilityElement(processIdentifier)
@@ -64,13 +83,15 @@ enum SystemWindowRelativeFocus {
     match?.performAction(.raise)
   }
 
-  private static func getWindows() -> [WindowModel] {
+  // MARK: Private methods
+
+  private func getWindows() -> [WindowModel] {
     let options: CGWindowListOption = [.optionOnScreenOnly, .excludeDesktopElements]
     let windowModels: [WindowModel] = ((try? WindowsInfo.getWindows(options)) ?? [])
     return windowModels
   }
 
-  private static func indexWindowsInStage(_ models: [WindowModel]) -> [WindowModel] {
+  private func indexWindowsInStage(_ models: [WindowModel]) -> [WindowModel] {
     let excluded = ["WindowManager", "Window Server"]
     let minimumSize = CGSize(width: 150, height: 150)
     let windows: [WindowModel] = models
diff --git a/App/Sources/Core/Runners/SystemCommandRunner.swift b/App/Sources/Core/Runners/SystemCommandRunner.swift
index d6a86691..84be9e94 100644
--- a/App/Sources/Core/Runners/SystemCommandRunner.swift
+++ b/App/Sources/Core/Runners/SystemCommandRunner.swift
@@ -10,8 +10,9 @@ import Windows
 final class SystemCommandRunner: @unchecked Sendable {
   var machPort: MachPortEventController?
 
-  private let applicationStore: ApplicationStore
   private let applicationActivityMonitor: ApplicationActivityMonitor<UserSpace.Application>
+  private let applicationStore: ApplicationStore
+  private let relativeFocus: SystemWindowRelativeFocus
   private let workspace: WorkspaceProviding
 
   private var flagsChangedSubscription: AnyCancellable?
@@ -31,10 +32,13 @@ final class SystemCommandRunner: @unchecked Sendable {
       .compactMap { $0 }
       .sink { [weak self] flags in
         guard let self else { return }
-        WindowStore.shared.state.interactive = flags != CGEventFlags.maskNonCoalesced
-        if WindowStore.shared.state.interactive == false {
-          self.frontMostIndex = 0
-          self.visibleMostIndex = 0
+        Task { @MainActor in
+          WindowStore.shared.state.interactive = flags != CGEventFlags.maskNonCoalesced
+          if WindowStore.shared.state.interactive == false {
+            self.frontMostIndex = 0
+            self.visibleMostIndex = 0
+            self.relativeFocus.reset()
+          }
         }
       }
   }
@@ -71,13 +75,13 @@ final class SystemCommandRunner: @unchecked Sendable {
       case .missionControl:
         Dock.run(.missionControl)
       case .moveFocusToNextWindowUpwards:
-        try SystemWindowRelativeFocus.run(.up, snapshot: snapshot)
+        try relativeFocus.run(.up, snapshot: snapshot)
       case .moveFocusToNextWindowDownwards:
-        try SystemWindowRelativeFocus.run(.down, snapshot: snapshot)
+        try relativeFocus.run(.down, snapshot: snapshot)
       case .moveFocusToNextWindowOnLeft:
-        try SystemWindowRelativeFocus.run(.left, snapshot: snapshot)
+        try relativeFocus.run(.left, snapshot: snapshot)
       case .moveFocusToNextWindowOnRight:
-        try SystemWindowRelativeFocus.run(.right, snapshot: snapshot)
+        try relativeFocus.run(.right, snapshot: snapshot)
       }
     }
   }