diff --git a/android_env/components/app_screen_checker.py b/android_env/components/app_screen_checker.py index 73f7ef0..ada7c33 100644 --- a/android_env/components/app_screen_checker.py +++ b/android_env/components/app_screen_checker.py @@ -59,18 +59,20 @@ def find_child( if not self.children: return None - try: - return next(x for x in self.children if predicate(x)) - except StopIteration: - logging.info('Failed to find child. max_levels: %i.', max_levels) - # Search children. - if max_levels: - for child in self.children: - child_result = child.find_child(predicate, max_levels - 1) - if child_result is not None: - return child_result + # Use a simple loop instead of `next(x for x in ...)` to avoid generator + # expression overhead and StopIteration exception handling. + for x in self.children: + if predicate(x): + return x - return None + # Search children if not found in direct children. + if max_levels > 0: + for child in self.children: + child_result = child.find_child(predicate, max_levels - 1) + if child_result is not None: + return child_result + + return None def __repr__(self): return self._data @@ -93,34 +95,26 @@ def build_tree_from_dumpsys_output(dumpsys_output: str) -> _DumpsysNode: Returns: _DumpsysNode The root of the tree. """ - lines = dumpsys_output.split('\n') # Split by lines. - lines = [x.rstrip(' \r') for x in lines] - lines = [x for x in lines if len(x)] # Remove empty lines. - root = _DumpsysNode('___root___') # The root of all nodes. parents_stack = [root] - for line in lines: + + # Iterate directly over splitlines() to avoid allocating intermediate lists + # for stripped/filtered lines. + for line in dumpsys_output.splitlines(): + line = line.rstrip(' \r') + if not line: + continue stripped_line = line.lstrip(' ') indent = len(line) - len(stripped_line) # Number of indent spaces. new_node = _DumpsysNode(stripped_line) # Create a node without indentation. - parent = parents_stack.pop() - if parent.data == '___root___': # The root is an exception for indentation. - parent_indent = -2 - else: - parent_indent = (len(parents_stack) - 1) * 2 - - if indent == parent_indent: # `new_node` is a sibiling. - parent = parents_stack.pop() - elif indent < parent_indent: # Indentation reduced (i.e. a block finished) - num_levels = (indent // 2) + 1 - parents_stack = parents_stack[:num_levels] - parent = parents_stack.pop() - elif indent > parent_indent: # `new_node` is a child. - pass # No need to change the current parent. + # Find parent using direct index calculation from indentation. + # We assume 2 spaces indentation. Slicing in-place avoids pop/push loop. + parent_idx = indent // 2 + del parents_stack[parent_idx + 1 :] + parent = parents_stack[-1] parent.children.append(new_node) - parents_stack.append(parent) parents_stack.append(new_node) return root @@ -155,13 +149,14 @@ def matches_path( return False current_node = view_hierarchy + # Inline the direct child search to avoid defining a nested function + # (regex_predicate) and calling `find_child` in every iteration. for i, regex in enumerate(expected_view_hierarchy_path): - - def regex_predicate(node, expr=regex): - matches = expr.match(node.data) - return matches is not None - - child = current_node.find_child(regex_predicate) + child = None + for x in current_node.children: + if regex.match(x.data) is not None: + child = x + break if child is None: logging.error('Mismatched regex (%i, %s). current_node: %s', i, regex.pattern, current_node) diff --git a/android_env/components/app_screen_checker_test.py b/android_env/components/app_screen_checker_test.py index dc7f941..a11ee0e 100644 --- a/android_env/components/app_screen_checker_test.py +++ b/android_env/components/app_screen_checker_test.py @@ -16,8 +16,10 @@ """Tests for android_env.components.app_screen_checker.""" import re +import timeit from unittest import mock +from absl import flags from absl.testing import absltest from android_env.components import adb_call_parser from android_env.components import app_screen_checker @@ -25,6 +27,12 @@ from android_env.proto import adb_pb2 from android_env.proto import task_pb2 +# Benchmarks take ~2 minutes to run, so they are disabled by default. +# Run with --test_arg=--run_benchmarks to enable. +_RUN_BENCHMARKS = flags.DEFINE_bool( + 'run_benchmarks', False, 'Whether to run microbenchmarks.' +) + def _flatten_tree( tree: app_screen_checker._DumpsysNode, flat_tree: list[str], indent: int = 2 @@ -262,5 +270,109 @@ def test_wait_for_app_screen_successful(self): self.assertLess(wait_time, timeout) +class AppScreenCheckerBenchmark(absltest.TestCase): + + def setUp(self): + super().setUp() + if not _RUN_BENCHMARKS.value: + self.skipTest('Benchmark disabled. Run with --test_arg=--run_benchmarks') + + def test_build_tree_from_dumpsys_output(self): + # A larger synthetic dumpsys output to simulate real-world scenario. + # We repeat the tree structure to make it larger. + single_tree = """ + View Hierarchy + Node_0 + Node_0_0 + Node_0_0_0 + Node_0_0_0_0 + Node_0_0_1 + Node_0_0_1_0 + Node_0_0_1_1 + Node_0_0_1_2 + Node_0_1 + Node_0_1_0 + Node_0_1_1 + Node_1 + Node_1_0 + Node_1_1 + """ + # Let's repeat it to make it ~1000 lines. + dumpsys_output = 'TASK\n ACTIVITY\n' + for i in range(50): + dumpsys_output += f' View Hierarchy {i}\n' + for line in single_tree.strip().split('\n'): + if line.strip(): + dumpsys_output += ' ' + line + '\n' + + setup = ( + 'from android_env.components import app_screen_checker; ' + f'dumpsys_output = """{dumpsys_output}"""' + ) + stmt = 'app_screen_checker.build_tree_from_dumpsys_output(dumpsys_output)' + t = timeit.Timer(stmt, setup=setup) + number = 100 + res = t.timeit(number=number) + print( + '\nbuild_tree_from_dumpsys_output (~1000 lines):' + f' {res / number * 1e3:.3f} ms per loop' + ) + + def test_matches_path(self): + # Use a large dumpsys output and match a path. + single_tree = """ + View Hierarchy + Node_0 + Node_0_0 + Node_0_0_0 + Node_0_0_0_0 + Node_0_0_1 + Node_0_0_1_0 + Node_0_0_1_1 + Node_0_0_1_2 + Node_0_1 + Node_0_1_0 + Node_0_1_1 + Node_1 + Node_1_0 + Node_1_1 + """ + dumpsys_output = 'TASK\n ACTIVITY\n' + for i in range(50): + dumpsys_output += f' View Hierarchy_{i}\n' + for line in single_tree.strip().split('\n'): + if line.strip(): + dumpsys_output += ' ' + line + '\n' + # Add target view hierarchy at the end + dumpsys_output += """ + View Hierarchy + Hirohito + Akihito + Naruhito + Aiko + Fumihito + Mako + Kako + Hisahito + Masahito + """ + + expected_path = ['^Hirohito$', 'Akihito$', 'Fumihito$', 'Kako$'] + setup = ( + 'from android_env.components import app_screen_checker; import re;' + f' dumpsys_output = """{dumpsys_output}"""; expected_path =' + f' {expected_path}; expected_view_hierarchy_path = [re.compile(regex)' + ' for regex in expected_path]' + ) + stmt = ( + 'app_screen_checker.matches_path(dumpsys_output,' + ' expected_view_hierarchy_path, max_levels=100)' + ) + t = timeit.Timer(stmt, setup=setup) + number = 100 + res = t.timeit(number=number) + print(f'\nmatches_path (~1000 lines): {res / number * 1e3:.3f} ms per loop') + + if __name__ == '__main__': absltest.main()