From 37f1f3e67139a460da6b550aad06289434a1c5b5 Mon Sep 17 00:00:00 2001 From: Kazuaki Matsuo Date: Mon, 28 Sep 2026 21:55:39 -0700 Subject: [PATCH] perf: reuse snapshot traversal within element response batches --- .../Categories/XCUIElementQuery+FBHelpers.h | 8 +++++ .../Categories/XCUIElementQuery+FBHelpers.m | 36 +++++++++++++++++-- WebDriverAgentLib/Routing/FBResponsePayload.m | 26 ++++++++++---- .../UnitTests/FBElementUtilitiesTests.m | 35 ++++++++++++++++++ 4 files changed, 96 insertions(+), 9 deletions(-) diff --git a/WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.h b/WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.h index 8c82890741..602c072263 100644 --- a/WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.h +++ b/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)root; +@end + @interface XCUIElementQuery (FBHelpers) /** @@ -23,6 +29,8 @@ NS_ASSUME_NONNULL_BEGIN */ - (nullable id)fb_cachedSnapshot; +- (nullable id)fb_cachedSnapshotWithContext:(nullable FBQuerySnapshotContext *)context; + @end NS_ASSUME_NONNULL_END diff --git a/WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.m b/WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.m index cbefcf16f0..d34c348c69 100644 --- a/WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.m +++ b/WebDriverAgentLib/Categories/XCUIElementQuery+FBHelpers.m @@ -13,9 +13,41 @@ #import "FBXCElementSnapshot.h" #import "XCTElementSetTransformer-Protocol.h" +@interface FBQuerySnapshotContext () +@property (nonatomic) NSMapTable *roots; +@end + +@implementation FBQuerySnapshotContext +- (instancetype)init +{ + if ((self = [super init])) { + _roots = [NSMapTable mapTableWithKeyOptions:(NSPointerFunctionsOptions)(NSPointerFunctionsStrongMemory | NSPointerFunctionsObjectPointerPersonality) + valueOptions:NSPointerFunctionsStrongMemory]; + } + return self; +} + +- (NSOrderedSet *)snapshotsForRoot:(id)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)fb_cachedSnapshot +{ + return [self fb_cachedSnapshotWithContext:nil]; +} + +- (nullable id)fb_cachedSnapshotWithContext:(nullable FBQuerySnapshotContext *)context { id 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 transformer in transformersChain) { matchingSnapshots = (NSOrderedSet *)[transformer transform:matchingSnapshots diff --git a/WebDriverAgentLib/Routing/FBResponsePayload.m b/WebDriverAgentLib/Routing/FBResponsePayload.m index 5d0098a89a..a5bcb0e8f7 100644 --- a/WebDriverAgentLib/Routing/FBResponsePayload.m +++ b/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 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 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 FBResponseWithCachedElement(XCUIElement *element, FBElementCache *elementCache, BOOL compact) { [elementCache storeElement:maybeStable(element)]; @@ -64,9 +72,10 @@ id FBResponseWithCachedElements(NSArray *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 snapshot = element.lastSnapshot - ?: element.fb_cachedSnapshot + ?: [element.query fb_cachedSnapshotWithContext:context] ?: [element fb_customSnapshot]; NSDictionary *compactResult = FBToElementDict((NSString *)[FBXCElementSnapshotWrapper wdUIDWithSnapshot:snapshot]); if (compact) { diff --git a/WebDriverAgentTests/UnitTests/FBElementUtilitiesTests.m b/WebDriverAgentTests/UnitTests/FBElementUtilitiesTests.m index 5dad1c79de..044f995a73 100644 --- a/WebDriverAgentTests/UnitTests/FBElementUtilitiesTests.m +++ b/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];