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() ); }