fix: Optimize class chain lookup

WebDriverAgentLib/Categories/XCUIElement+FBClassChain.m

@@ -11,6 +11,7 @@
 #import "FBClassChainQueryParser.h"
 #import "FBXCodeCompatibility.h"
 #import "FBExceptions.h"
+#import "XCUIElement+FBUtilities.h"
 
 @implementation XCUIElement (FBClassChain)
 
@@ -22,16 +23,46 @@ @implementation XCUIElement (FBClassChain)
     @throw [NSException exceptionWithName:FBClassChainQueryParseException reason:error.localizedDescription userInfo:error.userInfo];
     return nil;
   }
-  NSMutableArray<FBClassChainItem *> *lookupChain = parsedChain.elements.mutableCopy;
+  // The snapshot-walk strategy below only pays off when an intermediate
+  // (non-final) segment carries an explicit position: that is the only
+  // shape where the query-based strategy has to resolve a live element
+  // mid-chain - and pay an accessibility round trip - just to keep building
+  // the next segment's query. Every other shape (including the common case
+  // of zero or one position, on the final segment only) already resolves in
+  // a single round trip with the query-based strategy, while the snapshot
+  // walk always pays for one full upfront subtree snapshot regardless of
+  // whether it is actually needed - a bad trade in deep/large trees. See
+  // https://github.com/appium/WebDriverAgent/pull/1194#issuecomment-5156633352
+  NSArray<FBClassChainItem *> *chainItems = parsedChain.elements;
+  return [self.class fb_hasIntermediatePosition:chainItems]
+    ? [self fb_snapshotDescendantsMatchingChainItems:chainItems shouldReturnAfterFirstMatch:shouldReturnAfterFirstMatch]
+    : [self fb_queryDescendantsMatchingChainItems:chainItems shouldReturnAfterFirstMatch:shouldReturnAfterFirstMatch];
+}
+
++ (BOOL)fb_hasIntermediatePosition:(NSArray<FBClassChainItem *> *)chainItems
+{
+  for (NSUInteger i = 0; i + 1 < chainItems.count; i++) {
+    if (nil != chainItems[i].position) {
+      return YES;
+    }
+  }
+  return NO;
+}
+
+#pragma mark - Query-based strategy
+#pragma mark (single accessibility round trip; used whenever no intermediate
+#pragma mark  segment carries an explicit position)
+
+- (NSArray<XCUIElement *> *)fb_queryDescendantsMatchingChainItems:(NSArray<FBClassChainItem *> *)chainItems shouldReturnAfterFirstMatch:(BOOL)shouldReturnAfterFirstMatch
+{
+  NSMutableArray<FBClassChainItem *> *lookupChain = chainItems.mutableCopy;
   FBClassChainItem *chainItem = lookupChain.firstObject;
   XCUIElement *currentRoot = self;
   XCUIElementQuery *query = [currentRoot fb_queryWithChainItem:chainItem query:nil];
   [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,
-      // because predicates don't support search by indexes
       NSArray<XCUIElement *> *currentRootMatch = [self.class fb_matchingElementsWithItem:chainItem
                                                                                    query:query
                                                              shouldReturnAfterFirstMatch:nil];
@@ -95,4 +126,116 @@ - (XCUIElementQuery *)fb_queryWithChainItem:(FBClassChainItem *)item query:(null
   return @[];
 }
 
+#pragma mark - Snapshot-based strategy
+#pragma mark (single upfront snapshot walked in memory; used when an
+#pragma mark  intermediate segment has an explicit position, avoiding one
+#pragma mark  extra accessibility round trip per such segment)
+
+- (NSArray<XCUIElement *> *)fb_snapshotDescendantsMatchingChainItems:(NSArray<FBClassChainItem *> *)chainItems shouldReturnAfterFirstMatch:(BOOL)shouldReturnAfterFirstMatch
+{
+  NSMutableArray<FBClassChainItem *> *lookupChain = chainItems.mutableCopy;
+  // Reuse an already-taken snapshot of `self` if one is available (e.g. the
+  // caller just resolved/inspected this same element) instead of always
+  // paying for a fresh one.
+  NSArray<id<FBXCElementSnapshot>> *currentRoots = @[self.lastSnapshot ?: [self fb_customSnapshot]];
+  FBClassChainItem *chainItem = lookupChain.firstObject;
+  NSArray<id<FBXCElementSnapshot>> *candidates = [self.class fb_snapshotsMatchingItem:chainItem inRoots:currentRoots];
+  [lookupChain removeObjectAtIndex:0];
+  while (lookupChain.count > 0) {
+    if (nil != chainItem.position) {
+      // An explicit position always narrows the match set down to a single
+      // element, which becomes the sole root for the rest of the chain, so
+      // it has to be resolved now instead of being folded into `candidates`
+      // like an unindexed segment would be
+      NSArray<id<FBXCElementSnapshot>> *currentRootMatch = [self.class fb_matchingSnapshotsWithItem:chainItem
+                                                                                           candidates:candidates
+                                                                          shouldReturnAfterFirstMatch:nil];
+      if (0 == currentRootMatch.count) {
+        return @[];
+      }
+      currentRoots = @[currentRootMatch.firstObject];
+    } else {
+      currentRoots = candidates;
+    }
+    chainItem = lookupChain.firstObject;
+    candidates = [self.class fb_snapshotsMatchingItem:chainItem inRoots:currentRoots];
+    [lookupChain removeObjectAtIndex:0];
+  }
+  NSArray<id<FBXCElementSnapshot>> *matchedSnapshots = [self.class fb_matchingSnapshotsWithItem:chainItem
+                                                                                      candidates:candidates
+                                                                     shouldReturnAfterFirstMatch:@(shouldReturnAfterFirstMatch)];
+  return [self fb_filterDescendantsWithSnapshots:matchedSnapshots onlyChildren:NO];
+}
+
++ (NSArray<id<FBXCElementSnapshot>> *)fb_snapshotsMatchingItem:(FBClassChainItem *)item inRoots:(NSArray<id<FBXCElementSnapshot>> *)roots
+{
+  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 {
+      for (id<FBXCElementSnapshot> child in root.children) {
+        if (item.type == XCUIElementTypeAny || child.elementType == item.type) {
+          [typeMatches addObject:child];
+        }
+      }
+    }
+  }
+  if (roots.count > 1) {
+    // Overlapping roots (e.g. a previous segment matched both an ancestor
+    // and its own descendant) can otherwise yield the same snapshot twice,
+    // which would skew positional selection ([2], [-1], etc.) compared to
+    // the XCUIElementQuery-based matching this replaced, which always
+    // operated on a de-duplicated element set.
+    NSMutableArray<id<FBXCElementSnapshot>> *dedupedMatches = [NSMutableArray arrayWithCapacity:typeMatches.count];
+    NSHashTable<id<FBXCElementSnapshot>> *seenMatches = [NSHashTable hashTableWithOptions:NSHashTableObjectPointerPersonality];
+    for (id<FBXCElementSnapshot> match in typeMatches) {
+      if (![seenMatches containsObject:match]) {
+        [seenMatches addObject:match];
+        [dedupedMatches addObject:match];
+      }
+    }
+    typeMatches = dedupedMatches;
+  }
+  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 typeMatches.copy;
+}
+
++ (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)) {
+    id<FBXCElementSnapshot> result = candidates.firstObject;
+    return result ? @[result] : @[];
+  }
+  if (0 == item.position.integerValue) {
+    return candidates;
+  }
+  if (candidates.count >= (NSUInteger)ABS(item.position.integerValue)) {
+    return item.position.integerValue > 0
+      ? @[[candidates objectAtIndex:item.position.integerValue - 1]]
+      : @[[candidates objectAtIndex:candidates.count + item.position.integerValue]];
+  }
+  return @[];
+}
+
 @end

WebDriverAgentTests/IntegrationApp/Classes/ViewController.m

@@ -48,6 +48,40 @@ - (IBAction)didTapButton:(UIButton *)button
   button.selected = !button.selected;
 }
 
+- (IBAction)goToDeepHierarchy:(id)sender
+{
+  // Plain UIViews with fixed frames only - no Auto Layout constraints and no
+  // specialized subclasses (e.g. UITextView) that carry their own layout/text
+  // engines, which can make a deep nested chain pathologically expensive to
+  // lay out. This page exists purely as a fixture for exercising element
+  // lookups (e.g. class chain locators) against a deep accessibility tree.
+  UIViewController *deepHierarchyViewController = [UIViewController new];
+  deepHierarchyViewController.view.backgroundColor = UIColor.whiteColor;
+  deepHierarchyViewController.view.accessibilityIdentifier = @"DeepHierarchyPage";
+
+  NSInteger depth = 70;
+  // A plain UILabel sibling, not part of the nested chain below, so the
+  // fixture stays recognizable to a human glancing at the simulator instead
+  // of showing a blank white screen.
+  UILabel *titleLabel = [[UILabel alloc] initWithFrame:CGRectMake(20, 60, CGRectGetWidth(UIScreen.mainScreen.bounds) - 40, 60)];
+  titleLabel.text = [NSString stringWithFormat:@"Deep Hierarchy\n%ld nested elements", (long)depth];
+  titleLabel.numberOfLines = 2;
+  titleLabel.textAlignment = NSTextAlignmentCenter;
+  titleLabel.font = [UIFont systemFontOfSize:20];
+  [deepHierarchyViewController.view addSubview:titleLabel];
+
+  UIView *parent = deepHierarchyViewController.view;
+  for (NSInteger i = 0; i < depth; i++) {
+    UIView *view = [[UIView alloc] initWithFrame:CGRectMake(0, 0, 1, 1)];
+    view.accessibilityIdentifier = [NSString stringWithFormat:@"view_%ld", (long)i];
+    view.accessibilityLabel = [NSString stringWithFormat:@"View %ld", (long)i];
+    [parent addSubview:view];
+    parent = view;
+  }
+
+  [self.navigationController pushViewController:deepHierarchyViewController animated:NO];
+}
+
 - (void)viewDidLayoutSubviews
 {
   [super viewDidLayoutSubviews];

WebDriverAgentTests/IntegrationApp/Resources/Base.lproj/Main.storyboard

@@ -74,6 +74,16 @@
                                     <segue destination="XaE-eF-eIt" kind="show" id="q3u-B9-6R6"/>
                                 </connections>
                             </button>
+                            <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="dHi-01-Btn">
+                                <rect key="frame" x="150" y="380" width="114" height="30"/>
+                                <constraints>
+                                    <constraint firstAttribute="height" constant="30" id="dHi-02-Hgt"/>
+                                </constraints>
+                                <state key="normal" title="DeepHierarchy"/>
+                                <connections>
+                                    <action selector="goToDeepHierarchy:" destination="BYZ-38-t0r" eventType="touchUpInside" id="dHi-03-Act"/>
+                                </connections>
+                            </button>
                         </subviews>
                         <color key="backgroundColor" red="1" green="1" blue="1" alpha="1" colorSpace="custom" customColorSpace="sRGB"/>
                         <accessibility key="accessibilityConfiguration" label="MainView"/>
@@ -90,6 +100,8 @@
                             <constraint firstItem="uiD-R4-b34" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="nFx-Xr-rGC"/>
                             <constraint firstItem="S56-6U-3gG" firstAttribute="top" secondItem="M2N-Yn-ytb" secondAttribute="bottom" constant="8" id="rMW-LW-ejO"/>
                             <constraint firstItem="M2N-Yn-ytb" firstAttribute="top" secondItem="YgP-SF-TkS" secondAttribute="bottom" constant="8" id="tMN-gl-smg"/>
+                            <constraint firstItem="dHi-01-Btn" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="dHi-04-Cnt"/>
+                            <constraint firstItem="dHi-01-Btn" firstAttribute="top" secondItem="1Ht-AF-MGW" secondAttribute="bottom" constant="8" id="dHi-05-Top"/>
                         </constraints>
                     </view>
                     <navigationItem key="navigationItem" id="dmu-Fe-aoT"/>

WebDriverAgentTests/IntegrationTests/FBIntegrationTestCase.h

@@ -14,6 +14,13 @@ extern NSString *const FBShowAlertForceTouchButtonName;
 extern NSString *const FBTouchesCountLabelIdentifier;
 extern NSString *const FBTapsCountLabelIdentifier;
 
+/**
+ Labels of the buttons on the integration app's main page, in their on-screen
+ (top to bottom) order. Update this in one place - WebDriverAgentTests/IntegrationApp/Resources/Base.lproj/Main.storyboard -
+ and here, rather than duplicating the list/count across individual tests.
+ */
+extern NSArray<NSString *> *const FBMainViewButtonLabels;
+
 /**
  XCTestCase helper class used for integration tests
  */
@@ -62,6 +69,15 @@ extern NSString *const FBTapsCountLabelIdentifier;
  */
 - (void)goToScrollPageWithCells:(BOOL)showCells;
 
+/**
+ Navigates integration app to a page containing a 70-level-deep chain of
+ nested elements (otherElements[@"view_0"]...otherElements[@"view_69"]),
+ each nested directly inside the previous one. Intended as a fixture for
+ performance testing of element lookups (e.g. class chain locators) against
+ a deep accessibility tree.
+ */
+- (void)goToDeepHierarchyPage;
+
 /**
  Verifies no alerts are present on the page.
  If an alert exists then it is going to be dismissed.

WebDriverAgentTests/IntegrationTests/FBIntegrationTestCase.m

@@ -26,6 +26,14 @@
 NSString *const FBShowAlertForceTouchButtonName = @"Create Alert (Force Touch)";
 NSString *const FBTouchesCountLabelIdentifier = @"numberOfTouchesLabel";
 NSString *const FBTapsCountLabelIdentifier = @"numberOfTapsLabel";
+NSArray<NSString *> *const FBMainViewButtonLabels = @[
+  @"Alerts",
+  @"Deadlock app",
+  @"Attributes",
+  @"Scrolling",
+  @"Touch",
+  @"DeepHierarchy",
+];
 
 @interface FBIntegrationTestCase ()
 @property (nonatomic, strong) XCUIApplication *testedApplication;
@@ -129,6 +137,17 @@ - (void)goToScrollPageWithCells:(BOOL)showCells
   FBAssertWaitTillBecomesTrue(self.testedApplication.staticTexts[@"3"].fb_isVisible);
 }
 
+- (void)goToDeepHierarchyPage
+{
+  [self.testedApplication.buttons[@"DeepHierarchy"] tap];
+  [self.testedApplication fb_waitUntilStable];
+  // Not fb_isVisible: that runs a native visibility computation which is
+  // pathologically slow to resolve for a view nested 70 levels deep at a
+  // 1x1 frame. Existence in the accessibility tree is all this fixture
+  // actually needs to confirm navigation succeeded.
+  FBAssertWaitTillBecomesTrue(self.testedApplication.otherElements[@"view_0"].exists);
+}
+
 - (void)clearAlert
 {
   [self.testedApplication fb_waitUntilStable];

WebDriverAgentTests/IntegrationTests/FBXPathIntegrationTests.m

@@ -130,7 +130,7 @@ - (void)testSingleDescendantXMLRepresentationWithoutAttributes
 - (void)testFindMatchesInElement
 {
   NSArray<id<FBXCElementSnapshot>> *matchingSnapshots = [FBXPath matchesWithRootElement:self.testedApplication forQuery:@"//XCUIElementTypeButton"];
-  XCTAssertEqual([matchingSnapshots count], 5);
+  XCTAssertEqual([matchingSnapshots count], FBMainViewButtonLabels.count);
   for (id<FBXCElementSnapshot> element in matchingSnapshots) {
     XCTAssertTrue([[FBXCElementSnapshotWrapper ensureWrapped:element].wdType isEqualToString:@"XCUIElementTypeButton"]);
   }
@@ -174,7 +174,7 @@ - (void)testFindMatchesWithoutContextScopeLimit
 - (void)testFindMatchesInElementWithDotNotation
 {
   NSArray<id<FBXCElementSnapshot>> *matchingSnapshots = [FBXPath matchesWithRootElement:self.testedApplication forQuery:@".//XCUIElementTypeButton"];
-  XCTAssertEqual([matchingSnapshots count], 5);
+  XCTAssertEqual([matchingSnapshots count], FBMainViewButtonLabels.count);
   for (id<FBXCElementSnapshot> element in matchingSnapshots) {
     XCTAssertTrue([[FBXCElementSnapshotWrapper ensureWrapped:element].wdType isEqualToString:@"XCUIElementTypeButton"]);
   }
@@ -232,7 +232,7 @@ - (void)testFindMatchesWithExtensionFunctionsNoMatches
 - (void)testFindMultipleMatchesWithMatchesFunction
 {
   [self assertXPathQuery:@"//XCUIElementTypeButton[matches(@label, '.*')]"
-       findsButtonLabels:@[@"Alerts", @"Deadlock app", @"Attributes", @"Scrolling", @"Touch"]];
+       findsButtonLabels:FBMainViewButtonLabels];
 }
 
 - (void)testInvalidXPathExtensionFunctionViaElementLookup

WebDriverAgentTests/IntegrationTests/XCElementSnapshotHelperTests.m

@@ -36,13 +36,7 @@ - (void)setUp
 
 - (void)testDescendantsMatchingType
 {
-  NSSet<NSString *> *expectedLabels = [NSSet setWithArray:@[
-    @"Alerts",
-    @"Attributes",
-    @"Scrolling",
-    @"Deadlock app",
-    @"Touch",
-  ]];
+  NSSet<NSString *> *expectedLabels = [NSSet setWithArray:FBMainViewButtonLabels];
   NSArray<id<FBXCElementSnapshot>> *matchingSnapshots = [[FBXCElementSnapshotWrapper ensureWrapped:
                                                           [self.testedView fb_customSnapshot]]
                                                          fb_descendantsMatchingType:XCUIElementTypeButton];