chore: Optimize stable instance retrieval

WebDriverAgentLib/Categories/XCUIElement+FBFind.m

@@ -110,7 +110,6 @@ @implementation XCUIElement (FBFind)
     matchingSnapshots = @[snapshot];
   }
   return [self fb_filterDescendantsWithSnapshots:matchingSnapshots
-                                         selfUID:self.fb_uid
                                     onlyChildren:NO];
 }
 

WebDriverAgentLib/Categories/XCUIElement+FBPickerWheel.m

@@ -12,6 +12,7 @@
 #import "FBRunLoopSpinner.h"
 #import "FBXCElementSnapshot.h"
 #import "FBXCodeCompatibility.h"
+#import "XCUIElement+FBUID.h"
 #import "XCUICoordinate.h"
 #import "XCUIElement+FBCaching.h"
 #import "XCUIElement+FBResolve.h"
@@ -33,7 +34,7 @@ - (BOOL)fb_scrollWithOffset:(CGFloat)relativeHeightOffset error:(NSError **)erro
   // Fetching stable instance of an element allows it to be bounded to the
   // unique element identifier (UID), so it could be found next time even if its
   // id is different from the initial one. See https://github.com/appium/appium/issues/17569
-  XCUIElement *stableInstance = self.fb_stableInstance;
+  XCUIElement *stableInstance = [self fb_stableInstanceWithUid:[FBXCElementSnapshotWrapper wdUIDWithSnapshot:snapshot]];
   [endCoord tap];
   return [[[[FBRunLoopSpinner new]
      timeout:VALUE_CHANGE_TIMEOUT]

WebDriverAgentLib/Categories/XCUIElement+FBResolve.h

@@ -27,10 +27,11 @@ NS_ASSUME_NONNULL_BEGIN
  Although, if the cached element instance is the one returned by this API call then the same element
  is going to be matched and no staleness exception will be thrown.
 
+ @param uid Element UUID
  @return Either the same element instance if `fb_isResolvedNatively` was set to NO (usually the cache for elements
  matched by xpath locators) or the stable instance of the self element based on the query by element's UUID.
  */
-- (XCUIElement *)fb_stableInstance;
+- (XCUIElement *)fb_stableInstanceWithUid:(NSString *)uid;
 
 @end
 

WebDriverAgentLib/Categories/XCUIElement+FBResolve.m

@@ -32,23 +32,19 @@ - (NSNumber *)fb_isResolvedNatively
   return nil == result ? @YES : result;
 }
 
-- (XCUIElement *)fb_stableInstance
+- (XCUIElement *)fb_stableInstanceWithUid:(NSString *)uid
 {
-  if (![self.fb_isResolvedNatively boolValue]) {
+  if (nil == uid || ![self.fb_isResolvedNatively boolValue] || [self isKindOfClass:XCUIApplication.class]) {
     return self;
   }
-
-  XCUIElementQuery *query = [self isKindOfClass:XCUIApplication.class]
-    ? self.application.fb_query
-    : [self.application.fb_query descendantsMatchingType:XCUIElementTypeAny];
-  NSString *uid = nil == self.fb_cachedSnapshot
-    ? self.fb_uid
-    : [FBXCElementSnapshotWrapper wdUIDWithSnapshot:(id)self.fb_cachedSnapshot];
-  if (nil == uid) {
-    return self;
+  NSPredicate *predicate = [NSPredicate predicateWithFormat:@"%K = %@", FBStringify(FBXCElementSnapshotWrapper, fb_uid), uid];
+  XCUIElementQuery *query = [self.application.fb_query descendantsMatchingType:XCUIElementTypeAny];
+  XCUIElement *result = [query matchingPredicate:predicate].allElementsBoundByIndex.firstObject;
+  if (nil != result) {
+    result.fb_isResolvedNatively = @NO;
+    return result;
   }
-  NSPredicate *predicate = [NSPredicate predicateWithFormat:@"%K = %@",FBStringify(FBXCElementSnapshotWrapper, fb_uid), uid];
-  return [query matchingPredicate:predicate].allElementsBoundByIndex.firstObject ?: self;
+  return self;
 }
 
 @end

WebDriverAgentLib/Categories/XCUIElement+FBUtilities.h

@@ -41,14 +41,11 @@ NS_ASSUME_NONNULL_BEGIN
  Filters elements by matching them to snapshots from the corresponding array
 
  @param snapshots Array of snapshots to be matched with
- @param selfUID Optionally the unique identifier of the current element.
- Providing it as an argument improves the performance of the method.
  @param onlyChildren Whether to only look for direct element children
 
  @return Array of filtered elements, which have matches in snapshots array
  */
 - (NSArray<XCUIElement *> *)fb_filterDescendantsWithSnapshots:(NSArray<id<FBXCElementSnapshot>> *)snapshots
-                                                      selfUID:(nullable NSString *)selfUID
                                                  onlyChildren:(BOOL)onlyChildren;
 
 /**

WebDriverAgentLib/Categories/XCUIElement+FBUtilities.m

@@ -71,7 +71,6 @@ @implementation XCUIElement (FBUtilities)
 }
 
 - (NSArray<XCUIElement *> *)fb_filterDescendantsWithSnapshots:(NSArray<id<FBXCElementSnapshot>> *)snapshots
-                                                      selfUID:(NSString *)selfUID
                                                  onlyChildren:(BOOL)onlyChildren
 {
   if (0 == snapshots.count) {
@@ -85,13 +84,11 @@ @implementation XCUIElement (FBUtilities)
     }
   }
   NSMutableArray<XCUIElement *> *matchedElements = [NSMutableArray array];
-  NSString *uid = selfUID;
-  if (nil == uid) {
-    uid = self.fb_uid;
-  }
+  NSString *uid = nil == self.lastSnapshot
+    ? self.fb_uid
+    : [FBXCElementSnapshotWrapper wdUIDWithSnapshot:self.lastSnapshot];
   if (nil != uid && [matchedIds containsObject:uid]) {
-    XCUIElement *stableSelf = self.fb_stableInstance;
-    stableSelf.fb_isResolvedNatively = @NO;
+    XCUIElement *stableSelf = [self fb_stableInstanceWithUid:uid];
     if (1 == snapshots.count) {
       return @[stableSelf];
     }

WebDriverAgentLib/Commands/FBElementCommands.m

@@ -264,7 +264,9 @@ + (NSArray *)routes
   if (focusedElement != nil) {
     FBElementCache *elementCache = request.session.elementCache;
     BOOL useNativeCachingStrategy = request.session.useNativeCachingStrategy;
-    NSString *focusedUUID = [elementCache storeElement:(useNativeCachingStrategy ? focusedElement : focusedElement.fb_stableInstance)];
+    NSString *focusedUUID = [elementCache storeElement:(useNativeCachingStrategy
+                                                        ? focusedElement
+                                                        : [focusedElement fb_stableInstanceWithUid:focusedElement.fb_uid])];
     if (focusedUUID && [focusedUUID isEqualToString:(id)request.parameters[@"uuid"]]) {
       isFocused = YES;
     }

WebDriverAgentLib/Commands/FBFindElementCommands.m

@@ -87,7 +87,6 @@ + (NSArray *)routes
       && [FBXCElementSnapshotWrapper ensureWrapped:shot].wdVisible;
   }];
   NSArray *cells = [element fb_filterDescendantsWithSnapshots:visibleCellSnapshots
-                                                      selfUID:[FBXCElementSnapshotWrapper wdUIDWithSnapshot:snapshot]
                                                  onlyChildren:NO];
   return FBResponseWithCachedElements(cells, request.session.elementCache, FBConfiguration.shouldUseCompactResponses);
 }

WebDriverAgentLib/Routing/FBResponsePayload.m

@@ -34,23 +34,35 @@
   return FBResponseWithStatus([FBCommandStatus okWithValue:object]);
 }
 
-id<FBResponsePayload> FBResponseWithCachedElement(XCUIElement *element, FBElementCache *elementCache, BOOL compact)
+XCUIElement *maybeStable(XCUIElement *element)
 {
   BOOL useNativeCachingStrategy = nil == FBSession.activeSession
     ? YES
     : FBSession.activeSession.useNativeCachingStrategy;
-  [elementCache storeElement:(useNativeCachingStrategy ? element : element.fb_stableInstance)];
+  if (useNativeCachingStrategy) {
+    return element;
+  }
+
+  XCUIElement *result = element;
+  id<FBXCElementSnapshot> snapshot = element.lastSnapshot ?: [element fb_cachedSnapshot] ?: [element fb_takeSnapshot:NO];
+  NSString *uid = [FBXCElementSnapshotWrapper wdUIDWithSnapshot:snapshot];
+  if (nil != uid) {
+    result = [element fb_stableInstanceWithUid:uid];
+  }
+  return result;
+}
+
+id<FBResponsePayload> FBResponseWithCachedElement(XCUIElement *element, FBElementCache *elementCache, BOOL compact)
+{
+  [elementCache storeElement:maybeStable(element)];
   return FBResponseWithStatus([FBCommandStatus okWithValue:FBDictionaryResponseWithElement(element, compact)]);
 }
 
 id<FBResponsePayload> FBResponseWithCachedElements(NSArray<XCUIElement *> *elements, FBElementCache *elementCache, BOOL compact)
 {
   NSMutableArray *elementsResponse = [NSMutableArray array];
-  BOOL useNativeCachingStrategy = nil == FBSession.activeSession
-    ? YES
-    : FBSession.activeSession.useNativeCachingStrategy;
   for (XCUIElement *element in elements) {
-    [elementCache storeElement:(useNativeCachingStrategy ? element : element.fb_stableInstance)];
+    [elementCache storeElement:maybeStable(element)];
     [elementsResponse addObject:FBDictionaryResponseWithElement(element, compact)];
   }
   return FBResponseWithStatus([FBCommandStatus okWithValue:elementsResponse]);

WebDriverAgentTests/IntegrationTests/FBXPathIntegrationTests.m

@@ -43,9 +43,7 @@ - (void)setUp
 - (id<FBXCElementSnapshot>)destinationSnapshot
 {
   XCUIElement *matchingElement = self.testedView.buttons.allElementsBoundByIndex.firstObject;
-  FBAssertWaitTillBecomesTrue(nil != [matchingElement fb_takeSnapshot:YES]);
-
-  id<FBXCElementSnapshot> snapshot = matchingElement.lastSnapshot;
+  id<FBXCElementSnapshot> snapshot = [matchingElement fb_takeSnapshot:YES];
   // Over iOS13, snapshot returns a child.
   // The purpose of here is return a single element to replace children with an empty array for testing.
   snapshot.children = @[];

WebDriverAgentTests/IntegrationTests/XCUIElementFBFindTests.m

@@ -15,6 +15,7 @@
 #import "FBTestMacros.h"
 #import "XCUIElement.h"
 #import "XCUIElement+FBFind.h"
+#import "XCUIElement+FBUID.h"
 #import "FBXCElementSnapshotWrapper+Helpers.h"
 #import "XCUIElement+FBIsVisible.h"
 #import "XCUIElement+FBClassChain.h"
@@ -93,7 +94,10 @@ - (void)testStableInstance
   NSArray<XCUIElement *> *matchingSnapshots = [self.testedView fb_descendantsMatchingIdentifier:@"Alerts"
                                                                     shouldReturnAfterFirstMatch:YES];
   XCTAssertEqual(matchingSnapshots.count, 1);
-  for (XCUIElement *el in @[matchingSnapshots.lastObject, matchingSnapshots.lastObject.fb_stableInstance]) {
+  for (XCUIElement *el in @[
+    matchingSnapshots.lastObject,
+    [matchingSnapshots.lastObject fb_stableInstanceWithUid:[matchingSnapshots.lastObject fb_uid]]
+  ]) {
     XCTAssertEqual(el.elementType, XCUIElementTypeButton);
     XCTAssertEqualObjects(el.label, @"Alerts");
   }

WebDriverAgentTests/IntegrationTests/XCUIElementHelperIntegrationTests.m

@@ -44,7 +44,6 @@ - (void)testDescendantsFiltering
   [buttonSnapshots addObject:[buttons.firstObject fb_takeSnapshot:YES]];
 
   NSArray<XCUIElement *> *result = [self.testedApplication fb_filterDescendantsWithSnapshots:buttonSnapshots
-                                                                                     selfUID:nil
                                                                                 onlyChildren:NO];
   XCTAssertEqual(1, result.count);
   XCTAssertEqual([result.firstObject elementType], XCUIElementTypeButton);