Optimize AppScreenChecker performance when parsing dumpsys output and matching view hierarchy paths.

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)

android_env/components/app_screen_checker_test.py

@@ -16,15 +16,23 @@
 """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
 from android_env.components import errors
 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()