From d06a6c21683e0680ce6b69999c4a5a8a1dc51248 Mon Sep 17 00:00:00 2001 From: Lahfir Date: Tue, 14 Apr 2026 00:52:01 -0700 Subject: [PATCH] refactor: unify allocate_refs across snapshot and drill-down paths allocate_refs in snapshot.rs and allocate_refs_with_root in snapshot_ref.rs were near-exact copies that differed only in whether each allocated ref carried a root_ref tag. The duplication meant any bug fix or behavior change in one path risked silently drifting from the other (bot flagged this after the first round of fixes added even more duplicated skeleton-anchor logic to allocate_refs_with_root). Move the shared logic into ref_alloc.rs behind a new RefAllocConfig struct with Option<&str> for the root_ref_id, and delete the snapshot_ref.rs copy. Both snapshot::build and snapshot_ref::run_from_ref now route through the same function with their respective configs. append_surface_refs and the existing unit tests in both modules are updated to use the shared function + config. --- crates/core/src/ref_alloc.rs | 62 +++++++++++++- crates/core/src/snapshot.rs | 124 ++++++++++----------------- crates/core/src/snapshot_ref.rs | 143 ++++++++------------------------ 3 files changed, 139 insertions(+), 190 deletions(-) diff --git a/crates/core/src/ref_alloc.rs b/crates/core/src/ref_alloc.rs index e6c7150..e55783a 100644 --- a/crates/core/src/ref_alloc.rs +++ b/crates/core/src/ref_alloc.rs @@ -1,5 +1,5 @@ use crate::node::AccessibilityNode; -use crate::refs::RefEntry; +use crate::refs::{RefEntry, RefMap}; pub(crate) const INTERACTIVE_ROLES: &[&str] = &[ "button", @@ -61,3 +61,63 @@ pub(crate) fn is_collapsible(node: &AccessibilityNode) -> bool { && node.states.is_empty() && node.children.len() == 1 } + +pub(crate) struct RefAllocConfig<'a> { + pub include_bounds: bool, + pub interactive_only: bool, + pub compact: bool, + pub pid: i32, + pub source_app: Option<&'a str>, + pub root_ref_id: Option<&'a str>, +} + +pub(crate) fn allocate_refs( + mut node: AccessibilityNode, + refmap: &mut RefMap, + config: &RefAllocConfig, +) -> AccessibilityNode { + let root_ref_owned = config.root_ref_id.map(str::to_string); + let is_interactive = INTERACTIVE_ROLES.contains(&node.role.as_str()); + + if is_interactive { + let entry = + ref_entry_from_node(&node, config.pid, config.source_app, root_ref_owned.clone()); + node.ref_id = Some(refmap.allocate(entry)); + } + + let has_label = node.name.as_deref().is_some_and(|n| !n.is_empty()) + || node.description.as_deref().is_some_and(|d| !d.is_empty()); + let is_skeleton_anchor = !is_interactive && node.children_count.is_some() && has_label; + + if is_skeleton_anchor { + let mut entry = ref_entry_from_node(&node, config.pid, config.source_app, root_ref_owned); + entry.available_actions = vec![]; + node.ref_id = Some(refmap.allocate(entry)); + } + + if !config.include_bounds { + node.bounds = None; + } + + node.children = node + .children + .into_iter() + .filter_map(|child| { + let child = allocate_refs(child, refmap, config); + if config.compact && is_collapsible(&child) { + return child.children.into_iter().next(); + } + if config.interactive_only + && child.ref_id.is_none() + && child.children.is_empty() + && child.children_count.is_none() + { + None + } else { + Some(child) + } + }) + .collect(); + + node +} diff --git a/crates/core/src/snapshot.rs b/crates/core/src/snapshot.rs index 23ad5a1..62cadd9 100644 --- a/crates/core/src/snapshot.rs +++ b/crates/core/src/snapshot.rs @@ -2,7 +2,7 @@ use crate::{ adapter::{PlatformAdapter, SnapshotSurface, TreeOptions, WindowFilter}, error::AppError, node::{AccessibilityNode, WindowInfo}, - ref_alloc::{is_collapsible, ref_entry_from_node, INTERACTIVE_ROLES}, + ref_alloc::{self, RefAllocConfig}, refs::RefMap, }; @@ -82,15 +82,15 @@ pub fn build( } else { RefMap::new() }; - let mut tree = allocate_refs( - raw_tree, - &mut refmap, - opts.include_bounds, - opts.interactive_only, - opts.compact, - window.pid, - Some(window.app.as_str()), - ); + let config = RefAllocConfig { + include_bounds: opts.include_bounds, + interactive_only: opts.interactive_only, + compact: opts.compact, + pid: window.pid, + source_app: Some(window.app.as_str()), + root_ref_id: None, + }; + let mut tree = ref_alloc::allocate_refs(raw_tree, &mut refmap, &config); crate::hints::add_structural_hints(&mut tree); @@ -131,72 +131,19 @@ pub fn append_surface_refs( }; let raw_tree = adapter.get_tree(&window, &opts).ok()?; let mut refmap = RefMap::load().ok()?; - let tree = allocate_refs(raw_tree, &mut refmap, false, true, false, pid, source_app); + let config = RefAllocConfig { + include_bounds: false, + interactive_only: true, + compact: false, + pid, + source_app, + root_ref_id: None, + }; + let tree = ref_alloc::allocate_refs(raw_tree, &mut refmap, &config); refmap.save().ok()?; Some(tree) } -fn allocate_refs( - mut node: AccessibilityNode, - refmap: &mut RefMap, - include_bounds: bool, - interactive_only: bool, - compact: bool, - window_pid: i32, - source_app: Option<&str>, -) -> AccessibilityNode { - let is_interactive = INTERACTIVE_ROLES.contains(&node.role.as_str()); - - if is_interactive { - let entry = ref_entry_from_node(&node, window_pid, source_app, None); - node.ref_id = Some(refmap.allocate(entry)); - } - - let has_label = node.name.as_deref().is_some_and(|n| !n.is_empty()) - || node.description.as_deref().is_some_and(|d| !d.is_empty()); - let is_skeleton_anchor = !is_interactive && node.children_count.is_some() && has_label; - - if is_skeleton_anchor { - let mut entry = ref_entry_from_node(&node, window_pid, source_app, None); - entry.available_actions = vec![]; - node.ref_id = Some(refmap.allocate(entry)); - } - - if !include_bounds { - node.bounds = None; - } - - node.children = node - .children - .into_iter() - .filter_map(|child| { - let child = allocate_refs( - child, - refmap, - include_bounds, - interactive_only, - compact, - window_pid, - source_app, - ); - if compact && is_collapsible(&child) { - return child.children.into_iter().next(); - } - if interactive_only - && child.ref_id.is_none() - && child.children.is_empty() - && child.children_count.is_none() - { - None - } else { - Some(child) - } - }) - .collect(); - - node -} - #[cfg(test)] mod tests { use super::*; @@ -217,14 +164,25 @@ mod tests { } } + fn run_config(compact: bool, interactive_only: bool) -> RefAllocConfig<'static> { + RefAllocConfig { + include_bounds: false, + interactive_only, + compact, + pid: 1, + source_app: Some("Test"), + root_ref_id: None, + } + } + fn run_compact(tree: AccessibilityNode) -> AccessibilityNode { let mut refmap = RefMap::new(); - allocate_refs(tree, &mut refmap, false, false, true, 1, Some("Test")) + ref_alloc::allocate_refs(tree, &mut refmap, &run_config(true, false)) } fn run_compact_interactive(tree: AccessibilityNode) -> AccessibilityNode { let mut refmap = RefMap::new(); - allocate_refs(tree, &mut refmap, false, true, true, 1, Some("Test")) + ref_alloc::allocate_refs(tree, &mut refmap, &run_config(true, true)) } #[test] @@ -332,7 +290,7 @@ mod tests { root.children = vec![container]; let mut refmap = RefMap::new(); - let result = allocate_refs(root, &mut refmap, false, false, false, 1, Some("Test")); + let result = ref_alloc::allocate_refs(root, &mut refmap, &run_config(false, false)); assert!(result.children[0].ref_id.is_some()); assert_eq!(refmap.len(), 1); @@ -350,7 +308,7 @@ mod tests { root.children = vec![container]; let mut refmap = RefMap::new(); - let result = allocate_refs(root, &mut refmap, false, false, false, 1, Some("Test")); + let result = ref_alloc::allocate_refs(root, &mut refmap, &run_config(false, false)); assert!(result.children[0].ref_id.is_none()); assert_eq!(refmap.len(), 0); @@ -365,7 +323,7 @@ mod tests { root.children = vec![container]; let mut refmap = RefMap::new(); - let result = allocate_refs(root, &mut refmap, false, false, false, 1, Some("Test")); + let result = ref_alloc::allocate_refs(root, &mut refmap, &run_config(false, false)); assert!(result.children[0].ref_id.is_some()); assert_eq!(refmap.len(), 1); @@ -380,7 +338,7 @@ mod tests { root.children = vec![container]; let mut refmap = RefMap::new(); - let result = allocate_refs(root, &mut refmap, false, true, false, 1, Some("Test")); + let result = ref_alloc::allocate_refs(root, &mut refmap, &run_config(false, true)); assert_eq!(result.children.len(), 1); assert_eq!(result.children[0].children_count, Some(10)); @@ -412,7 +370,15 @@ mod tests { root.children = vec![sidebar, described, content]; let mut refmap = RefMap::new(); - let result = allocate_refs(root, &mut refmap, false, false, false, 42, Some("Fixture")); + let config = RefAllocConfig { + include_bounds: false, + interactive_only: false, + compact: false, + pid: 42, + source_app: Some("Fixture"), + root_ref_id: None, + }; + let result = ref_alloc::allocate_refs(root, &mut refmap, &config); assert_eq!(refmap.len(), 4, "should allocate 4 refs total"); let result_value = serde_json::to_value(&result).unwrap(); diff --git a/crates/core/src/snapshot_ref.rs b/crates/core/src/snapshot_ref.rs index d61a30c..f7ff21a 100644 --- a/crates/core/src/snapshot_ref.rs +++ b/crates/core/src/snapshot_ref.rs @@ -1,21 +1,12 @@ use crate::{ adapter::{PlatformAdapter, TreeOptions}, error::AppError, - node::{AccessibilityNode, WindowInfo}, - ref_alloc::{is_collapsible, ref_entry_from_node, INTERACTIVE_ROLES}, + node::WindowInfo, + ref_alloc::{self, RefAllocConfig}, refs::RefMap, snapshot::SnapshotResult, }; -struct DrillDownConfig<'a> { - include_bounds: bool, - interactive_only: bool, - compact: bool, - pid: i32, - source_app: Option<&'a str>, - root_ref_id: &'a str, -} - pub fn run_from_ref( adapter: &dyn PlatformAdapter, opts: &TreeOptions, @@ -34,16 +25,17 @@ pub fn run_from_ref( refmap.remove_by_root_ref(root_ref_id); - let config = DrillDownConfig { + let source_app = entry.source_app.as_deref(); + let config = RefAllocConfig { include_bounds: opts.include_bounds, interactive_only: opts.interactive_only, compact: opts.compact, pid: entry.pid, - source_app: entry.source_app.as_deref(), - root_ref_id, + source_app, + root_ref_id: Some(root_ref_id), }; - let mut tree = allocate_refs_with_root(raw_tree, &mut refmap, &config); + let mut tree = ref_alloc::allocate_refs(raw_tree, &mut refmap, &config); crate::hints::add_structural_hints(&mut tree); @@ -65,65 +57,6 @@ pub fn run_from_ref( }) } -fn allocate_refs_with_root( - mut node: AccessibilityNode, - refmap: &mut RefMap, - config: &DrillDownConfig, -) -> AccessibilityNode { - let is_interactive = INTERACTIVE_ROLES.contains(&node.role.as_str()); - - if is_interactive { - let entry = ref_entry_from_node( - &node, - config.pid, - config.source_app, - Some(config.root_ref_id.to_string()), - ); - node.ref_id = Some(refmap.allocate(entry)); - } - - let has_label = node.name.as_deref().is_some_and(|n| !n.is_empty()) - || node.description.as_deref().is_some_and(|d| !d.is_empty()); - let is_skeleton_anchor = !is_interactive && node.children_count.is_some() && has_label; - - if is_skeleton_anchor { - let mut entry = ref_entry_from_node( - &node, - config.pid, - config.source_app, - Some(config.root_ref_id.to_string()), - ); - entry.available_actions = vec![]; - node.ref_id = Some(refmap.allocate(entry)); - } - - if !config.include_bounds { - node.bounds = None; - } - - node.children = node - .children - .into_iter() - .filter_map(|child| { - let child = allocate_refs_with_root(child, refmap, config); - if config.compact && is_collapsible(&child) { - return child.children.into_iter().next(); - } - if config.interactive_only - && child.ref_id.is_none() - && child.children.is_empty() - && child.children_count.is_none() - { - None - } else { - Some(child) - } - }) - .collect(); - - node -} - #[cfg(test)] mod tests { use super::*; @@ -131,6 +64,7 @@ mod tests { use crate::adapter::{NativeHandle, PermissionStatus, PlatformAdapter}; use crate::error::AdapterError; use crate::node::AccessibilityNode; + use crate::ref_alloc::ref_entry_from_node; use crate::refs::HomeGuard; use std::cell::Cell; @@ -404,6 +338,23 @@ mod tests { ); } + fn drill_config<'a>( + source_app: Option<&'a str>, + pid: i32, + root_ref_id: &'a str, + interactive_only: bool, + compact: bool, + ) -> RefAllocConfig<'a> { + RefAllocConfig { + include_bounds: false, + interactive_only, + compact, + pid, + source_app, + root_ref_id: Some(root_ref_id), + } + } + #[test] fn test_allocate_refs_with_root_tags_entries() { let mut btn = node("button"); @@ -412,15 +363,8 @@ mod tests { root.children = vec![btn]; let mut refmap = RefMap::new(); - let config = DrillDownConfig { - include_bounds: false, - interactive_only: false, - compact: false, - pid: 42, - source_app: Some("TestApp"), - root_ref_id: "@e5", - }; - let tree = allocate_refs_with_root(root, &mut refmap, &config); + let config = drill_config(Some("TestApp"), 42, "@e5", false, false); + let tree = ref_alloc::allocate_refs(root, &mut refmap, &config); assert_eq!(refmap.len(), 1); let btn_ref = tree.children[0] @@ -441,15 +385,8 @@ mod tests { root.children = vec![btn, text]; let mut refmap = RefMap::new(); - let config = DrillDownConfig { - include_bounds: false, - interactive_only: true, - compact: false, - pid: 1, - source_app: None, - root_ref_id: "@e1", - }; - let tree = allocate_refs_with_root(root, &mut refmap, &config); + let config = drill_config(None, 1, "@e1", true, false); + let tree = ref_alloc::allocate_refs(root, &mut refmap, &config); assert_eq!(tree.children.len(), 1); assert_eq!(tree.children[0].role, "button"); @@ -464,15 +401,8 @@ mod tests { root.children = vec![container]; let mut refmap = RefMap::new(); - let config = DrillDownConfig { - include_bounds: false, - interactive_only: true, - compact: false, - pid: 1, - source_app: None, - root_ref_id: "@e1", - }; - let tree = allocate_refs_with_root(root, &mut refmap, &config); + let config = drill_config(None, 1, "@e1", true, false); + let tree = ref_alloc::allocate_refs(root, &mut refmap, &config); assert_eq!(tree.children.len(), 1); assert_eq!(tree.children[0].children_count, Some(4)); @@ -488,15 +418,8 @@ mod tests { root.children = vec![wrapper]; let mut refmap = RefMap::new(); - let config = DrillDownConfig { - include_bounds: false, - interactive_only: false, - compact: true, - pid: 1, - source_app: None, - root_ref_id: "@e1", - }; - let tree = allocate_refs_with_root(root, &mut refmap, &config); + let config = drill_config(None, 1, "@e1", false, true); + let tree = ref_alloc::allocate_refs(root, &mut refmap, &config); assert_eq!(tree.children.len(), 1); assert_eq!(tree.children[0].role, "button");