perf: reuse snapshot traversal within element response batches

WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.h

@@ -11,6 +11,12 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+/** Request-local scratch data. Discard after element response conversion; never
+ share across UI actions or commands. Root identity separates snapshot generations. */
+@interface FBQuerySnapshotContext : NSObject
+- (NSOrderedSet *)snapshotsForRoot:(id<FBXCElementSnapshot>)root;
+@end
+
 @interface XCUIElementQuery (FBHelpers)
 
 /**
@@ -23,6 +29,8 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (nullable id<FBXCElementSnapshot>)fb_cachedSnapshot;
 
+- (nullable id<FBXCElementSnapshot>)fb_cachedSnapshotWithContext:(nullable FBQuerySnapshotContext *)context;
+
 @end
 
 NS_ASSUME_NONNULL_END

WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.m

@@ -13,9 +13,41 @@
 #import "FBXCElementSnapshot.h"
 #import "XCTElementSetTransformer-Protocol.h"
 
+@interface FBQuerySnapshotContext ()
+@property (nonatomic) NSMapTable<id, NSOrderedSet *> *roots;
+@end
+
+@implementation FBQuerySnapshotContext
+- (instancetype)init
+{
+  if ((self = [super init])) {
+    _roots = [NSMapTable mapTableWithKeyOptions:(NSPointerFunctionsOptions)(NSPointerFunctionsStrongMemory | NSPointerFunctionsObjectPointerPersonality)
+                                  valueOptions:NSPointerFunctionsStrongMemory];
+  }
+  return self;
+}
+
+- (NSOrderedSet *)snapshotsForRoot:(id<FBXCElementSnapshot>)root
+{
+  NSOrderedSet *result = [self.roots objectForKey:root];
+  if (nil == result) {
+    NSMutableArray *snapshots = [NSMutableArray arrayWithObject:root];
+    [snapshots addObjectsFromArray:root._allDescendants];
+    result = [NSOrderedSet orderedSetWithArray:snapshots];
+    [self.roots setObject:result forKey:root];
+  }
+  return result;
+}
+@end
+
 @implementation XCUIElementQuery (FBHelpers)
 
 - (nullable id<FBXCElementSnapshot>)fb_cachedSnapshot
+{
+  return [self fb_cachedSnapshotWithContext:nil];
+}
+
+- (nullable id<FBXCElementSnapshot>)fb_cachedSnapshotWithContext:(nullable FBQuerySnapshotContext *)context
 {
   id<FBXCElementSnapshot> rootElementSnapshot = self.rootElementSnapshot;
   if (nil == rootElementSnapshot) {
@@ -29,9 +61,7 @@ @implementation XCUIElementQuery (FBHelpers)
     inputQuery = inputQuery.inputQuery;
   }
 
-  NSMutableArray *snapshots = [NSMutableArray arrayWithObject:rootElementSnapshot];
-  [snapshots addObjectsFromArray:rootElementSnapshot._allDescendants];
-  NSOrderedSet *matchingSnapshots = [NSOrderedSet orderedSetWithArray:snapshots];
+  NSOrderedSet *matchingSnapshots = [(context ?: [FBQuerySnapshotContext new]) snapshotsForRoot:rootElementSnapshot];
   @try {
     for (id<XCTElementSetTransformer> transformer in transformersChain) {
       matchingSnapshots = (NSOrderedSet *)[transformer transform:matchingSnapshots

WebDriverAgentLib/Routing/FBResponsePayload.m

@@ -16,11 +16,14 @@
 #import "FBMacros.h"
 #import "FBProtocolHelpers.h"
 #import "XCUIElementQuery.h"
+#import "XCUIElementQuery+FBHelpers.h"
 #import "XCUIElement+FBResolve.h"
 #import "XCUIElement+FBUID.h"
 #import "XCUIElement+FBUtilities.h"
 #import "XCUIElement+FBWebDriverAttributes.h"
 
+static NSDictionary *FBDictionaryResponseWithContext(XCUIElement *element, BOOL compact, FBQuerySnapshotContext *context);
+
 NSString *arbitraryAttrPrefix = @"attribute/";
 
 id<FBResponsePayload> FBResponseWithOK(void)
@@ -33,7 +36,7 @@
   return FBResponseWithStatus([FBCommandStatus okWithValue:object]);
 }
 
-XCUIElement *maybeStable(XCUIElement *element)
+static XCUIElement *FBMaybeStableWithContext(XCUIElement *element, FBQuerySnapshotContext *context)
 {
   BOOL useNativeCachingStrategy = nil == FBSession.activeSession
     ? YES
@@ -44,7 +47,7 @@
 
   XCUIElement *result = element;
   id<FBXCElementSnapshot> snapshot = element.lastSnapshot
-    ?: element.fb_cachedSnapshot
+    ?: [element.query fb_cachedSnapshotWithContext:context]
     ?: [element fb_standardSnapshot];
   NSString *uid = [FBXCElementSnapshotWrapper wdUIDWithSnapshot:snapshot];
   if (nil != uid) {
@@ -53,6 +56,11 @@
   return result;
 }
 
+XCUIElement *maybeStable(XCUIElement *element)
+{
+  return FBMaybeStableWithContext(element, nil);
+}
+
 id<FBResponsePayload> FBResponseWithCachedElement(XCUIElement *element, FBElementCache *elementCache, BOOL compact)
 {
   [elementCache storeElement:maybeStable(element)];
@@ -64,9 +72,10 @@
 id<FBResponsePayload> FBResponseWithCachedElements(NSArray<XCUIElement *> *elements, FBElementCache *elementCache, BOOL compact)
 {
   NSMutableArray *elementsResponse = [NSMutableArray array];
+  FBQuerySnapshotContext *context = [FBQuerySnapshotContext new];
   for (XCUIElement *element in elements) {
-    [elementCache storeElement:maybeStable(element)];
-    [elementsResponse addObject:FBDictionaryResponseWithElement(element, compact)];
+    [elementCache storeElement:FBMaybeStableWithContext(element, context)];
+    [elementsResponse addObject:FBDictionaryResponseWithContext(element, compact, context)];
     element.lastSnapshot = nil;
   }
   return FBResponseWithStatus([FBCommandStatus okWithValue:elementsResponse]);
@@ -105,12 +114,17 @@
                                             httpStatusCode:status.statusCode];
 }
 
-inline NSDictionary *FBDictionaryResponseWithElement(XCUIElement *element, BOOL compact)
+NSDictionary *FBDictionaryResponseWithElement(XCUIElement *element, BOOL compact)
+{
+  return FBDictionaryResponseWithContext(element, compact, nil);
+}
+
+static NSDictionary *FBDictionaryResponseWithContext(XCUIElement *element, BOOL compact, FBQuerySnapshotContext *context)
 {
   __block NSDictionary *elementResponse = nil;
   @autoreleasepool {
     id<FBXCElementSnapshot> snapshot = element.lastSnapshot
-      ?: element.fb_cachedSnapshot
+      ?: [element.query fb_cachedSnapshotWithContext:context]
       ?: [element fb_customSnapshot];
     NSDictionary *compactResult = FBToElementDict((NSString *)[FBXCElementSnapshotWrapper wdUIDWithSnapshot:snapshot]);
     if (compact) {

WebDriverAgentTests/UnitTests/FBElementUtilitiesTests.m

@@ -12,12 +12,47 @@
 #import "FBElement.h"
 #import "XCUIElementDouble.h"
 #import "FBElementUtils.h"
+#import "XCUIElementQuery+FBHelpers.h"
+#import "XCElementSnapshotDouble.h"
+
+@interface FBCountedDescendantsSnapshot : XCElementSnapshotDouble
+@property (nonatomic) NSUInteger descendantReads;
+@end
+@implementation FBCountedDescendantsSnapshot
+- (NSArray *)_allDescendants { self.descendantReads++; return @[]; }
+@end
 
 @interface FBElementUtilitiesTests : XCTestCase
 @end
 
 @implementation FBElementUtilitiesTests
 
+- (void)testSnapshotContextFlattensEachRootOnlyOnce
+{
+  FBQuerySnapshotContext *context = [FBQuerySnapshotContext new];
+  FBCountedDescendantsSnapshot *first = [FBCountedDescendantsSnapshot new];
+  FBCountedDescendantsSnapshot *next = [FBCountedDescendantsSnapshot new];
+  XCTAssertEqual([context snapshotsForRoot:(id)first], [context snapshotsForRoot:(id)first]);
+  XCTAssertEqual(first.descendantReads, 1u);
+  XCTAssertEqual([context snapshotsForRoot:(id)next].firstObject, next);
+  XCTAssertEqual(next.descendantReads, 1u);
+  FBQuerySnapshotContext *nextRequest = [FBQuerySnapshotContext new];
+  [nextRequest snapshotsForRoot:(id)first];
+  XCTAssertEqual(first.descendantReads, 2u);
+}
+
+- (void)testSnapshotContextReleasesRootsAtTheEndOfTheRequest
+{
+  __weak FBCountedDescendantsSnapshot *weakRoot;
+  @autoreleasepool {
+    FBQuerySnapshotContext *context = [FBQuerySnapshotContext new];
+    FBCountedDescendantsSnapshot *root = [FBCountedDescendantsSnapshot new];
+    weakRoot = root;
+    [context snapshotsForRoot:(id)root];
+  }
+  XCTAssertNil(weakRoot);
+}
+
 - (void)testTypesFiltering {
   NSMutableArray *elements = [NSMutableArray new];
   XCUIElementDouble *el1 = [XCUIElementDouble new];