jotpluggler: fix flashing while searching (#36128)

* modify in place instead of recreating nodes

* don't delete DataTreeNodes and simplify code

* faster: more efficient state tracking, better handler deletion

tools/jotpluggler/datatree.py

@@ -2,7 +2,6 @@ import os
 import re
 import threading
 import numpy as np
-from collections import deque
 import dearpygui.dearpygui as dpg
 
 
@@ -12,8 +11,9 @@ class DataTreeNode:
     self.full_path = full_path
     self.parent = parent
     self.children: dict[str, DataTreeNode] = {}
+    self.filtered_children: dict[str, DataTreeNode] = {}
+    self.created_children: dict[str, DataTreeNode] = {}
     self.is_leaf = False
-    self.child_count = 0
     self.is_plottable: bool | None = None
     self.ui_created = False
     self.children_ui_created = False
@@ -28,13 +28,17 @@ class DataTree:
     self.playback_manager = playback_manager
     self.current_search = ""
     self.data_tree = DataTreeNode(name="root")
-    self._build_queue: deque[tuple[DataTreeNode, str | None, str | int]] = deque()
-    self._all_paths_cache: set[str] = set()
-    self._item_handlers: set[str] = set()
+    self._build_queue: dict[str, tuple[DataTreeNode, DataTreeNode, str | int]] = {} # full_path -> (node, parent, before_tag)
+    self._current_created_paths: set[str] = set()
+    self._current_filtered_paths: set[str] = set()
+    self._path_to_node: dict[str, DataTreeNode] = {}  # full_path -> node
+    self._expanded_tags: set[str] = set()
+    self._item_handlers: dict[str, str] = {}  # ui_tag -> handler_tag
     self._avg_char_width = None
     self._queued_search = None
     self._new_data = False
     self._ui_lock = threading.RLock()
+    self._handlers_to_delete = []
     self.data_manager.add_observer(self._on_data_loaded)
 
   def create_ui(self, parent_tag: str):
@@ -48,134 +52,185 @@ class DataTree:
 
   def _on_data_loaded(self, data: dict):
     with self._ui_lock:
-      if data.get('segment_added'):
+      if data.get('segment_added') or data.get('reset'):
         self._new_data = True
-      elif data.get('reset'):
-        self._all_paths_cache = set()
-        self._new_data = True
-
-
-  def _populate_tree(self):
-    self._clear_ui()
-    self.data_tree = self._add_paths_to_tree(self._all_paths_cache, incremental=False)
-    if self.data_tree:
-      self._request_children_build(self.data_tree)
-
-  def _add_paths_to_tree(self, paths, incremental=False):
-    search_term = self.current_search.strip().lower()
-    filtered_paths = [path for path in paths if self._should_show_path(path, search_term)]
-    target_tree = self.data_tree if incremental else DataTreeNode(name="root")
-
-    if not filtered_paths:
-      return target_tree
-
-    parent_nodes_to_recheck = set()
-    for path in sorted(filtered_paths):
-      parts = path.split('/')
-      current_node = target_tree
-      current_path_prefix = ""
-
-      for i, part in enumerate(parts):
-        current_path_prefix = f"{current_path_prefix}/{part}" if current_path_prefix else part
-        if i < len(parts) - 1:
-          parent_nodes_to_recheck.add(current_node)  # for incremental changes from new data
-        if part not in current_node.children:
-          current_node.children[part] = DataTreeNode(name=part, full_path=current_path_prefix, parent=current_node)
-        current_node = current_node.children[part]
-
-      if not current_node.is_leaf:
-        current_node.is_leaf = True
-
-    self._calculate_child_counts(target_tree)
-    if incremental:
-      for p_node in parent_nodes_to_recheck:
-        p_node.children_ui_created = False
-        self._request_children_build(p_node)
-    return target_tree
 
   def update_frame(self, font):
+    if self._handlers_to_delete:  # we need to do everything in main thread, frame callbacks are flaky
+      dpg.render_dearpygui_frame()  # wait a frame to ensure queued callbacks are done
+      with self._ui_lock:
+        for handler in self._handlers_to_delete:
+          dpg.delete_item(handler)
+        self._handlers_to_delete.clear()
+
     with self._ui_lock:
       if self._avg_char_width is None and dpg.is_dearpygui_running():
         self._avg_char_width = self.calculate_avg_char_width(font)
 
       if self._new_data:
-        current_paths = set(self.data_manager.get_all_paths())
-        new_paths = current_paths - self._all_paths_cache
-        all_paths_empty = not self._all_paths_cache
-        self._all_paths_cache = current_paths
-        if all_paths_empty:
-          self._populate_tree()
-        elif new_paths:
-          self._add_paths_to_tree(new_paths, incremental=True)
+        self._process_path_change()
         self._new_data = False
         return
 
       if self._queued_search is not None:
         self.current_search = self._queued_search
-        self._all_paths_cache = set(self.data_manager.get_all_paths())
-        self._populate_tree()
+        self._process_path_change()
         self._queued_search = None
         return
 
       nodes_processed = 0
       while self._build_queue and nodes_processed < self.MAX_NODES_PER_FRAME:
-        child_node, parent_tag, before_tag = self._build_queue.popleft()
+        child_node, parent, before_tag = self._build_queue.pop(next(iter(self._build_queue)))
+        parent_tag = "data_tree_container" if parent.name == "root" else parent.ui_tag
         if not child_node.ui_created:
           if child_node.is_leaf:
             self._create_leaf_ui(child_node, parent_tag, before_tag)
           else:
             self._create_tree_node_ui(child_node, parent_tag, before_tag)
+        parent.created_children[child_node.name] = parent.children[child_node.name]
+        self._current_created_paths.add(child_node.full_path)
         nodes_processed += 1
 
-  def search_data(self):
-    self._queued_search = dpg.get_value("search_input")
+  def _process_path_change(self):
+    self._build_queue.clear()
+    search_term = self.current_search.strip().lower()
+    all_paths = set(self.data_manager.get_all_paths())
+    new_filtered_leafs = {path for path in all_paths if self._should_show_path(path, search_term)}
+    new_filtered_paths = set(new_filtered_leafs)
+    for path in new_filtered_leafs:
+      parts = path.split('/')
+      for i in range(1, len(parts)):
+        prefix = '/'.join(parts[:i])
+        new_filtered_paths.add(prefix)
+    created_paths_to_remove = self._current_created_paths - new_filtered_paths
+    filtered_paths_to_remove = self._current_filtered_paths - new_filtered_leafs
+
+    if created_paths_to_remove or filtered_paths_to_remove:
+      self._remove_paths_from_tree(created_paths_to_remove, filtered_paths_to_remove)
+      self._apply_expansion_to_tree(self.data_tree, search_term)
+
+    paths_to_add = new_filtered_leafs - self._current_created_paths
+    if paths_to_add:
+      self._add_paths_to_tree(paths_to_add)
+      self._apply_expansion_to_tree(self.data_tree, search_term)
+    self._current_filtered_paths = new_filtered_paths
+
+  def _remove_paths_from_tree(self, created_paths_to_remove, filtered_paths_to_remove):
+    for path in sorted(created_paths_to_remove, reverse=True):
+      current_node = self._path_to_node[path]
+
+      if len(current_node.created_children) == 0:
+        self._current_created_paths.remove(current_node.full_path)
+        if item_handler_tag := self._item_handlers.get(current_node.ui_tag):
+          dpg.configure_item(item_handler_tag, show=False)
+          self._handlers_to_delete.append(item_handler_tag)
+          del self._item_handlers[current_node.ui_tag]
+        dpg.delete_item(current_node.ui_tag)
+        current_node.ui_created = False
+        current_node.ui_tag = None
+        current_node.children_ui_created = False
+        del current_node.parent.created_children[current_node.name]
+        del current_node.parent.filtered_children[current_node.name]
+
+    for path in filtered_paths_to_remove:
+      parts = path.split('/')
+      current_node = self._path_to_node[path]
 
-  def _clear_ui(self):
-    for handler_tag in self._item_handlers:
-      dpg.configure_item(handler_tag, show=False)
-    dpg.set_frame_callback(dpg.get_frame_count() + 1, callback=self._delete_handlers, user_data=list(self._item_handlers))
-    self._item_handlers.clear()
+      part_array_index = -1
+      while len(current_node.filtered_children) == 0 and part_array_index >= -len(parts):
+        current_node = current_node.parent
+        if parts[part_array_index] in current_node.filtered_children:
+          del current_node.filtered_children[parts[part_array_index]]
+        part_array_index -= 1
 
-    if dpg.does_item_exist("data_tree_container"):
-      dpg.delete_item("data_tree_container", children_only=True)
+  def _add_paths_to_tree(self, paths):
+    parent_nodes_to_recheck = set()
+    for path in sorted(paths):
+      parts = path.split('/')
+      current_node = self.data_tree
+      current_path_prefix = ""
 
-    self._build_queue.clear()
+      for i, part in enumerate(parts):
+        current_path_prefix = f"{current_path_prefix}/{part}" if current_path_prefix else part
+        if i < len(parts) - 1:
+          parent_nodes_to_recheck.add(current_node)  # for incremental changes from new data
+        if part not in current_node.children:
+          current_node.children[part] = DataTreeNode(name=part, full_path=current_path_prefix, parent=current_node)
+          self._path_to_node[current_path_prefix] = current_node.children[part]
+        current_node.filtered_children[part] = current_node.children[part]
+        current_node = current_node.children[part]
 
-  def _delete_handlers(self, sender, app_data, user_data):
-    for handler in user_data:
-      dpg.delete_item(handler)
+      if not current_node.is_leaf:
+        current_node.is_leaf = True
 
-  def _calculate_child_counts(self, node: DataTreeNode):
-    if node.is_leaf:
-      node.child_count = 0
-    else:
-      node.child_count = len(node.children)
-      for child in node.children.values():
-        self._calculate_child_counts(child)
+    for p_node in parent_nodes_to_recheck:
+      p_node.children_ui_created = False
+      self._request_children_build(p_node)
+
+  def _get_node_label_and_expand(self, node: DataTreeNode, search_term: str):
+    label = f"{node.name} ({len(node.filtered_children)} fields)"
+    expand = len(search_term) > 0 and any(search_term in path for path in self._get_descendant_paths(node))
+    if expand and node.parent and len(node.parent.filtered_children) > 100 and len(node.filtered_children) > 2:
+      label += " (+)"  # symbol for large lists which aren't fully expanded for performance (only affects procLog rn)
+      expand = False
+    return label, expand
+
+  def _apply_expansion_to_tree(self, node: DataTreeNode, search_term: str):
+    if node.ui_created and not node.is_leaf and node.ui_tag and dpg.does_item_exist(node.ui_tag):
+      label, expand = self._get_node_label_and_expand(node, search_term)
+      if expand:
+        self._expanded_tags.add(node.ui_tag)
+        dpg.set_value(node.ui_tag, expand)
+      elif node.ui_tag in self._expanded_tags:  # not expanded and was expanded
+        self._expanded_tags.remove(node.ui_tag)
+        dpg.set_value(node.ui_tag, expand)
+        dpg.delete_item(node.ui_tag, children_only=True)  # delete children (not visible since collapsed)
+        self._reset_ui_state_recursive(node)
+        node.children_ui_created = False
+      dpg.set_item_label(node.ui_tag, label)
+    for child in node.created_children.values():
+      self._apply_expansion_to_tree(child, search_term)
+
+  def _reset_ui_state_recursive(self, node: DataTreeNode):
+    for child in node.created_children.values():
+      if child.ui_tag is not None:
+        if item_handler_tag := self._item_handlers.get(child.ui_tag):
+          self._handlers_to_delete.append(item_handler_tag)
+          dpg.configure_item(item_handler_tag, show=False)
+          del self._item_handlers[child.ui_tag]
+        self._reset_ui_state_recursive(child)
+        child.ui_created = False
+        child.ui_tag = None
+        child.children_ui_created = False
+        self._current_created_paths.remove(child.full_path)
+    node.created_children.clear()
+
+  def search_data(self):
+    with self._ui_lock:
+      self._queued_search = dpg.get_value("search_input")
 
   def _create_tree_node_ui(self, node: DataTreeNode, parent_tag: str, before: str | int):
-    tag = f"tree_{node.full_path}"
-    node.ui_tag = tag
-    label = f"{node.name} ({node.child_count} fields)"
+    node.ui_tag = f"tree_{node.full_path}"
     search_term = self.current_search.strip().lower()
-    expand = bool(search_term) and len(search_term) > 1 and any(search_term in path for path in self._get_descendant_paths(node))
-    if expand and node.parent and node.parent.child_count > 100 and node.child_count > 2:  # don't fully autoexpand large lists (only affects procLog rn)
-      label += " (+)"
-      expand = False
+    label, expand = self._get_node_label_and_expand(node, search_term)
+    if expand:
+      self._expanded_tags.add(node.ui_tag)
+    elif node.ui_tag in self._expanded_tags:
+      self._expanded_tags.remove(node.ui_tag)
 
     with dpg.tree_node(
-      label=label, parent=parent_tag, tag=tag, default_open=expand, open_on_arrow=True, open_on_double_click=True, before=before, delay_search=True
+      label=label, parent=parent_tag, tag=node.ui_tag, default_open=expand, open_on_arrow=True, open_on_double_click=True, before=before, delay_search=True
     ):
       with dpg.item_handler_registry() as handler_tag:
         dpg.add_item_toggled_open_handler(callback=lambda s, a, u: self._request_children_build(node))
         dpg.add_item_visible_handler(callback=lambda s, a, u: self._request_children_build(node))
-      dpg.bind_item_handler_registry(tag, handler_tag)
-      self._item_handlers.add(handler_tag)
-
+      dpg.bind_item_handler_registry(node.ui_tag, handler_tag)
+      self._item_handlers[node.ui_tag] = handler_tag
     node.ui_created = True
 
   def _create_leaf_ui(self, node: DataTreeNode, parent_tag: str, before: str | int):
-    with dpg.group(parent=parent_tag, tag=f"leaf_{node.full_path}", before=before, delay_search=True) as draggable_group:
+    node.ui_tag = f"leaf_{node.full_path}"
+    with dpg.group(parent=parent_tag, tag=node.ui_tag, before=before, delay_search=True):
       with dpg.table(header_row=False, policy=dpg.mvTable_SizingStretchProp, delay_search=True):
         dpg.add_table_column(init_width_or_weight=0.5)
         dpg.add_table_column(init_width_or_weight=0.5)
@@ -186,21 +241,21 @@ class DataTree:
     if node.is_plottable is None:
       node.is_plottable = self.data_manager.is_plottable(node.full_path)
     if node.is_plottable:
-      with dpg.drag_payload(parent=draggable_group, drag_data=node.full_path, payload_type="TIMESERIES_PAYLOAD"):
+      with dpg.drag_payload(parent=node.ui_tag, drag_data=node.full_path, payload_type="TIMESERIES_PAYLOAD"):
         dpg.add_text(f"Plot: {node.full_path}")
 
     with dpg.item_handler_registry() as handler_tag:
       dpg.add_item_visible_handler(callback=self._on_item_visible, user_data=node.full_path)
-    dpg.bind_item_handler_registry(draggable_group, handler_tag)
-    self._item_handlers.add(handler_tag)
-
+    dpg.bind_item_handler_registry(node.ui_tag, handler_tag)
+    self._item_handlers[node.ui_tag] = handler_tag
     node.ui_created = True
-    node.ui_tag = f"value_{node.full_path}"
 
   def _on_item_visible(self, sender, app_data, user_data):
     with self._ui_lock:
       path = user_data
       value_tag = f"value_{path}"
+      if not dpg.does_item_exist(value_tag):
+        return
       value_column_width = dpg.get_item_rect_size("sidebar_window")[0] // 2
       value = self.data_manager.get_value_at(path, self.playback_manager.current_time_s)
       if value is not None:
@@ -212,8 +267,7 @@ class DataTree:
   def _request_children_build(self, node: DataTreeNode):
     with self._ui_lock:
       if not node.children_ui_created and (node.name == "root" or (node.ui_tag is not None and dpg.get_value(node.ui_tag))):  # check root or node expanded
-        parent_tag = "data_tree_container" if node.name == "root" else node.ui_tag
-        sorted_children = sorted(node.children.values(), key=self._natural_sort_key)
+        sorted_children = sorted(node.filtered_children.values(), key=self._natural_sort_key)
         next_existing: list[int | str] = [0] * len(sorted_children)
         current_before_tag: int | str = 0
 
@@ -228,7 +282,7 @@ class DataTree:
         for i, child_node in enumerate(sorted_children):
           if not child_node.ui_created:
             before_tag = next_existing[i]
-            self._build_queue.append((child_node, parent_tag, before_tag))
+            self._build_queue[child_node.full_path] = (child_node, node, before_tag)
         node.children_ui_created = True
 
   def _should_show_path(self, path: str, search_term: str) -> bool:
@@ -242,7 +296,7 @@ class DataTree:
     return (node_type_key, parts)
 
   def _get_descendant_paths(self, node: DataTreeNode):
-    for child_name, child_node in node.children.items():
+    for child_name, child_node in node.filtered_children.items():
       child_name_lower = child_name.lower()
       if child_node.is_leaf:
         yield child_name_lower