Improve iOS a11y sync behavior

compose/ui/ui/src/uikitMain/kotlin/androidx/compose/ui/platform/Accessibility.uikit.kt

@@ -35,7 +35,7 @@ import kotlinx.cinterop.ExportObjCClass
 import kotlinx.cinterop.readValue
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
-import kotlinx.coroutines.delay
+import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.launch
 import platform.CoreGraphics.CGRect
 import platform.CoreGraphics.CGRectMake
@@ -85,6 +85,57 @@ interface AccessibilityDebugLogger {
     fun log(message: Any?)
 }
 
+/**
+ * Enum class representing different kinds of accessibility invalidation.
+ */
+private sealed interface SemanticsTreeInvalidation {
+    /**
+     * Represents a complete invalidation of the semantics tree.
+     */
+    object SemanticsChanged : SemanticsTreeInvalidation
+
+    /**
+     * Represents a change in the bounds of a node with the given [id].
+     */
+    class BoundsChanged(val id: Int) : SemanticsTreeInvalidation
+}
+
+private sealed interface SemanticsTreeSyncStrategy {
+    /**
+     * Returns the updated strategy with the given [invalidation] applied.
+     */
+    fun reduce(invalidation: SemanticsTreeInvalidation): SemanticsTreeSyncStrategy
+
+    object Complete : SemanticsTreeSyncStrategy {
+        override fun reduce(invalidation: SemanticsTreeInvalidation): Complete {
+            return this
+        }
+    }
+
+    class BoundsOnly(id: Int) : SemanticsTreeSyncStrategy {
+        val ids = mutableSetOf(id)
+
+        override fun reduce(invalidation: SemanticsTreeInvalidation): SemanticsTreeSyncStrategy {
+            when (invalidation) {
+                is SemanticsTreeInvalidation.SemanticsChanged -> return Complete
+                is SemanticsTreeInvalidation.BoundsChanged -> {
+                    ids.add(invalidation.id)
+                    return this
+                }
+            }
+        }
+    }
+
+    companion object {
+        fun from(invalidation: SemanticsTreeInvalidation): SemanticsTreeSyncStrategy {
+            return when (invalidation) {
+                is SemanticsTreeInvalidation.SemanticsChanged -> Complete
+                is SemanticsTreeInvalidation.BoundsChanged -> BoundsOnly(invalidation.id)
+            }
+        }
+    }
+}
+
 private class CachedAccessibilityPropertyKey<V>
 
 private object CachedAccessibilityPropertyKeys {
@@ -226,11 +277,27 @@ private class AccessibilityElement(
         return null
     }
 
+    /**
+     * Discards the cache for the properties that are computed from the [SemanticsNode.config].
+     */
     fun discardCache() {
         _cachedConfig = null
         cachedProperties.clear()
     }
 
+    /**
+     * Discards the cached accessibilityFrame for this element and all its descendants recursively.
+     */
+    fun discardCachedAccessibilityFrameRecursively() {
+        if (cachedProperties.remove(CachedAccessibilityPropertyKeys.accessibilityFrame) != null) {
+            for (child in children) {
+                child.discardCachedAccessibilityFrameRecursively()
+            }
+        } else {
+            // Not calculated yet, or the subtree was already discarded. Do nothing.
+        }
+    }
+
     fun dispose() {
         check(isAlive) {
             "AccessibilityElement is already disposed"
@@ -816,10 +883,9 @@ private class AccessibilityContainer(
     }
 }
 
-private sealed interface NodesSyncResult {
-    object NoChanges : NodesSyncResult
-    data class Success(val newElementToFocus: Any?) : NodesSyncResult
-}
+private class NodesSyncResult(
+    val newElementToFocus: Any?
+)
 
 /**
  * A sealed class that represents the options for syncing the Compose SemanticsNode tree with the iOS UIAccessibility tree.
@@ -895,6 +961,8 @@ internal class AccessibilityMediator(
     private var needsInitialRefocusing = true
     private var isAlive = true
 
+    private val invalidationChannel = Channel<SemanticsTreeInvalidation>(capacity = Channel.UNLIMITED)
+
     /**
      * Remembered [AccessibilityDebugLogger] after last sync, if logging is enabled according to
      * [AccessibilitySyncOptions].
@@ -904,12 +972,6 @@ internal class AccessibilityMediator(
 
     var rootSemanticsNodeId: Int = -1
 
-    /**
-     * A value of true indicates that the Compose accessible tree is dirty, meaning that compose
-     * semantics tree was modified since last sync, false otherwise.
-     */
-    private var isCurrentComposeAccessibleTreeDirty = false
-
     /**
      * Job to cancel tree syncing when the mediator is disposed.
      */
@@ -929,13 +991,18 @@ internal class AccessibilityMediator(
     init {
         getAccessibilitySyncOptions().debugLoggerIfEnabled?.log("AccessibilityMediator for ${view} created")
 
-        val updateIntervalMillis = 50L
-        // TODO: this approach was copied from desktop implementation, obviously it has a [updateIntervalMillis] lag
-        //  between the actual change in the semantics tree and the change in the accessibility tree.
-        //  should we use some other approach?
         coroutineScope.launch {
             while (isAlive) {
-                var result: NodesSyncResult
+                var strategy = SemanticsTreeSyncStrategy.from(
+                    invalidationChannel.receive()
+                )
+
+                while (true) {
+                    val result = invalidationChannel.tryReceive()
+
+                    val invalidation = result.getOrNull() ?: break
+                    strategy = strategy.reduce(invalidation)
+                }
 
                 val syncOptions = getAccessibilitySyncOptions()
 
@@ -948,24 +1015,16 @@ internal class AccessibilityMediator(
                 }
 
                 if (shouldPerformSync) {
+                    var result: NodesSyncResult
+
                     val time = measureTime {
-                        result = sync()
+                        result = sync(strategy)
                     }
 
-                    when (val immutableResult = result) {
-                        is NodesSyncResult.NoChanges -> {
-                            // Do nothing
-                        }
-
-                        is NodesSyncResult.Success -> {
-                            debugLogger?.log("AccessibilityMediator.sync took $time")
+                    debugLogger?.log("AccessibilityMediator.sync took $time")
 
-                            UIAccessibilityPostNotification(UIAccessibilityLayoutChangedNotification, immutableResult.newElementToFocus)
-                        }
-                    }
+                    UIAccessibilityPostNotification(UIAccessibilityLayoutChangedNotification, result.newElementToFocus)
                 }
-
-                delay(updateIntervalMillis)
             }
         }
     }
@@ -979,15 +1038,17 @@ internal class AccessibilityMediator(
     fun onSemanticsChange() {
         debugLogger?.log("onSemanticsChange")
 
-        isCurrentComposeAccessibleTreeDirty = true
+        invalidationKind = SemanticsTreeInvalidationKind.COMPLETE
+        invalidationChannel.trySend(Unit)
     }
 
     fun onLayoutChange(nodeId: Int) {
         debugLogger?.log("onLayoutChange (nodeId=$nodeId)")
 
-        // TODO: Only recompute the layout-related properties of the node
-        isCurrentComposeAccessibleTreeDirty = true
+        invalidatedBoundsNodeIds.add(nodeId)
 
+        // unprocessedInvalidationKind will be set to BOUNDS in sync(), it's a strict subset of COMPLETE
+        invalidationChannel.trySend(Unit)
     }
 
     fun dispose() {
@@ -1075,22 +1136,37 @@ internal class AccessibilityMediator(
 
     /**
      * Syncs the accessibility tree with the current semantics tree.
-     * TODO: Does a full tree traversal on every sync, expect changes from Google, they are also aware
-     *  of the issue and associated performance overhead.
      */
-    private fun sync(): NodesSyncResult {
-        // TODO: investigate what needs to be done to reflect that this hierarchy is probably covered
-        //   by sibling overlay
+    private fun sync(strategy: SemanticsTreeSyncStrategy): NodesSyncResult {
+        when (strategy) {
+            is SemanticsTreeSyncStrategy.Complete -> {
+                return completeSync()
+            }
 
-        if (!isCurrentComposeAccessibleTreeDirty) {
-            return NodesSyncResult.NoChanges
+            is SemanticsTreeSyncStrategy.BoundsOnly -> {
+                for (id in strategy.ids) {
+                    val element = accessibilityElementsMap[id]
+                    element?.discardCachedAccessibilityFrameRecursively()
+                }
+
+                return NodesSyncResult(null)
+            }
         }
+    }
 
+    /**
+     * Performs a complete sync of the accessibility tree with the current semantics tree.
+     *
+     * TODO: Does a full tree traversal on every sync, expect changes from Google, they are also aware
+     *  of the issue and associated performance overhead.
+     */
+
+    private fun completeSync(): NodesSyncResult {
+        // TODO: investigate what needs to be done to reflect that this hierarchy is probably covered
+        //   by sibling overlay or another UIView hierarchy represented by other mediator
         val rootSemanticsNode = owner.rootSemanticsNode
         rootSemanticsNodeId = rootSemanticsNode.id
 
-        isCurrentComposeAccessibleTreeDirty = false
-
         check(!view.isAccessibilityElement) {
             "Root view must not be an accessibility element"
         }
@@ -1108,11 +1184,13 @@ internal class AccessibilityMediator(
         // TODO: in future the focused element could be the interop UIView that is detached from the
         //  hierarchy, but still maintains the focus until the GC collects it, or AX services detect
         //  that it's not reachable anymore through containment chain
-        val isFocusedElementDead = focusedElement?.let {
+        val isFocusedElementAlive = focusedElement?.let {
             val accessibilityElement = it as? AccessibilityElement
             accessibilityElement?.isAlive ?: false
         } ?: false
 
+        val isFocusedElementDead = !isFocusedElementAlive
+
         val needsRefocusing = needsInitialRefocusing || isFocusedElementDead
 
         val newElementToFocus = if (needsRefocusing) {
@@ -1132,7 +1210,7 @@ internal class AccessibilityMediator(
             null
         }
 
-        return NodesSyncResult.Success(newElementToFocus)
+        return NodesSyncResult(newElementToFocus)
     }
 }