fix: Optimize class chain lookup

Author: mykola-mokhnach

WebDriverAgentLib/Categories/XCUIElement+FBClassChain.m

@@ -11,9 +11,15 @@
 #import "FBClassChainQueryParser.h"
 #import "FBXCodeCompatibility.h"
 #import "FBExceptions.h"
+#import "XCUIElement+FBUtilities.h"
 
 @implementation XCUIElement (FBClassChain)
 
+// Walks a single upfront snapshot of `self` in memory to resolve every
+// non-indexed chain segment, instead of resolving an intermediate live
+// XCUIElement (and paying an accessibility round trip) just to obtain a
+// query root for the next segment. A live element is only ever resolved
+// once, for the final match(es), via a uid-predicate lookup.
 - (NSArray<XCUIElement *> *)fb_descendantsMatchingClassChain:(NSString *)classChainQuery shouldReturnAfterFirstMatch:(BOOL)shouldReturnAfterFirstMatch
 {
   NSError *error;
@@ -23,74 +29,85 @@ @implementation XCUIElement (FBClassChain)
     return nil;
   }
   NSMutableArray<FBClassChainItem *> *lookupChain = parsedChain.elements.mutableCopy;
+  NSArray<id<FBXCElementSnapshot>> *currentRoots = @[[self fb_customSnapshot]];
   FBClassChainItem *chainItem = lookupChain.firstObject;
-  XCUIElement *currentRoot = self;
-  XCUIElementQuery *query = [currentRoot fb_queryWithChainItem:chainItem query:nil];
+  NSArray<id<FBXCElementSnapshot>> *candidates = [self.class fb_snapshotsMatchingItem:chainItem inRoots:currentRoots];
   [lookupChain removeObjectAtIndex:0];
   while (lookupChain.count > 0) {
-    BOOL isRootChanged = NO;
     if (nil != chainItem.position) {
-      // It is necessary to resolve the query if intermediate element index is not zero or one,
+      // It is necessary to resolve the position if intermediate element index is not zero or one,
       // because predicates don't support search by indexes
-      NSArray<XCUIElement *> *currentRootMatch = [self.class fb_matchingElementsWithItem:chainItem
-                                                                                   query:query
-                                                             shouldReturnAfterFirstMatch:nil];
+      NSArray<id<FBXCElementSnapshot>> *currentRootMatch = [self.class fb_matchingSnapshotsWithItem:chainItem
+                                                                                           candidates:candidates
+                                                                          shouldReturnAfterFirstMatch:nil];
       if (0 == currentRootMatch.count) {
         return @[];
       }
-      currentRoot = currentRootMatch.firstObject;
-      isRootChanged = YES;
+      currentRoots = @[currentRootMatch.firstObject];
+    } else {
+      currentRoots = candidates;
     }
-    chainItem = [lookupChain firstObject];
-    query = [currentRoot fb_queryWithChainItem:chainItem query:isRootChanged ? nil : query];
+    chainItem = lookupChain.firstObject;
+    candidates = [self.class fb_snapshotsMatchingItem:chainItem inRoots:currentRoots];
     [lookupChain removeObjectAtIndex:0];
   }
-  return [self.class fb_matchingElementsWithItem:chainItem
-                                           query:query
-                     shouldReturnAfterFirstMatch:@(shouldReturnAfterFirstMatch)];
+  NSArray<id<FBXCElementSnapshot>> *matchedSnapshots = [self.class fb_matchingSnapshotsWithItem:chainItem
+                                                                                      candidates:candidates
+                                                                     shouldReturnAfterFirstMatch:@(shouldReturnAfterFirstMatch)];
+  return [self fb_filterDescendantsWithSnapshots:matchedSnapshots onlyChildren:NO];
 }
 
-- (XCUIElementQuery *)fb_queryWithChainItem:(FBClassChainItem *)item query:(nullable XCUIElementQuery *)query
++ (NSArray<id<FBXCElementSnapshot>> *)fb_snapshotsMatchingItem:(FBClassChainItem *)item inRoots:(NSArray<id<FBXCElementSnapshot>> *)roots
 {
-  if (item.isDescendant) {
-    if (query) {
-      query = [query descendantsMatchingType:item.type];
-    } else {
-      query = [self.fb_query descendantsMatchingType:item.type];
-    }
-  } else {
-    if (query) {
-      query = [query childrenMatchingType:item.type];
+  NSMutableArray<id<FBXCElementSnapshot>> *typeMatches = [NSMutableArray array];
+  for (id<FBXCElementSnapshot> root in roots) {
+    if (item.isDescendant) {
+      // descendantsByFilteringWithBlock: includes the receiver itself if it
+      // matches the filter, unlike XCUIElementQuery's descendantsMatchingType:,
+      // so the root has to be excluded explicitly here.
+      [typeMatches addObjectsFromArray:[root descendantsByFilteringWithBlock:^BOOL(id<FBXCElementSnapshot> snapshot) {
+        return snapshot != root && (item.type == XCUIElementTypeAny || snapshot.elementType == item.type);
+      }]];
     } else {
-      query = [self.fb_query childrenMatchingType:item.type];
+      for (id<FBXCElementSnapshot> child in root.children) {
+        if (item.type == XCUIElementTypeAny || child.elementType == item.type) {
+          [typeMatches addObject:child];
+        }
+      }
     }
   }
-  if (item.predicates) {
-    for (FBAbstractPredicateItem *predicate in item.predicates) {
-      if ([predicate isKindOfClass:FBSelfPredicateItem.class]) {
-        query = [query matchingPredicate:predicate.value];
-      } else if ([predicate isKindOfClass:FBDescendantPredicateItem.class]) {
-        query = [query containingPredicate:predicate.value];
+  for (FBAbstractPredicateItem *predicateItem in item.predicates) {
+    if ([predicateItem isKindOfClass:FBSelfPredicateItem.class]) {
+      typeMatches = [[typeMatches filteredArrayUsingPredicate:predicateItem.value] mutableCopy];
+    } else if ([predicateItem isKindOfClass:FBDescendantPredicateItem.class]) {
+      NSMutableArray<id<FBXCElementSnapshot>> *containingMatches = [NSMutableArray array];
+      for (id<FBXCElementSnapshot> candidate in typeMatches) {
+        NSArray<id<FBXCElementSnapshot>> *matchingDescendants = [candidate descendantsByFilteringWithBlock:^BOOL(id<FBXCElementSnapshot> descendant) {
+          return descendant != candidate && [predicateItem.value evaluateWithObject:descendant];
+        }];
+        if (matchingDescendants.count > 0) {
+          [containingMatches addObject:candidate];
+        }
       }
+      typeMatches = containingMatches;
     }
   }
-  return query;
+  return typeMatches.copy;
 }
 
-+ (NSArray<XCUIElement *> *)fb_matchingElementsWithItem:(FBClassChainItem *)item query:(XCUIElementQuery *)query shouldReturnAfterFirstMatch:(nullable NSNumber *)shouldReturnAfterFirstMatch
++ (NSArray<id<FBXCElementSnapshot>> *)fb_matchingSnapshotsWithItem:(FBClassChainItem *)item candidates:(NSArray<id<FBXCElementSnapshot>> *)candidates shouldReturnAfterFirstMatch:(nullable NSNumber *)shouldReturnAfterFirstMatch
 {
   if (1 == item.position.integerValue || (0 == item.position.integerValue && shouldReturnAfterFirstMatch.boolValue)) {
-    XCUIElement *result = query.fb_firstMatch;
+    id<FBXCElementSnapshot> result = candidates.firstObject;
     return result ? @[result] : @[];
   }
-  NSArray<XCUIElement *> *allMatches = query.fb_allMatches;
   if (0 == item.position.integerValue) {
-    return allMatches;
+    return candidates;
   }
-  if (allMatches.count >= (NSUInteger)ABS(item.position.integerValue)) {
+  if (candidates.count >= (NSUInteger)ABS(item.position.integerValue)) {
     return item.position.integerValue > 0
-      ? @[[allMatches objectAtIndex:item.position.integerValue - 1]]
-      : @[[allMatches objectAtIndex:allMatches.count + item.position.integerValue]];
+      ? @[[candidates objectAtIndex:item.position.integerValue - 1]]
+      : @[[candidates objectAtIndex:candidates.count + item.position.integerValue]];
   }
   return @[];
 }