AX Reverse relationship maps shouldn't keep growing memory usage.

We weren't aggressively removing entries from our map that were empty.

Bug: 826422
Change-Id: Ifc6e6c24457ea7b19c775ef64c783e0bf38b42a6
Reviewed-on: https://chromium-review.googlesource.com/1008271
Commit-Queue: Dominic Mazzoni <dmazzoni@chromium.org>
Reviewed-by: David Tseng <dtseng@chromium.org>
Cr-Commit-Position: refs/heads/master@{#551849}

ui/accessibility/ax_tree.cc

@@ -628,26 +628,37 @@ void AXTree::UpdateReverseRelations(AXNode* node, const AXNodeData& new_data) {
   const AXNodeData& old_data = node->data();
   int id = new_data.id;
   auto int_callback = [this, node, id](ax::mojom::IntAttribute attr,
-                                       const int& old_int, const int& new_int) {
+                                       const int& old_id, const int& new_id) {
     if (!IsNodeIdIntAttribute(attr))
       return;
 
-    int_reverse_relations_[attr][old_int].erase(id);
-    int_reverse_relations_[attr][new_int].insert(id);
+    auto& map = int_reverse_relations_[attr];
+    if (map.find(old_id) != map.end()) {
+      map[old_id].erase(id);
+      if (map[old_id].empty())
+        map.erase(old_id);
+    }
+    map[new_id].insert(id);
   };
   CallIfAttributeValuesChanged(old_data.int_attributes, new_data.int_attributes,
                                0, int_callback);
 
   auto intlist_callback = [this, node, id](
                               ax::mojom::IntListAttribute attr,
-                              const std::vector<int32_t>& old_intlist,
-                              const std::vector<int32_t>& new_intlist) {
+                              const std::vector<int32_t>& old_idlist,
+                              const std::vector<int32_t>& new_idlist) {
     if (!IsNodeIdIntListAttribute(attr))
       return;
 
-    for (int32_t old_id : old_intlist)
-      intlist_reverse_relations_[attr][old_id].erase(id);
-    for (int32_t new_id : new_intlist)
+    auto& map = intlist_reverse_relations_[attr];
+    for (int32_t old_id : old_idlist) {
+      if (map.find(old_id) != map.end()) {
+        map[old_id].erase(id);
+        if (map[old_id].empty())
+          map.erase(old_id);
+      }
+    }
+    for (int32_t new_id : new_idlist)
       intlist_reverse_relations_[attr][new_id].insert(id);
   };
   CallIfAttributeValuesChanged(old_data.intlist_attributes,
@@ -669,6 +680,11 @@ void AXTree::DestroySubtree(AXNode* node,
 
 void AXTree::DestroyNodeAndSubtree(AXNode* node,
                                    AXTreeUpdateState* update_state) {
+  // Clear out any reverse relations.
+  AXNodeData empty_data;
+  empty_data.id = node->id();
+  UpdateReverseRelations(node, empty_data);
+
   if (delegate_) {
     if (!update_state || !update_state->HasChangedNode(node))
       delegate_->OnNodeWillBeDeleted(this, node);

ui/accessibility/ax_tree_unittest.cc

@@ -1317,6 +1317,80 @@ TEST(AXTreeTest, IntListReverseRelations) {
   EXPECT_TRUE(base::ContainsKey(reverse_labelled_by, 1));
 }
 
+TEST(AXTreeTest, DeletingNodeUpdatesReverseRelations) {
+  AXTreeUpdate initial_state;
+  initial_state.root_id = 1;
+  initial_state.nodes.resize(3);
+  initial_state.nodes[0].id = 1;
+  initial_state.nodes[0].child_ids = {2, 3};
+  initial_state.nodes[1].id = 2;
+  initial_state.nodes[2].id = 3;
+  initial_state.nodes[2].AddIntAttribute(
+      ax::mojom::IntAttribute::kActivedescendantId, 2);
+  AXTree tree(initial_state);
+
+  auto reverse_active_descendant =
+      tree.GetReverseRelations(ax::mojom::IntAttribute::kActivedescendantId, 2);
+  ASSERT_EQ(1U, reverse_active_descendant.size());
+  EXPECT_TRUE(base::ContainsKey(reverse_active_descendant, 3));
+
+  AXTreeUpdate update;
+  update.root_id = 1;
+  update.nodes.resize(1);
+  update.nodes[0].id = 1;
+  update.nodes[0].child_ids = {2};
+  EXPECT_TRUE(tree.Unserialize(update));
+
+  reverse_active_descendant =
+      tree.GetReverseRelations(ax::mojom::IntAttribute::kActivedescendantId, 2);
+  ASSERT_EQ(0U, reverse_active_descendant.size());
+}
+
+TEST(AXTreeTest, ReverseRelationsDoNotKeepGrowing) {
+  // The number of total entries in int_reverse_relations and
+  // intlist_reverse_relations should not keep growing as the tree
+  // changes.
+
+  AXTreeUpdate initial_state;
+  initial_state.root_id = 1;
+  initial_state.nodes.resize(2);
+  initial_state.nodes[0].id = 1;
+  initial_state.nodes[0].AddIntAttribute(
+      ax::mojom::IntAttribute::kActivedescendantId, 2);
+  initial_state.nodes[0].AddIntListAttribute(
+      ax::mojom::IntListAttribute::kLabelledbyIds, {2});
+  initial_state.nodes[0].child_ids.push_back(2);
+  initial_state.nodes[1].id = 2;
+  AXTree tree(initial_state);
+
+  for (int i = 0; i < 1000; ++i) {
+    AXTreeUpdate update;
+    update.root_id = 1;
+    update.nodes.resize(2);
+    update.nodes[0].id = 1;
+    update.nodes[1].id = i + 3;
+    update.nodes[0].AddIntAttribute(
+        ax::mojom::IntAttribute::kActivedescendantId, update.nodes[1].id);
+    update.nodes[0].AddIntListAttribute(
+        ax::mojom::IntListAttribute::kLabelledbyIds, {update.nodes[1].id});
+    update.nodes[0].child_ids.push_back(update.nodes[1].id);
+    EXPECT_TRUE(tree.Unserialize(update));
+  }
+
+  size_t int_size = 0;
+  for (auto& iter : tree.int_reverse_relations())
+    int_size += iter.second.size() + 1;
+  // Note: 10 is arbitary, the idea here is just that we mutated a note
+  // with a relationship 1000 times, so if we have fewer than 10 entries
+  // in the map then clearly the map isn't growing / leaking. Same below.
+  EXPECT_LT(int_size, 10U);
+
+  size_t intlist_size = 0;
+  for (auto& iter : tree.intlist_reverse_relations())
+    intlist_size += iter.second.size() + 1;
+  EXPECT_LT(intlist_size, 10U);
+}
+
 TEST(AXTreeTest, SkipIgnoredNodes) {
   AXTreeUpdate tree_update;
   tree_update.root_id = 1;