Refactoring to unify tree maps

For some cases, we are going to use window ID instead of task ID to
identify a11y tree. To avoid having additional map between window ID and
tree, the CL does refactoring to unify maps for multiple key types.

Bug: b:146591913
Test: ArcAccessibilityHelperBridgeTest
Change-Id: If735dbeb8d13b979809afbe9b9af2b7025ffd4cd
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1976144
Auto-Submit: Daichi Hirono <hirono@chromium.org>
Commit-Queue: Daichi Hirono <hirono@chromium.org>
Reviewed-by: Sara Kato <sarakato@chromium.org>
Cr-Commit-Position: refs/heads/master@{#726708}

chrome/browser/chromeos/arc/accessibility/arc_accessibility_helper_bridge.cc

@@ -152,6 +152,19 @@ std::string GetARGBFromPrefs(PrefService* prefs,
                             base::NumberToString(opacity / 100.0).c_str());
 }
 
+ArcAccessibilityHelperBridge::TreeKey KeyForTaskId(int32_t value) {
+  return {ArcAccessibilityHelperBridge::TreeKeyType::kTaskId, value, {}};
+}
+
+ArcAccessibilityHelperBridge::TreeKey KeyForNotification(std::string value) {
+  return {ArcAccessibilityHelperBridge::TreeKeyType::kNotificationKey, 0,
+          std::move(value)};
+}
+
+ArcAccessibilityHelperBridge::TreeKey KeyForInputMethod() {
+  return {ArcAccessibilityHelperBridge::TreeKeyType::kInputMethod, 0, {}};
+}
+
 }  // namespace
 
 arc::mojom::CaptionStylePtr GetCaptionStyleFromPrefs(PrefService* prefs) {
@@ -197,6 +210,12 @@ ArcAccessibilityHelperBridge::GetForBrowserContext(
   return ArcAccessibilityHelperBridgeFactory::GetForBrowserContext(context);
 }
 
+// static
+ArcAccessibilityHelperBridge::TreeKey
+ArcAccessibilityHelperBridge::KeyForNotification(std::string notification_key) {
+  return arc::KeyForNotification(std::move(notification_key));
+}
+
 void ArcAccessibilityHelperBridge::UpdateCaptionSettings() const {
   arc::mojom::CaptionStylePtr caption_style =
       GetCaptionStyleFromPrefs(profile_->GetPrefs());
@@ -290,7 +309,7 @@ void ArcAccessibilityHelperBridge::OnSetNativeChromeVoxArcSupportProcessed(
   DCHECK_NE(task_id, kNoTaskId);
 
   if (!enabled) {
-    task_id_to_tree_.erase(task_id);
+    trees_.erase(KeyForTaskId(task_id));
 
     exo::Surface* surface = exo::GetShellMainSurface(window);
     if (surface) {
@@ -380,7 +399,7 @@ void ArcAccessibilityHelperBridge::HandleFilterTypeAllEvent(
 
     // This bridge must receive OnNotificationStateChanged call for the
     // notification_key before this receives an accessibility event for it.
-    tree_source = GetFromNotificationKey(notification_key);
+    tree_source = GetFromKey(KeyForNotification(notification_key));
     DCHECK(tree_source);
   } else if (event_data->is_input_method_window) {
     exo::InputMethodSurface* input_method_surface =
@@ -389,15 +408,14 @@ void ArcAccessibilityHelperBridge::HandleFilterTypeAllEvent(
     if (!input_method_surface)
       return;
 
-    if (!input_method_tree_) {
-      input_method_tree_ = std::make_unique<AXTreeSourceArc>(this);
-
+    if (!trees_.count(KeyForInputMethod())) {
+      auto* tree = CreateFromKey(KeyForInputMethod());
       ui::AXTreeData tree_data;
-      input_method_tree_->GetTreeData(&tree_data);
+      tree->GetTreeData(&tree_data);
       input_method_surface->SetChildAxTreeId(tree_data.tree_id);
     }
 
-    tree_source = input_method_tree_.get();
+    tree_source = GetFromKey(KeyForInputMethod());
   } else {
     if (event_data->task_id == kNoTaskId)
       return;
@@ -410,10 +428,10 @@ void ArcAccessibilityHelperBridge::HandleFilterTypeAllEvent(
     if (task_id != event_data->task_id)
       return;
 
-    tree_source = GetFromTaskId(event_data->task_id);
+    tree_source = GetFromKey(KeyForTaskId(task_id));
 
     if (!tree_source) {
-      tree_source = CreateFromTaskId(event_data->task_id);
+      tree_source = CreateFromKey(KeyForTaskId(event_data->task_id));
 
       ui::AXTreeData tree_data;
       tree_source->GetTreeData(&tree_data);
@@ -481,11 +499,10 @@ void ArcAccessibilityHelperBridge::OnAccessibilityEvent(
 void ArcAccessibilityHelperBridge::OnNotificationStateChanged(
     const std::string& notification_key,
     arc::mojom::AccessibilityNotificationStateType state) {
+  auto key = KeyForNotification(notification_key);
   switch (state) {
     case arc::mojom::AccessibilityNotificationStateType::SURFACE_CREATED: {
-      AXTreeSourceArc* tree_source =
-          CreateFromNotificationKey(notification_key);
-
+      AXTreeSourceArc* tree_source = CreateFromKey(std::move(key));
       ui::AXTreeData tree_data;
       if (!tree_source->GetTreeData(&tree_data)) {
         NOTREACHED();
@@ -495,40 +512,26 @@ void ArcAccessibilityHelperBridge::OnNotificationStateChanged(
       break;
     }
     case arc::mojom::AccessibilityNotificationStateType::SURFACE_REMOVED:
-      notification_key_to_tree_.erase(notification_key);
+      trees_.erase(key);
       UpdateTreeIdOfNotificationSurface(notification_key,
                                         ui::AXTreeIDUnknown());
       break;
   }
 }
 
-AXTreeSourceArc* ArcAccessibilityHelperBridge::GetFromTaskId(int32_t task_id) {
-  auto tree_it = task_id_to_tree_.find(task_id);
-  if (tree_it == task_id_to_tree_.end())
+AXTreeSourceArc* ArcAccessibilityHelperBridge::GetFromKey(const TreeKey& key) {
+  auto tree_it = trees_.find(key);
+  if (tree_it == trees_.end())
     return nullptr;
 
   return tree_it->second.get();
 }
 
-AXTreeSourceArc* ArcAccessibilityHelperBridge::CreateFromTaskId(
-    int32_t task_id) {
-  task_id_to_tree_[task_id].reset(new AXTreeSourceArc(this));
-  return task_id_to_tree_[task_id].get();
-}
-
-AXTreeSourceArc* ArcAccessibilityHelperBridge::CreateFromNotificationKey(
-    const std::string& notification_key) {
-  notification_key_to_tree_[notification_key].reset(new AXTreeSourceArc(this));
-  return notification_key_to_tree_[notification_key].get();
-}
-
-AXTreeSourceArc* ArcAccessibilityHelperBridge::GetFromNotificationKey(
-    const std::string& notification_key) {
-  const auto tree_it = notification_key_to_tree_.find(notification_key);
-  if (tree_it == notification_key_to_tree_.end())
-    return nullptr;
-
-  return tree_it->second.get();
+AXTreeSourceArc* ArcAccessibilityHelperBridge::CreateFromKey(TreeKey key) {
+  auto tree = std::make_unique<AXTreeSourceArc>(this);
+  auto* tree_ptr = tree.get();
+  trees_.insert(std::make_pair(std::move(key), std::move(tree)));
+  return tree_ptr;
 }
 
 void ArcAccessibilityHelperBridge::UpdateTreeIdOfNotificationSurface(
@@ -555,21 +558,12 @@ void ArcAccessibilityHelperBridge::UpdateTreeIdOfNotificationSurface(
 
 AXTreeSourceArc* ArcAccessibilityHelperBridge::GetFromTreeId(
     ui::AXTreeID tree_id) const {
-  for (auto it = task_id_to_tree_.begin(); it != task_id_to_tree_.end(); ++it) {
+  for (auto it = trees_.begin(); it != trees_.end(); ++it) {
     ui::AXTreeData tree_data;
     it->second->GetTreeData(&tree_data);
     if (tree_data.tree_id == tree_id)
       return it->second.get();
   }
-
-  for (auto notification_it = notification_key_to_tree_.begin();
-       notification_it != notification_key_to_tree_.end(); ++notification_it) {
-    ui::AXTreeData tree_data;
-    notification_it->second->GetTreeData(&tree_data);
-    if (tree_data.tree_id == tree_id)
-      return notification_it->second.get();
-  }
-
   return nullptr;
 }
 
@@ -801,8 +795,8 @@ void ArcAccessibilityHelperBridge::UpdateWindowProperties(
   // Do a lookup for the tree source. A tree source may not exist because the
   // app isn't whitelisted Android side or no data has been received for the
   // app.
-  auto it = task_id_to_tree_.find(task_id);
-  bool use_talkback = it == task_id_to_tree_.end();
+  auto* tree = GetFromKey(KeyForTaskId(task_id));
+  bool use_talkback = !tree;
 
   window->SetProperty(aura::client::kAccessibilityTouchExplorationPassThrough,
                       use_talkback);
@@ -837,14 +831,14 @@ void ArcAccessibilityHelperBridge::OnWindowActivated(
 }
 
 void ArcAccessibilityHelperBridge::OnTaskDestroyed(int32_t task_id) {
-  task_id_to_tree_.erase(task_id);
+  trees_.erase(KeyForTaskId(task_id));
 }
 
 void ArcAccessibilityHelperBridge::OnNotificationSurfaceAdded(
     ArcNotificationSurface* surface) {
   const std::string& notification_key = surface->GetNotificationKey();
 
-  auto* const tree = GetFromNotificationKey(notification_key);
+  auto* const tree = GetFromKey(KeyForNotification(notification_key));
   if (!tree)
     return;
 
@@ -869,7 +863,7 @@ void ArcAccessibilityHelperBridge::OnNotificationSurfaceAdded(
 void ArcAccessibilityHelperBridge::OnAndroidVirtualKeyboardVisibilityChanged(
     bool visible) {
   if (!visible)
-    input_method_tree_.reset();
+    trees_.erase(KeyForInputMethod());
 }
 
 }  // namespace arc

chrome/browser/chromeos/arc/accessibility/arc_accessibility_helper_bridge.h

@@ -103,15 +103,18 @@ class ArcAccessibilityHelperBridge
   void OnNotificationSurfaceRemoved(
       ash::ArcNotificationSurface* surface) override {}
 
-  const std::map<int32_t, std::unique_ptr<AXTreeSourceArc>>&
-  task_id_to_tree_for_test() const {
-    return task_id_to_tree_;
-  }
+  enum class TreeKeyType {
+    kTaskId,
+    kNotificationKey,
+    kInputMethod,
+  };
 
-  const std::map<std::string, std::unique_ptr<AXTreeSourceArc>>&
-  notification_key_to_tree_for_test() const {
-    return notification_key_to_tree_;
-  }
+  using TreeKey = std::tuple<TreeKeyType, int32_t, std::string>;
+  using TreeMap = std::map<TreeKey, std::unique_ptr<AXTreeSourceArc>>;
+
+  static TreeKey KeyForNotification(std::string notification_key);
+
+  const TreeMap& trees_for_test() const { return trees_; }
 
   void set_filter_type_all_for_test() { use_filter_type_all_for_test_ = true; }
 
@@ -144,20 +147,15 @@ class ArcAccessibilityHelperBridge
   void HandleFilterTypeFocusEvent(mojom::AccessibilityEventDataPtr event_data);
   void HandleFilterTypeAllEvent(mojom::AccessibilityEventDataPtr event_data);
 
-  AXTreeSourceArc* GetFromTaskId(int32_t task_id);
-  AXTreeSourceArc* CreateFromTaskId(int32_t task_id);
-  AXTreeSourceArc* GetFromNotificationKey(const std::string& notification_key);
-  AXTreeSourceArc* CreateFromNotificationKey(
-      const std::string& notification_key);
+  AXTreeSourceArc* CreateFromKey(TreeKey);
+  AXTreeSourceArc* GetFromKey(const TreeKey&);
   AXTreeSourceArc* GetFromTreeId(ui::AXTreeID tree_id) const;
 
   bool activation_observer_added_ = false;
   Profile* const profile_;
   ArcBridgeService* const arc_bridge_service_;
-  std::map<int32_t, std::unique_ptr<AXTreeSourceArc>> task_id_to_tree_;
-  std::map<std::string, std::unique_ptr<AXTreeSourceArc>>
-      notification_key_to_tree_;
-  std::unique_ptr<AXTreeSourceArc> input_method_tree_;
+  TreeMap trees_;
+
   std::unique_ptr<chromeos::AccessibilityStatusSubscription>
       accessibility_status_subscription_;
   bool use_filter_type_all_for_test_ = false;

chrome/browser/chromeos/arc/accessibility/arc_accessibility_helper_bridge_unittest.cc

@@ -206,8 +206,8 @@ TEST_F(ArcAccessibilityHelperBridgeTest, TaskAndAXTreeLifecycle) {
       accessibility_helper_bridge();
   helper_bridge->set_filter_type_all_for_test();
 
-  const auto& task_id_to_tree = helper_bridge->task_id_to_tree_for_test();
-  ASSERT_EQ(0U, task_id_to_tree.size());
+  const auto& key_to_tree = helper_bridge->trees_for_test();
+  ASSERT_EQ(0U, key_to_tree.size());
 
   auto event1 = arc::mojom::AccessibilityEventData::New();
   event1->source_id = 1;
@@ -231,12 +231,12 @@ TEST_F(ArcAccessibilityHelperBridgeTest, TaskAndAXTreeLifecycle) {
 
   // There's no active window.
   helper_bridge->OnAccessibilityEvent(event1.Clone());
-  ASSERT_EQ(0U, task_id_to_tree.size());
+  ASSERT_EQ(0U, key_to_tree.size());
 
   // Let's make task 1 active by activating the window.
   helper_bridge->SetActiveWindowId(std::string("org.chromium.arc.1"));
   helper_bridge->OnAccessibilityEvent(event1.Clone());
-  ASSERT_EQ(1U, task_id_to_tree.size());
+  ASSERT_EQ(1U, key_to_tree.size());
 
   // Same package name, different task.
   auto event2 = arc::mojom::AccessibilityEventData::New();
@@ -261,12 +261,12 @@ TEST_F(ArcAccessibilityHelperBridgeTest, TaskAndAXTreeLifecycle) {
 
   // Active window is still task 1.
   helper_bridge->OnAccessibilityEvent(event2.Clone());
-  ASSERT_EQ(1U, task_id_to_tree.size());
+  ASSERT_EQ(1U, key_to_tree.size());
 
   // Now make task 2 active.
   helper_bridge->SetActiveWindowId(std::string("org.chromium.arc.2"));
   helper_bridge->OnAccessibilityEvent(event2.Clone());
-  ASSERT_EQ(2U, task_id_to_tree.size());
+  ASSERT_EQ(2U, key_to_tree.size());
 
   // Same task id, different package name.
   event2->node_data.clear();
@@ -282,13 +282,13 @@ TEST_F(ArcAccessibilityHelperBridgeTest, TaskAndAXTreeLifecycle) {
 
   // No new tasks tree mappings should have occurred.
   helper_bridge->OnAccessibilityEvent(event2.Clone());
-  ASSERT_EQ(2U, task_id_to_tree.size());
+  ASSERT_EQ(2U, key_to_tree.size());
 
   helper_bridge->OnTaskDestroyed(1);
-  ASSERT_EQ(1U, task_id_to_tree.size());
+  ASSERT_EQ(1U, key_to_tree.size());
 
   helper_bridge->OnTaskDestroyed(2);
-  ASSERT_EQ(0U, task_id_to_tree.size());
+  ASSERT_EQ(0U, key_to_tree.size());
 }
 
 TEST_F(ArcAccessibilityHelperBridgeTest, EventAnnouncement) {
@@ -326,9 +326,8 @@ TEST_F(ArcAccessibilityHelperBridgeTest, NotificationEventArriveFirst) {
       accessibility_helper_bridge();
   arc_notification_surface_manager_->AddObserver(helper_bridge);
 
-  const auto& notification_key_to_tree_ =
-      helper_bridge->notification_key_to_tree_for_test();
-  ASSERT_EQ(0U, notification_key_to_tree_.size());
+  const auto& key_to_tree_ = helper_bridge->trees_for_test();
+  ASSERT_EQ(0U, key_to_tree_.size());
 
   // mojo: notification 1 created
   helper_bridge->OnNotificationStateChanged(
@@ -340,15 +339,17 @@ TEST_F(ArcAccessibilityHelperBridgeTest, NotificationEventArriveFirst) {
   event1->node_data.push_back(arc::mojom::AccessibilityNodeInfoData::New());
   helper_bridge->OnAccessibilityEvent(event1.Clone());
 
-  EXPECT_EQ(1U, notification_key_to_tree_.size());
+  EXPECT_EQ(1U, key_to_tree_.size());
 
   // wayland: surface 1 added
   MockArcNotificationSurface test_surface(kNotificationKey);
   arc_notification_surface_manager_->AddSurface(&test_surface);
 
   // Confirm that axtree id is set to the surface.
-  auto it = notification_key_to_tree_.find(kNotificationKey);
-  EXPECT_NE(notification_key_to_tree_.end(), it);
+  auto treeKey =
+      ArcAccessibilityHelperBridge::KeyForNotification(kNotificationKey);
+  auto it = key_to_tree_.find(treeKey);
+  EXPECT_NE(key_to_tree_.end(), it);
   AXTreeSourceArc* tree = it->second.get();
   ui::AXTreeData tree_data;
   tree->GetTreeData(&tree_data);
@@ -362,7 +363,7 @@ TEST_F(ArcAccessibilityHelperBridgeTest, NotificationEventArriveFirst) {
   // Ax tree of the surface should be reset as the tree no longer exists.
   EXPECT_EQ(ui::AXTreeIDUnknown(), test_surface.GetAXTreeId());
 
-  EXPECT_EQ(0U, notification_key_to_tree_.size());
+  EXPECT_EQ(0U, key_to_tree_.size());
 
   // mojo: notification 2 created
   helper_bridge->OnNotificationStateChanged(
@@ -374,12 +375,12 @@ TEST_F(ArcAccessibilityHelperBridgeTest, NotificationEventArriveFirst) {
   event3->node_data.push_back(arc::mojom::AccessibilityNodeInfoData::New());
   helper_bridge->OnAccessibilityEvent(event3.Clone());
 
-  EXPECT_EQ(1U, notification_key_to_tree_.size());
+  EXPECT_EQ(1U, key_to_tree_.size());
 
   // Ax tree from the second event is attached to the first surface. This is
   // expected behavior.
-  auto it2 = notification_key_to_tree_.find(kNotificationKey);
-  EXPECT_NE(notification_key_to_tree_.end(), it2);
+  auto it2 = key_to_tree_.find(treeKey);
+  EXPECT_NE(key_to_tree_.end(), it2);
   AXTreeSourceArc* tree2 = it2->second.get();
   ui::AXTreeData tree_data2;
   tree2->GetTreeData(&tree_data2);
@@ -389,7 +390,7 @@ TEST_F(ArcAccessibilityHelperBridgeTest, NotificationEventArriveFirst) {
   arc_notification_surface_manager_->RemoveSurface(&test_surface);
 
   // Tree shouldn't be removed as a surface for the second one will come.
-  EXPECT_EQ(1U, notification_key_to_tree_.size());
+  EXPECT_EQ(1U, key_to_tree_.size());
 
   // wayland: surface 2 added
   MockArcNotificationSurface test_surface_2(kNotificationKey);
@@ -402,7 +403,7 @@ TEST_F(ArcAccessibilityHelperBridgeTest, NotificationEventArriveFirst) {
       kNotificationKey,
       arc::mojom::AccessibilityNotificationStateType::SURFACE_REMOVED);
 
-  EXPECT_EQ(0U, notification_key_to_tree_.size());
+  EXPECT_EQ(0U, key_to_tree_.size());
 
   // wayland: surface 2 removed
   arc_notification_surface_manager_->RemoveSurface(&test_surface_2);
@@ -419,9 +420,8 @@ TEST_F(ArcAccessibilityHelperBridgeTest, NotificationSurfaceArriveFirst) {
       accessibility_helper_bridge();
   arc_notification_surface_manager_->AddObserver(helper_bridge);
 
-  const auto& notification_key_to_tree_ =
-      helper_bridge->notification_key_to_tree_for_test();
-  ASSERT_EQ(0U, notification_key_to_tree_.size());
+  const auto& key_to_tree_ = helper_bridge->trees_for_test();
+  ASSERT_EQ(0U, key_to_tree_.size());
 
   // wayland: surface 1 added
   MockArcNotificationSurface test_surface(kNotificationKey);
@@ -440,14 +440,14 @@ TEST_F(ArcAccessibilityHelperBridgeTest, NotificationSurfaceArriveFirst) {
   event1->node_data.push_back(arc::mojom::AccessibilityNodeInfoData::New());
   helper_bridge->OnAccessibilityEvent(event1.Clone());
 
-  EXPECT_EQ(1U, notification_key_to_tree_.size());
+  EXPECT_EQ(1U, key_to_tree_.size());
 
   // mojo: notification 2 removed
   helper_bridge->OnNotificationStateChanged(
       kNotificationKey,
       arc::mojom::AccessibilityNotificationStateType::SURFACE_REMOVED);
 
-  EXPECT_EQ(0U, notification_key_to_tree_.size());
+  EXPECT_EQ(0U, key_to_tree_.size());
 }
 
 TEST_F(ArcAccessibilityHelperBridgeTest,