From 7333b68f7f47454ec79430373bcd2343355d5417 Mon Sep 17 00:00:00 2001 From: caipira113 Date: Wed, 5 Aug 2026 16:19:41 +0900 Subject: [PATCH] perf(locator): stop searching once the caller's matches are found Resolution collected every match before discarding all but the ones the caller asked for. On an application that materialises accessibility children on demand this dominates runtime: with a large document open, `snapshot --depth 4` on Word took 139s because the walk continued through the document canvas long after the answer was known. Each step now receives the number of candidates the following steps can actually consume, computed backwards from the caller's request. `nth=N` needs N+1, `first` needs one, and `last`, a negative `nth`, or a following search step still need the full set, so those keep walking. `resolve_one` asks for two: one to act on, one to prove ambiguity. Measured on Word with a 74k-character document open: snapshot --depth 4 139.55s -> 0.14s snapshot splitgroup 94.43s -> 0.11s get role splitgroup >200s -> 1.98s Results are unchanged. The one visible difference is the ambiguity error, which now reports "at least N" rather than a total: proving ambiguity no longer requires counting every match, and paying a full traversal for a number the caller cannot act on is what made the error slower to produce than the successful path. --- src/accessibility.rs | 243 +++++++++++++++++++++++++++++++++++++++++-- src/actions.rs | 7 +- src/error.rs | 32 +++++- src/locator.rs | 121 +++++++++++++++++---- src/main.rs | 10 +- 5 files changed, 376 insertions(+), 37 deletions(-) diff --git a/src/accessibility.rs b/src/accessibility.rs index 040b729..f625cfa 100644 --- a/src/accessibility.rs +++ b/src/accessibility.rs @@ -475,6 +475,53 @@ pub fn find_all( results } +/// Find up to `limit` elements matching `query`, searching up to `max_depth` +/// levels. +/// +/// Traversal stops as soon as `limit` matches have been collected. On a small +/// tree this is indistinguishable from [`find_all`], but on an application that +/// materialises accessibility children on demand — Word and other document +/// editors build one element per text run when asked — walking the whole tree +/// to satisfy a query that only needs the first hit costs tens of seconds of +/// IPC. A `limit` of `usize::MAX` is exactly [`find_all`]. +pub fn find_limited( + root: &AXUIElement, + query: &AXQuery, + max_depth: usize, + limit: usize, +) -> Vec> { + let mut results = Vec::new(); + if limit == 0 { + return results; + } + find_limited_inner(root, query, max_depth, limit, &mut results); + results +} + +fn find_limited_inner( + root: &AXUIElement, + query: &AXQuery, + max_depth: usize, + limit: usize, + results: &mut Vec>, +) -> bool { + if max_depth == 0 { + return false; + } + for child in children(root) { + if query.matches(&child) { + results.push(child.clone()); + if results.len() >= limit { + return true; + } + } + if find_limited_inner(&child, query, max_depth - 1, limit, results) { + return true; + } + } + false +} + fn find_all_inner( root: &AXUIElement, query: &AXQuery, @@ -604,6 +651,20 @@ impl AXNode { .collect() } + /// Find at most `limit` elements matching a locator string. + /// + /// Callers that only need to know "one match" or "more than one match" — a + /// uniqueness check before acting on an element, for instance — should ask + /// for the smallest number that answers the question. On an application + /// that materialises accessibility children on demand, that is the + /// difference between milliseconds and tens of seconds. + pub fn locate_limited(&self, locator: &str, limit: usize) -> Vec { + resolve_locator_limited(&self.0, locator, limit) + .into_iter() + .map(AXNode::new) + .collect() + } + /// Generate a locator string that uniquely identifies `target` within this node's subtree. pub fn locator(&self, target: &AXNode) -> String { generate_locator(&self.0, &target.0) @@ -1533,6 +1594,21 @@ fn parse_locator_steps(locator: &str) -> Vec> { pub fn resolve_locator_all( root: &AXUIElement, locator: &str, +) -> Vec> { + resolve_locator_limited(root, locator, usize::MAX) +} + +/// Resolve a locator, collecting at most `limit` final matches. +/// +/// Each pipeline step is given the number of candidates the steps after it can +/// actually consume, so a search stops early when the extra matches would be +/// discarded anyway. `nth=N` needs `N + 1` candidates, `first` needs one, and +/// `last`, a negative `nth`, or a following search step all need the full set. +/// A `limit` of `usize::MAX` reproduces the unbounded behaviour exactly. +pub fn resolve_locator_limited( + root: &AXUIElement, + locator: &str, + limit: usize, ) -> Vec> { let steps = parse_locator_steps(locator); @@ -1540,15 +1616,17 @@ pub fn resolve_locator_all( return Vec::new(); } + let budgets = step_budgets(&steps, limit); + // First step let mut current = match &steps[0] { - LocatorStep::Descendant(sel) => collect_matching(root, sel, 50), + LocatorStep::Descendant(sel) => collect_matching_limited(root, sel, 50, budgets[0]), LocatorStep::DirectChild(sel) => collect_direct_children_matching(root, sel), }; // Pipeline: each subsequent step searches within current results - for step in &steps[1..] { - current = apply_step_typed(¤t, step); + for (step, budget) in steps[1..].iter().zip(budgets[1..].iter().copied()) { + current = apply_step_typed(¤t, step, budget); if current.is_empty() { break; } @@ -1556,17 +1634,57 @@ pub fn resolve_locator_all( current } +/// Work out how many candidates each step must produce. +/// +/// Walking backwards: a selection step tells the step before it how few +/// candidates are actually needed, while any step that has to see the whole set +/// resets the budget to unbounded. +fn step_budgets(steps: &[LocatorStep<'_>], final_budget: usize) -> Vec { + let mut budgets = vec![usize::MAX; steps.len()]; + let mut budget = final_budget; + for (i, step) in steps.iter().enumerate().rev() { + budgets[i] = budget; + budget = match step { + LocatorStep::Descendant(sel) | LocatorStep::DirectChild(sel) => { + if let Some(n_str) = sel.strip_prefix("nth=") { + match n_str.parse::() { + // A negative index counts from the end, so every + // candidate has to be collected first. + Ok(n) if n >= 0 => (n as usize).saturating_add(1), + _ => usize::MAX, + } + } else if *sel == "first" { + 1 + } else { + // "last" and ordinary selectors both consume the whole set. + usize::MAX + } + } + }; + } + budgets +} + /// Apply a typed pipeline step (descendant or direct child). -fn apply_step_typed(elements: &[CFRetained], step: &LocatorStep<'_>) -> Vec> { +fn apply_step_typed( + elements: &[CFRetained], + step: &LocatorStep<'_>, + budget: usize, +) -> Vec> { match step { - LocatorStep::Descendant(sel) => apply_step_inner(elements, sel, false), - LocatorStep::DirectChild(sel) => apply_step_inner(elements, sel, true), + LocatorStep::Descendant(sel) => apply_step_inner(elements, sel, false, budget), + LocatorStep::DirectChild(sel) => apply_step_inner(elements, sel, true, budget), } } /// Apply a single pipeline step to a set of elements. /// If `direct_only` is true, only search direct children (not descendants). -fn apply_step_inner(elements: &[CFRetained], step: &str, direct_only: bool) -> Vec> { +fn apply_step_inner( + elements: &[CFRetained], + step: &str, + direct_only: bool, + budget: usize, +) -> Vec> { // nth=N — pick Nth element from current set (supports negative index) if let Some(n_str) = step.strip_prefix("nth=") { if let Ok(n) = n_str.parse::() { @@ -1590,14 +1708,21 @@ fn apply_step_inner(elements: &[CFRetained], step: &str, direct_onl // Normal selector: search within each element let mut results = Vec::new(); for el in elements { + let remaining = budget.saturating_sub(results.len()); + if remaining == 0 { + break; + } if direct_only { for child in children(el) { if element_matches_selector(&child, step) { results.push(child); + if results.len() >= budget { + return results; + } } } } else { - collect_matching_inner(el, step, 50, &mut results); + collect_matching_inner(el, step, 50, budget, &mut results); } } results @@ -1619,9 +1744,29 @@ pub fn collect_matching( root: &AXUIElement, selector: &str, max_depth: usize, +) -> Vec> { + collect_matching_limited(root, selector, max_depth, usize::MAX) +} + +/// Collect up to `limit` elements matching a selector string (DFS). +/// +/// Traversal stops as soon as `limit` matches have been collected. This matters +/// because some applications build accessibility elements on demand: a document +/// editor materialises one element per text run when the tree is walked, so a +/// full DFS over an open document costs tens of seconds of IPC even when the +/// caller only wants the first hit. A `limit` of `usize::MAX` is exactly +/// [`collect_matching`]. +pub fn collect_matching_limited( + root: &AXUIElement, + selector: &str, + max_depth: usize, + limit: usize, ) -> Vec> { let mut results = Vec::new(); - collect_matching_inner(root, selector, max_depth, &mut results); + if limit == 0 { + return results; + } + collect_matching_inner(root, selector, max_depth, limit, &mut results); results } @@ -1629,16 +1774,23 @@ fn collect_matching_inner( root: &AXUIElement, selector: &str, depth: usize, + limit: usize, results: &mut Vec>, ) { - if depth == 0 { + if depth == 0 || results.len() >= limit { return; } for child in children(root) { if element_matches_selector(&child, selector) { results.push(child.clone()); + if results.len() >= limit { + return; + } + } + collect_matching_inner(&child, selector, depth - 1, limit, results); + if results.len() >= limit { + return; } - collect_matching_inner(&child, selector, depth - 1, results); } } @@ -1778,6 +1930,75 @@ impl std::fmt::Debug for AXNode { mod tests { use super::*; + // --- step_budgets ----------------------------------------------------- + // + // These pin the contract the search bound relies on: a step may only be + // cut short when the steps after it provably cannot use the extra + // matches. Getting this wrong changes results rather than just timing. + + #[test] + fn budget_unbounded_request_stays_unbounded() { + let steps = parse_locator_steps("button >> nth=0"); + let b = step_budgets(&steps, usize::MAX); + // Asking for every match must never bound an earlier step, otherwise + // locate_all would silently start returning fewer elements. + assert_eq!(b.last().copied(), Some(usize::MAX)); + } + + #[test] + fn budget_nth_needs_n_plus_one() { + let steps = parse_locator_steps("button >> nth=3"); + let b = step_budgets(&steps, 1); + assert_eq!(b[0], 4, "nth=3 needs four candidates to pick the fourth"); + assert_eq!(b[1], 1); + } + + #[test] + fn budget_first_needs_one() { + let steps = parse_locator_steps("button >> first"); + let b = step_budgets(&steps, 1); + assert_eq!(b[0], 1); + } + + #[test] + fn budget_last_needs_everything() { + let steps = parse_locator_steps("button >> last"); + let b = step_budgets(&steps, 1); + assert_eq!( + b[0], + usize::MAX, + "last can only be resolved once every candidate is known" + ); + } + + #[test] + fn budget_negative_nth_needs_everything() { + let steps = parse_locator_steps("button >> nth=-1"); + let b = step_budgets(&steps, 1); + assert_eq!( + b[0], + usize::MAX, + "a negative index counts from the end, so nothing may be skipped" + ); + } + + #[test] + fn budget_search_step_needs_all_preceding_candidates() { + // Each candidate of the first step is a separate search root for the + // second, so the first step must still produce all of them. + let steps = parse_locator_steps("group >> button >> nth=0"); + let b = step_budgets(&steps, 1); + assert_eq!(b[0], usize::MAX); + assert_eq!(b[1], 1); + } + + #[test] + fn budget_final_request_propagates_to_last_step() { + let steps = parse_locator_steps("button"); + assert_eq!(step_budgets(&steps, 2)[0], 2); + assert_eq!(step_budgets(&steps, 1)[0], 1); + } + #[test] fn parse_exact_match() { let ab = parse_attr_bracket(r#"title="Send""#).unwrap(); diff --git a/src/actions.rs b/src/actions.rs index 60b2c9e..4fb9aef 100644 --- a/src/actions.rs +++ b/src/actions.rs @@ -34,13 +34,18 @@ impl ExecutionContext { /// Resolve a locator to exactly one element. pub fn resolve_one(&self, locator: &str) -> Result { - let nodes = self.app.locate_all(locator); + // Two matches are enough: one to act on, one to prove ambiguity. + // The exact total is deliberately not computed — on a tree that + // materialises children on demand, counting the rest can cost tens of + // seconds to produce a number the caller cannot act on anyway. + let nodes = self.app.locate_limited(locator, 2); match nodes.len() { 0 => Err(AxError::LocatorNotFound(locator.to_string())), 1 => Ok(nodes.into_iter().next().unwrap()), n => Err(AxError::LocatorAmbiguous { locator: locator.to_string(), count: n, + at_least: true, }), } } diff --git a/src/error.rs b/src/error.rs index bfa4efb..551c5b6 100644 --- a/src/error.rs +++ b/src/error.rs @@ -15,7 +15,13 @@ pub enum AxError { /// Locator matched zero elements. LocatorNotFound(String), /// Locator matched multiple elements when exactly one was expected. - LocatorAmbiguous { locator: String, count: usize }, + /// + /// `count` is the number of matches actually collected, which may be a + /// lower bound: resolution stops early once ambiguity is proven, so it + /// does not pay to enumerate the rest. `at_least` records whether that + /// happened so the message can say "at least N" rather than claim a + /// total it did not compute. + LocatorAmbiguous { locator: String, count: usize, at_least: bool }, /// Element has zero size (not visible on screen). ElementZeroSize, /// An AX action (e.g. AXPress) failed. @@ -37,8 +43,13 @@ impl fmt::Display for AxError { Self::AppNotFound(name) => write!(f, "app not found: {name}"), Self::LocatorInvalid(msg) => write!(f, "invalid locator: {msg}"), Self::LocatorNotFound(loc) => write!(f, "locator not found: {loc}"), - Self::LocatorAmbiguous { locator, count } => { - write!(f, "locator matched {count} elements, must be unique for actions\nhint: use '{locator} >> nth=N' to select one") + Self::LocatorAmbiguous { locator, count, at_least } => { + let qty = if *at_least { + format!("at least {count}") + } else { + count.to_string() + }; + write!(f, "locator matched {qty} elements, must be unique for actions\nhint: use '{locator} >> nth=N' to select one") } Self::ElementZeroSize => write!(f, "element has zero size (not visible)"), Self::ActionFailed(action) => write!(f, "{action} failed"), @@ -98,12 +109,27 @@ mod tests { let e = AxError::LocatorAmbiguous { locator: "AXButton".to_string(), count: 3, + at_least: false, }; let s = e.to_string(); assert!(s.contains("3 elements")); assert!(s.contains("nth=N")); } + #[test] + fn display_locator_ambiguous_lower_bound() { + // Resolution stops as soon as ambiguity is proven, so the count is a + // floor rather than a total and the message must not imply otherwise. + let e = AxError::LocatorAmbiguous { + locator: "AXButton".to_string(), + count: 2, + at_least: true, + }; + let s = e.to_string(); + assert!(s.contains("at least 2 elements"), "got: {s}"); + assert!(s.contains("nth=N")); + } + #[test] fn display_timeout() { let e = AxError::Timeout { diff --git a/src/locator.rs b/src/locator.rs index 70cf3ef..3627003 100644 --- a/src/locator.rs +++ b/src/locator.rs @@ -24,7 +24,7 @@ //! let ok = btn.click(); //! ``` -use crate::accessibility::{attr_string, children, find_all, AXNode, AXQuery}; +use crate::accessibility::{attr_string, children, find_all, find_limited, AXNode, AXQuery}; use crate::error::AxError; use objc2_application_services::AXUIElement; use objc2_core_foundation::CFRetained; @@ -284,14 +284,59 @@ impl Locator { const MAX_DEPTH: usize = 30; +/// How many candidates a step must produce for the steps after it to be +/// answerable. +/// +/// Search steps are the expensive part of resolution: each one walks a subtree +/// through the accessibility IPC boundary. When the caller only wants the first +/// hit, or an `nth`/`first` step later discards everything else, collecting +/// every match first is wasted work — and on an application that materialises +/// accessibility children on demand it is the difference between milliseconds +/// and tens of seconds. +/// +/// `budget` is the number of candidates the *remaining* steps can consume: +/// +/// - `Nth(n)` needs `n + 1` candidates and drops the rest. +/// - `First` needs exactly one. +/// - `Last`, `Filter` and any further search step can consume all of them, so +/// the budget stays unbounded from that point backwards. +/// +/// The final budget is supplied by the caller: `resolve` asks for 1, +/// `resolve_one` asks for 2 so it can still detect ambiguity, and `resolve_all` +/// asks for `usize::MAX`, which reproduces the previous behaviour exactly. +fn budget_for(steps: &[LocatorStep], final_budget: usize) -> Vec { + let mut budgets = vec![usize::MAX; steps.len()]; + let mut budget = final_budget; + for (i, step) in steps.iter().enumerate().rev() { + budgets[i] = budget; + budget = match step { + // Selection steps bound what the preceding step has to produce. + LocatorStep::Nth(n) => n.saturating_add(1), + LocatorStep::First => 1, + // Last and Filter both need to see every candidate to be correct, + // and a preceding search step feeds them directly. + LocatorStep::Last | LocatorStep::Filter(_) => usize::MAX, + // A search step consumes each candidate it is given, so the step + // before it must still produce all of them. + _ => usize::MAX, + }; + } + budgets +} + impl Locator { /// Resolve all matching elements. pub fn resolve_all(&self) -> Vec { + self.resolve_with_budget(usize::MAX) + } + + fn resolve_with_budget(&self, final_budget: usize) -> Vec { let root_node = AXNode::new(self.root.clone()); let mut candidates = vec![root_node]; + let budgets = budget_for(&self.steps, final_budget); - for step in &self.steps { - candidates = apply_step(step, candidates); + for (step, budget) in self.steps.iter().zip(budgets) { + candidates = apply_step(step, candidates, budget); if candidates.is_empty() { return Vec::new(); } @@ -302,20 +347,22 @@ impl Locator { /// Resolve the first matching element. pub fn resolve(&self) -> Option { - self.resolve_all().into_iter().next() + self.resolve_with_budget(1).into_iter().next() } /// Resolve exactly one element. /// /// Returns `Err(LocatorNotFound)` if no matches, `Err(LocatorAmbiguous)` if more than one. pub fn resolve_one(&self) -> Result { - let results = self.resolve_all(); + // Two is enough: one to return, one to prove ambiguity. + let results = self.resolve_with_budget(2); match results.len() { 0 => Err(AxError::LocatorNotFound("Locator matched 0 elements".to_string())), 1 => Ok(results.into_iter().next().unwrap()), n => Err(AxError::LocatorAmbiguous { locator: format!("Locator({} steps)", self.steps.len()), count: n, + at_least: true, }), } } @@ -397,23 +444,28 @@ impl Locator { // Step evaluation // --------------------------------------------------------------------------- -fn apply_step(step: &LocatorStep, candidates: Vec) -> Vec { +fn apply_step(step: &LocatorStep, candidates: Vec, budget: usize) -> Vec { match step { // --- Factory steps: search within each candidate's subtree --- LocatorStep::Role(role) => { let q = AXQuery::new().role(role); - search_descendants(candidates, &q) + search_descendants(candidates, &q, budget) } LocatorStep::RoleWithName(role, name) => { let mut results = Vec::new(); for c in &candidates { let q = AXQuery::new().role(role); + // The name test happens after the role search, so the role + // search itself cannot be bounded by the caller's budget. let matches = find_all(&c.0, &q, MAX_DEPTH); for m in matches { let t = attr_string(&m, "AXTitle").unwrap_or_default(); let d = attr_string(&m, "AXDescription").unwrap_or_default(); if t == *name || d == *name { results.push(AXNode::new(m)); + if results.len() >= budget { + return results; + } } } } @@ -421,31 +473,37 @@ fn apply_step(step: &LocatorStep, candidates: Vec) -> Vec { } LocatorStep::Text(text) => { let q = AXQuery::new().has_text(text); - search_descendants(candidates, &q) + search_descendants(candidates, &q, budget) } LocatorStep::Title(title) => { let q = AXQuery::new().title(title); - search_descendants(candidates, &q) + search_descendants(candidates, &q, budget) } LocatorStep::Description(desc) => { let mut results = Vec::new(); for c in &candidates { - find_by_description(&c.0, desc, MAX_DEPTH, &mut results); + find_by_description(&c.0, desc, MAX_DEPTH, budget, &mut results); + if results.len() >= budget { + break; + } } results } LocatorStep::DomId(id) => { let mut results = Vec::new(); for c in &candidates { - find_by_dom_id(&c.0, id, MAX_DEPTH, &mut results); + find_by_dom_id(&c.0, id, MAX_DEPTH, budget, &mut results); + if results.len() >= budget { + break; + } } results } LocatorStep::DomClass(class) => { let q = AXQuery::new().dom_class(class); - search_descendants(candidates, &q) + search_descendants(candidates, &q, budget) } - LocatorStep::Query(q) => search_descendants(candidates, q), + LocatorStep::Query(q) => search_descendants(candidates, q, budget), // --- Filter step --- LocatorStep::Filter(criteria) => candidates @@ -460,10 +518,14 @@ fn apply_step(step: &LocatorStep, candidates: Vec) -> Vec { } } -fn search_descendants(candidates: Vec, q: &AXQuery) -> Vec { +fn search_descendants(candidates: Vec, q: &AXQuery, budget: usize) -> Vec { let mut results = Vec::new(); for c in &candidates { - let matches = find_all(&c.0, q, MAX_DEPTH); + let remaining = budget.saturating_sub(results.len()); + if remaining == 0 { + break; + } + let matches = find_limited(&c.0, q, MAX_DEPTH, remaining); results.extend(matches.into_iter().map(AXNode::new)); } results @@ -473,28 +535,47 @@ fn find_by_description( root: &AXUIElement, desc: &str, max_depth: usize, + limit: usize, results: &mut Vec, ) { - if max_depth == 0 { + if max_depth == 0 || results.len() >= limit { return; } for child in children(root) { if attr_string(&child, "AXDescription").as_deref() == Some(desc) { results.push(AXNode::new(child.clone())); + if results.len() >= limit { + return; + } + } + find_by_description(&child, desc, max_depth - 1, limit, results); + if results.len() >= limit { + return; } - find_by_description(&child, desc, max_depth - 1, results); } } -fn find_by_dom_id(root: &AXUIElement, id: &str, max_depth: usize, results: &mut Vec) { - if max_depth == 0 { +fn find_by_dom_id( + root: &AXUIElement, + id: &str, + max_depth: usize, + limit: usize, + results: &mut Vec, +) { + if max_depth == 0 || results.len() >= limit { return; } for child in children(root) { if attr_string(&child, "AXDOMIdentifier").as_deref() == Some(id) { results.push(AXNode::new(child.clone())); + if results.len() >= limit { + return; + } + } + find_by_dom_id(&child, id, max_depth - 1, limit, results); + if results.len() >= limit { + return; } - find_by_dom_id(&child, id, max_depth - 1, results); } } diff --git a/src/main.rs b/src/main.rs index de596e3..0a329a2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -847,7 +847,13 @@ fn cmd_snapshot(ctx: &ExecutionContext, locator: Option<&str>, depth: usize, all if let Some(loc) = locator { validate_locator(loc)?; - let nodes = ctx.app.locate_all(loc); + // Without --all only the first match is printed, so stop after two: + // the second is what tells the user more matches exist. + let nodes = if all { + ctx.app.locate_all(loc) + } else { + ctx.app.locate_limited(loc, 2) + }; if nodes.is_empty() { return Err(AxError::LocatorNotFound(loc.to_string())); } @@ -862,7 +868,7 @@ fn cmd_snapshot(ctx: &ExecutionContext, locator: Option<&str>, depth: usize, all let node = &nodes[0]; if nodes.len() > 1 { eprintln!( - "Matched {} elements, showing first. Use --all to see all.", + "Matched at least {} elements, showing first. Use --all to see all.", nodes.len() ); }