From 883f9fffb4d87cbec1abf8a26425ddadd7075eec Mon Sep 17 00:00:00 2001
From: Matt Bruce <matt.bruce@verizon.com>
Date: Thu, 24 Aug 2023 16:02:20 -0500
Subject: [PATCH] updated rotor logic

Signed-off-by: Matt Bruce <matt.bruce@verizon.com>
---
 VDSSample/Protocols/CustomRotorable.swift     | 47 ++++++++++++-------
 .../ViewControllers/BaseViewController.swift  |  1 -
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/VDSSample/Protocols/CustomRotorable.swift b/VDSSample/Protocols/CustomRotorable.swift
index 6b303d3..acecd19 100644
--- a/VDSSample/Protocols/CustomRotorable.swift
+++ b/VDSSample/Protocols/CustomRotorable.swift
@@ -9,7 +9,6 @@ import Foundation
 import UIKit
 
 public protocol CustomRotorable: UIViewController {
-    var customRotorCache: [String: [UIView]] { get set }
     var customRotors: [CustomRotorType] { get set }
 }
 
@@ -24,15 +23,41 @@ extension CustomRotorable {
         accessibilityCustomRotors = (accessibilityCustomRotors ?? []).filter { $0.name != name }
         
         //create new rotor
-        let newRotor = UIAccessibilityCustomRotor(name: name) { [weak self] predicate in
-            guard let self else { return nil }
+        let newRotor = AccessibilityCustomRotor(with: name, for: trait, rootView: self.view)
+        
+        //append rotor
+        accessibilityCustomRotors?.append(newRotor)
+    }
+        
+    /// Loads all of the custom rotors for the screen.
+    public func loadCustomRotors() {
+        customRotors.forEach { addCustomRotor(with: $0.name, for: $0.trait) }
+    }
+}
+
+private class AccessibilityCustomRotor: UIAccessibilityCustomRotor {
+    var views: [UIView]?
+    var trait: UIAccessibilityTraits
+    weak var rootView: UIView?
+    
+    init (with name: String, for trait: UIAccessibilityTraits, rootView: UIView){
+        self.rootView = rootView
+        self.trait = trait
+        super.init(name: name, itemSearch: { _ in return nil })
+        self.updateSearch()
+    }
+    
+    func updateSearch() {
+        itemSearchBlock = { [weak self] predicate in
+            
+            guard let self, let rootView = self.rootView else { return nil }
             
             //create the view accessibleElements cache if it doesn't exist
-            if self.customRotorCache[name] == nil {
-                self.customRotorCache[name] = self.view.accessibleElements(with: trait)
+            if self.views == nil {
+                self.views = rootView.accessibleElements(with: trait)
             }
             
-            guard let views = self.customRotorCache[name], !views.isEmpty else { return nil }
+            guard let views = self.views, !views.isEmpty else { return nil }
 
             let currentIndex = views.firstIndex(where: { $0 === predicate.currentItem.targetElement })
             let count = views.count
@@ -63,16 +88,6 @@ extension CustomRotorable {
             
             return UIAccessibilityCustomRotorItemResult(targetElement: views[nextIndex], targetRange: nil)
         }
-        
-        //append rotor
-        accessibilityCustomRotors?.append(newRotor)
-    }
-        
-    /// Loads all of the custom rotors for the screen.
-    public func loadCustomRotors() {
-        customRotorCache.removeAll()
-        UIAccessibility.post(notification: .screenChanged, argument: nil)
-        customRotors.forEach { addCustomRotor(with: $0.name, for: $0.trait) }
     }
 }
 
diff --git a/VDSSample/ViewControllers/BaseViewController.swift b/VDSSample/ViewControllers/BaseViewController.swift
index f268703..1449e03 100644
--- a/VDSSample/ViewControllers/BaseViewController.swift
+++ b/VDSSample/ViewControllers/BaseViewController.swift
@@ -82,7 +82,6 @@ public class BaseViewController<Component: UIView>: UIViewController, Initable ,
     //--------------------------------------------------
     public var subscribers = Set<AnyCancellable>()
     
-    public var customRotorCache: [String: [UIView]] = [:]
     public var customRotors: [CustomRotorType] = [
         CustomRotorType(name: "Links", trait: .link),
         CustomRotorType(name: "Buttons", trait: .button)