[fuchsia][a11y] Implements NodeIDMapper.

This change implements a class that maps between AXNode IDs to unique
Fuchsia Node IDs.

Bug: fuchsia:63607
Test: AXTreeConverterTest.*
Change-Id: I75352da5b310112d07bd353ea4af7e23c52eac06
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2530981Reviewed-by: Sergey Ulanov <sergeyu@chromium.org>
Reviewed-by: Sharon Yang <yangsharon@chromium.org>
Commit-Queue: Lucas Radaelli <lucasradaelli@google.com>
Cr-Commit-Position: refs/heads/master@{#828450}

fuchsia/engine/browser/accessibility_bridge.cc

@@ -12,7 +12,6 @@
 #include "base/fuchsia/fuchsia_logging.h"
 #include "base/logging.h"
 #include "content/public/browser/render_widget_host_view.h"
-#include "fuchsia/engine/browser/ax_tree_converter.h"
 #include "ui/accessibility/ax_action_data.h"
 #include "ui/gfx/geometry/rect_conversions.h"
 
@@ -114,7 +113,8 @@ void AccessibilityBridge::AccessibilityEventReceived(
         pending_hit_test_callbacks_.find(event.action_request_id) !=
             pending_hit_test_callbacks_.end()) {
       fuchsia::accessibility::semantics::Hit hit;
-      hit.set_node_id(ConvertToFuchsiaNodeId(event.id, root_id_));
+      hit.set_node_id(id_mapper_->ToFuchsiaNodeID(ax_tree_.data().tree_id,
+                                                  event.id, false));
 
       // Run the pending callback with the hit.
       pending_hit_test_callbacks_[event.action_request_id](std::move(hit));
@@ -138,7 +138,14 @@ void AccessibilityBridge::OnAccessibilityActionRequested(
     return;
   }
 
-  action_data.target_node_id = ConvertToAxNodeId(node_id, root_id_);
+  auto ax_id = id_mapper_->ToAXNodeID(node_id);
+  if (!ax_id) {
+    // Fuchsia is targeting a node that does not exist.
+    callback(false);
+    return;
+  }
+
+  action_data.target_node_id = ax_id->second;
 
   if (action == fuchsia::accessibility::semantics::Action::SHOW_ON_SCREEN) {
     ui::AXNode* node = ax_tree_.GetFromId(action_data.target_node_id);
@@ -192,6 +199,7 @@ void AccessibilityBridge::OnSemanticsModeChanged(
 
   enable_semantic_updates_ = updates_enabled;
   if (updates_enabled) {
+    id_mapper_ = std::make_unique<NodeIDMapper>();
     // The first call to AccessibilityEventReceived after this call will be
     // the entire semantic tree.
     web_contents_->EnableWebContentsOnlyAccessibilityMode();
@@ -214,7 +222,8 @@ void AccessibilityBridge::OnSemanticsModeChanged(
 
 void AccessibilityBridge::OnNodeWillBeDeleted(ui::AXTree* tree,
                                               ui::AXNode* node) {
-  to_delete_.push_back(ConvertToFuchsiaNodeId(node->id(), root_id_));
+  to_delete_.push_back(
+      id_mapper_->ToFuchsiaNodeID(ax_tree_.data().tree_id, node->id(), false));
 }
 
 void AccessibilityBridge::OnAtomicUpdateFinished(
@@ -229,8 +238,9 @@ void AccessibilityBridge::OnAtomicUpdateFinished(
   // deleted are already gone, which means that all updates collected here in
   // |to_update_| are going to be executed after |to_delete_|.
   for (const ui::AXTreeObserver::Change& change : changes) {
-    to_update_.push_back(
-        AXNodeDataToSemanticNode(change.node->data(), root_id_));
+    const auto& node = change.node->data();
+    to_update_.push_back(AXNodeDataToSemanticNode(
+        node, ax_tree_.data().tree_id, node.id == root_id_, id_mapper_.get()));
   }
   // TODO(https://crbug.com/1134737): Separate updates of atomic updates and
   // don't allow all of them to be in the same commit.

fuchsia/engine/browser/accessibility_bridge.h

@@ -16,6 +16,7 @@
 #include "content/public/browser/ax_event_notification_details.h"
 #include "content/public/browser/web_contents_delegate.h"
 #include "content/public/browser/web_contents_observer.h"
+#include "fuchsia/engine/browser/ax_tree_converter.h"
 #include "fuchsia/engine/web_engine_export.h"
 #include "ui/accessibility/ax_serializable_tree.h"
 #include "ui/accessibility/ax_tree_id.h"
@@ -125,6 +126,9 @@ class WEB_ENGINE_EXPORT AccessibilityBridge
   // The root id of |ax_tree_|.
   int32_t root_id_ = 0;
 
+  // Maps node IDs from one platform to another.
+  std::unique_ptr<NodeIDMapper> id_mapper_;
+
   base::OnceClosure event_received_callback_for_test_;
 
   // If set, the scale factor for this device for use in tests.

fuchsia/engine/browser/ax_tree_converter.cc

@@ -24,11 +24,6 @@ namespace {
 // Fuchsia's default root node ID.
 constexpr uint32_t kFuchsiaRootNodeId = 0;
 
-// Remapped value for AXNode::kInvalidAXID.
-// Value is chosen to be outside the range of a 32-bit signed int, so as to not
-// conflict with other AXIDs.
-constexpr uint32_t kInvalidIdRemappedForFuchsia = 1u + INT32_MAX;
-
 fuchsia::accessibility::semantics::Attributes ConvertAttributes(
     const ui::AXNodeData& node) {
   fuchsia::accessibility::semantics::Attributes attributes;
@@ -184,11 +179,12 @@ std::vector<fuchsia::accessibility::semantics::Action> ConvertActions(
 }
 
 std::vector<uint32_t> ConvertChildIds(std::vector<int32_t> ids,
-                                      int32_t ax_root_id) {
+                                      const ui::AXTreeID& tree_id,
+                                      NodeIDMapper* id_mapper) {
   std::vector<uint32_t> child_ids;
   child_ids.reserve(ids.size());
-  for (auto i : ids) {
-    child_ids.push_back(ConvertToFuchsiaNodeId(i, ax_root_id));
+  for (const auto& i : ids) {
+    child_ids.push_back(id_mapper->ToFuchsiaNodeID(tree_id, i, false));
   }
   return child_ids;
 }
@@ -217,14 +213,18 @@ fuchsia::ui::gfx::mat4 ConvertTransform(gfx::Transform* transform) {
 
 fuchsia::accessibility::semantics::Node AXNodeDataToSemanticNode(
     const ui::AXNodeData& node,
-    int32_t ax_root_id) {
+    const ui::AXTreeID& tree_id,
+    bool is_root,
+    NodeIDMapper* id_mapper) {
   fuchsia::accessibility::semantics::Node fuchsia_node;
-  fuchsia_node.set_node_id(ConvertToFuchsiaNodeId(node.id, ax_root_id));
+  fuchsia_node.set_node_id(
+      id_mapper->ToFuchsiaNodeID(tree_id, node.id, is_root));
   fuchsia_node.set_role(AxRoleToFuchsiaSemanticRole(node.role));
   fuchsia_node.set_states(ConvertStates(node));
   fuchsia_node.set_attributes(ConvertAttributes(node));
   fuchsia_node.set_actions(ConvertActions(node));
-  fuchsia_node.set_child_ids(ConvertChildIds(node.child_ids, ax_root_id));
+  fuchsia_node.set_child_ids(
+      ConvertChildIds(node.child_ids, tree_id, id_mapper));
   fuchsia_node.set_location(ConvertBoundingBox(node.relative_bounds.bounds));
   if (node.relative_bounds.transform) {
     fuchsia_node.set_transform(
@@ -261,24 +261,68 @@ bool ConvertAction(fuchsia::accessibility::semantics::Action fuchsia_action,
   }
 }
 
-uint32_t ConvertToFuchsiaNodeId(int32_t ax_node_id, int32_t ax_root_node_id) {
-  if (ax_node_id == ax_root_node_id)
-    return kFuchsiaRootNodeId;
+NodeIDMapper::NodeIDMapper()
+    : root_(std::make_pair(ui::AXTreeIDUnknown(), 0)) {}
+
+NodeIDMapper::~NodeIDMapper() = default;
+
+uint32_t NodeIDMapper::ToFuchsiaNodeID(const ui::AXTreeID& ax_tree_id,
+                                       int32_t ax_node_id,
+                                       bool is_tree_root) {
+  const bool should_change_root =
+      is_tree_root && (root_.first != ax_tree_id || root_.second != ax_node_id);
+
+  CHECK_LE(next_fuchsia_id_, UINT32_MAX);
+  uint32_t fuchsia_node_id;
+  if (should_change_root) {
+    // The node that points to the root is changing. Update the old root to
+    // receive an unique ID, and make the new root receive the default value.
+    if (root_.first != ui::AXTreeIDUnknown())
+      id_map_[root_.first][root_.second] = next_fuchsia_id_++;
+    root_ = std::make_pair(ax_tree_id, ax_node_id);
+    fuchsia_node_id = kFuchsiaRootNodeId;
+  } else {
+    auto it = id_map_.find(ax_tree_id);
+    if (it != id_map_.end()) {
+      auto node_id_it = it->second.find(ax_node_id);
+      if (node_id_it != it->second.end())
+        return node_id_it->second;
+    }
 
-  // kInvalidAXID has the same value as the Fuchsia root ID. It is remapped to
-  // avoid a conflict.
-  if (ax_node_id == ui::AXNode::kInvalidAXID)
-    return kInvalidIdRemappedForFuchsia;
+    // The ID is not in the map yet, so give it a new value.
+    fuchsia_node_id = next_fuchsia_id_++;
+  }
 
-  return base::checked_cast<uint32_t>(ax_node_id);
+  id_map_[ax_tree_id][ax_node_id] = fuchsia_node_id;
+  return fuchsia_node_id;
 }
 
-int32_t ConvertToAxNodeId(uint32_t fuchsia_node_id, int32_t ax_root_node_id) {
-  if (fuchsia_node_id == kFuchsiaRootNodeId)
-    return ax_root_node_id;
+base::Optional<std::pair<ui::AXTreeID, int32_t>> NodeIDMapper::ToAXNodeID(
+    uint32_t fuchsia_node_id) {
+  for (const auto& tree_id_to_node_ids : id_map_) {
+    for (const auto& ax_id_to_fuchsia_id : tree_id_to_node_ids.second) {
+      if (ax_id_to_fuchsia_id.second == fuchsia_node_id)
+        return std::make_pair(tree_id_to_node_ids.first,
+                              ax_id_to_fuchsia_id.first);
+    }
+  }
 
-  if (fuchsia_node_id == kInvalidIdRemappedForFuchsia)
-    return ui::AXNode::kInvalidAXID;
+  return base::nullopt;
+}
 
-  return base::checked_cast<int32_t>(fuchsia_node_id);
+bool NodeIDMapper::UpdateAXTreeIDForCachedNodeIDs(
+    const ui::AXTreeID& old_ax_tree_id,
+    const ui::AXTreeID& new_ax_tree_id) {
+  if (old_ax_tree_id == new_ax_tree_id)
+    return false;
+
+  auto it = id_map_.find(old_ax_tree_id);
+  if (it == id_map_.end())
+    return false;
+
+  // The iterator is not stable, so the order of operations here is important.
+  auto data = std::move(it->second);
+  id_map_.erase(it);
+  id_map_[new_ax_tree_id] = std::move(data);
+  return true;
 }

fuchsia/engine/browser/ax_tree_converter.h

@@ -5,18 +5,69 @@
 #ifndef FUCHSIA_ENGINE_BROWSER_AX_TREE_CONVERTER_H_
 #define FUCHSIA_ENGINE_BROWSER_AX_TREE_CONVERTER_H_
 
+#include <base/containers/flat_map.h>
+#include <base/optional.h>
 #include <fuchsia/accessibility/semantics/cpp/fidl.h>
 
+#include <unordered_map>
+
 #include "content/public/browser/ax_event_notification_details.h"
 #include "fuchsia/engine/web_engine_export.h"
 
+// Maps AXNode IDs to Fuchsia Node IDs.
+// This class saves the remapped values.
+class WEB_ENGINE_EXPORT NodeIDMapper {
+ public:
+  NodeIDMapper();
+  virtual ~NodeIDMapper();
+
+  // Maps |ax_tree_id| and the signed |ax_node_id| to a unique unsigned
+  // |fuchsia_node_id|, with special handling of root IDs. A Fuchsia Node ID of
+  // 0 indicates the root, and is always returned if |is_tree_root| is true.
+  // No value is returned if this mapper can't produce more unique IDs.
+  virtual uint32_t ToFuchsiaNodeID(const ui::AXTreeID& ax_tree_id,
+                                   int32_t ax_node_id,
+                                   bool is_tree_root);
+
+  // From a Fuchsia Node ID, returns the pair of the AXTreeID and the AXNode ID
+  // that maps to it. If the Fuchsia Node ID is not in the map, returns no
+  // value.
+  virtual base::Optional<std::pair<ui::AXTreeID, int32_t>> ToAXNodeID(
+      uint32_t fuchsia_node_id);
+
+  // Updates the  AXNode IDs to point to the new |ax_tree_id|. This method
+  // must be called whenever an AXTree changes its AXTreeID, so that stored
+  // values here will point to the correct tree.
+  // Returns true if the update was applied successfully.
+  virtual bool UpdateAXTreeIDForCachedNodeIDs(
+      const ui::AXTreeID& old_ax_tree_id,
+      const ui::AXTreeID& new_ax_tree_id);
+
+ private:
+  // Keeps track of the next unused, unique node ID.
+  uint32_t next_fuchsia_id_ = 1;
+
+  // Pair that represents the current root.
+  std::pair<ui::AXTreeID, int32_t> root_;
+
+  // Stores the node ID mappings. Note that because converting from AXNode IDs
+  // to Fuchsia Node IDs is more common, the map makes access in this
+  // direction O(1). The storage representation is chosen here to use a map to
+  // hold the AXTreeIDs (which are just a few), but an unordered_map to hold the
+  // IDs (which can be thousands).
+  base::flat_map<ui::AXTreeID, std::unordered_map<int32_t, uint32_t>> id_map_;
+};
+
 // Converts an AXNodeData to a Fuchsia Semantic Node.
 // Both data types represent a single node, and no additional state is needed.
 // AXNodeData is used to convey partial updates, so not all fields may be
 // present. Those that are will be converted. The Fuchsia SemanticsManager
 // accepts partial updates, so |node| does not require all fields to be set.
 WEB_ENGINE_EXPORT fuchsia::accessibility::semantics::Node
-AXNodeDataToSemanticNode(const ui::AXNodeData& node, int32_t ax_root_id);
+AXNodeDataToSemanticNode(const ui::AXNodeData& node,
+                         const ui::AXTreeID& tree_id,
+                         bool is_root,
+                         NodeIDMapper* id_mapper);
 
 // Converts Fuchsia action of type |fuchsia_action| to an ax::mojom::Action of
 // type |mojom_action|. Function will return true if |fuchsia_action| is
@@ -24,14 +75,4 @@ AXNodeDataToSemanticNode(const ui::AXNodeData& node, int32_t ax_root_id);
 bool ConvertAction(fuchsia::accessibility::semantics::Action fuchsia_action,
                    ax::mojom::Action* mojom_action);
 
-// Converts between the signed |ax_node_id| and the unsigned |fuchsia_node_id|,
-// with special handling of root and invalid node ids.
-// A Fuchsia node id of 0 indicates the root.
-// An AXNode node id of 0 indicates an invalid node, and is remapped to
-// MAX(int32_t) + 1 to avoid conflict with other node ids.
-WEB_ENGINE_EXPORT uint32_t ConvertToFuchsiaNodeId(int32_t ax_node_id,
-                                                  int32_t ax_root_node_id);
-WEB_ENGINE_EXPORT int32_t ConvertToAxNodeId(uint32_t fuchsia_node_id,
-                                            int32_t ax_root_node_id);
-
 #endif  // FUCHSIA_ENGINE_BROWSER_AX_TREE_CONVERTER_H_

fuchsia/engine/browser/ax_tree_converter_unittest.cc

@@ -10,6 +10,8 @@
 #include "testing/gtest/include/gtest/gtest.h"
 #include "ui/gfx/transform.h"
 
+namespace {
+
 using fuchsia::accessibility::semantics::Action;
 using fuchsia::accessibility::semantics::Attributes;
 using fuchsia::accessibility::semantics::CheckedState;
@@ -17,8 +19,6 @@ using fuchsia::accessibility::semantics::Node;
 using fuchsia::accessibility::semantics::Role;
 using fuchsia::accessibility::semantics::States;
 
-namespace {
-
 const char kLabel1[] = "label nodes, not people";
 const char kLabel2[] = "fancy stickers";
 const char kDescription1[] = "this node does some stuff";
@@ -26,15 +26,25 @@ const char kValue1[] = "user entered value";
 const int32_t kChildId1 = 23901;
 const int32_t kChildId2 = 484345;
 const int32_t kChildId3 = 4156877;
-const int32_t kRootId = 5;
 const int32_t kRectX = 1;
 const int32_t kRectY = 2;
 const int32_t kRectWidth = 7;
 const int32_t kRectHeight = 8;
-const uint32_t kInvalidIdRemappedForFuchsia = 1u + INT32_MAX;
 const std::array<float, 16> k4DIdentityMatrix = {1, 0, 0, 0, 0, 1, 0, 0,
                                                  0, 0, 1, 0, 0, 0, 0, 1};
 
+class MockNodeIDMapper : public NodeIDMapper {
+ public:
+  MockNodeIDMapper() = default;
+  ~MockNodeIDMapper() override = default;
+
+  uint32_t ToFuchsiaNodeID(const ui::AXTreeID& ax_tree_id,
+                           int32_t ax_node_id,
+                           bool is_tree_root) override {
+    return base::checked_cast<uint32_t>(ax_node_id);
+  }
+};
+
 ui::AXNodeData CreateAXNodeData(ax::mojom::Role role,
                                 ax::mojom::Action action,
                                 std::vector<int32_t> child_ids,
@@ -110,9 +120,11 @@ std::pair<ui::AXNodeData, Node> CreateSemanticNodeAllFieldsSet() {
   states.set_hidden(false);
   states.set_selected(false);
   states.set_viewport_offset({10, 20});
+  MockNodeIDMapper mapper;
   auto fuchsia_node = CreateSemanticNode(
-      ConvertToFuchsiaNodeId(ax_node_data.id, kRootId), Role::BUTTON,
-      std::move(attributes), std::move(states),
+      mapper.ToFuchsiaNodeID(ui::AXTreeID::CreateNewAXTreeID(), ax_node_data.id,
+                             false),
+      Role::BUTTON, std::move(attributes), std::move(states),
       std::vector<Action>{Action::SET_FOCUS},
       std::vector<uint32_t>{kChildId1, kChildId2, kChildId3}, box, mat.value);
 
@@ -132,7 +144,9 @@ TEST_F(AXTreeConverterTest, AllFieldsSetAndEqual) {
   auto& source_node_data = nodes.first;
   auto& expected_node = nodes.second;
 
-  auto converted_node = AXNodeDataToSemanticNode(source_node_data, kRootId);
+  MockNodeIDMapper mapper;
+  auto converted_node = AXNodeDataToSemanticNode(
+      source_node_data, ui::AXTreeID::CreateNewAXTreeID(), false, &mapper);
   EXPECT_TRUE(fidl::Equals(converted_node, expected_node));
 }
 
@@ -145,15 +159,16 @@ TEST_F(AXTreeConverterTest, SomeFieldsSetAndEqual) {
   source_node_data.role = ax::mojom::Role::kImage;
   source_node_data.AddStringAttribute(ax::mojom::StringAttribute::kValue,
                                       kValue1);
-  auto converted_node = AXNodeDataToSemanticNode(source_node_data, kRootId);
+
+  MockNodeIDMapper mapper;
+  auto converted_node = AXNodeDataToSemanticNode(
+      source_node_data, ui::AXTreeID::CreateNewAXTreeID(), false, &mapper);
 
   Node expected_node;
-  expected_node.set_node_id(
-      ConvertToFuchsiaNodeId(source_node_data.id, kRootId));
+  expected_node.set_node_id(0);
   expected_node.set_actions(
       std::vector<Action>{Action::SET_FOCUS, Action::SET_VALUE});
-  expected_node.set_child_ids(
-      std::vector<uint32_t>{ConvertToFuchsiaNodeId(kChildId1, kRootId)});
+  expected_node.set_child_ids(std::vector<uint32_t>{kChildId1});
   expected_node.set_role(Role::IMAGE);
   States states;
   states.set_hidden(false);
@@ -177,7 +192,10 @@ TEST_F(AXTreeConverterTest, FieldMismatch) {
       ax::mojom::Role::kHeader, ax::mojom::Action::kSetValue,
       std::vector<int32_t>{kChildId1, kChildId2, kChildId3}, relative_bounds,
       kLabel1, kDescription1, ax::mojom::CheckedState::kFalse);
-  auto converted_node = AXNodeDataToSemanticNode(source_node_data, kRootId);
+
+  MockNodeIDMapper mapper;
+  auto converted_node = AXNodeDataToSemanticNode(
+      source_node_data, ui::AXTreeID::CreateNewAXTreeID(), false, &mapper);
 
   Attributes attributes;
   attributes.set_label(kLabel1);
@@ -204,13 +222,16 @@ TEST_F(AXTreeConverterTest, FieldMismatch) {
   auto modified_node_data = source_node_data;
   modified_node_data.AddStringAttribute(ax::mojom::StringAttribute::kName,
                                         kLabel2);
-  converted_node = AXNodeDataToSemanticNode(modified_node_data, kRootId);
+
+  converted_node = AXNodeDataToSemanticNode(
+      modified_node_data, ui::AXTreeID::CreateNewAXTreeID(), false, &mapper);
   EXPECT_FALSE(fidl::Equals(converted_node, expected_node));
 
   // The same as above, this time changing |child_ids|.
   modified_node_data = source_node_data;
   modified_node_data.child_ids = std::vector<int32_t>{};
-  converted_node = AXNodeDataToSemanticNode(modified_node_data, kRootId);
+  converted_node = AXNodeDataToSemanticNode(
+      modified_node_data, ui::AXTreeID::CreateNewAXTreeID(), false, &mapper);
   EXPECT_FALSE(fidl::Equals(converted_node, expected_node));
 }
 
@@ -224,7 +245,10 @@ TEST_F(AXTreeConverterTest, LocationFieldRespectsTypeInvariants) {
       ax::mojom::Role::kHeader, ax::mojom::Action::kSetValue,
       std::vector<int32_t>{kChildId1, kChildId2, kChildId3}, relative_bounds,
       kLabel1, kDescription1, ax::mojom::CheckedState::kFalse);
-  auto converted_node = AXNodeDataToSemanticNode(source_node_data, kRootId);
+
+  MockNodeIDMapper mapper;
+  auto converted_node = AXNodeDataToSemanticNode(
+      source_node_data, ui::AXTreeID::CreateNewAXTreeID(), false, &mapper);
 
   // The type definition of the location field requires that in order to be
   // interpreted as having non-zero length in a dimension, the min must be less
@@ -246,79 +270,111 @@ TEST_F(AXTreeConverterTest, DefaultAction) {
       expected_node.mutable_actions()->begin(),
       fuchsia::accessibility::semantics::Action::DEFAULT);
 
-  auto converted_node = AXNodeDataToSemanticNode(source_node_data, kRootId);
-  EXPECT_EQ(ConvertToFuchsiaNodeId(source_node_data.id, kRootId),
-            converted_node.node_id());
+  MockNodeIDMapper mapper;
+  auto converted_node = AXNodeDataToSemanticNode(
+      source_node_data, ui::AXTreeID::CreateNewAXTreeID(), false, &mapper);
 
   EXPECT_TRUE(fidl::Equals(converted_node, expected_node));
 }
 
-TEST_F(AXTreeConverterTest, ConvertToFuchsiaNodeId) {
-  // Root AxNode is 0, Fuchsia is also 0.
-  EXPECT_EQ(0u, ConvertToFuchsiaNodeId(0, 0));
-
-  // Root AxNode is not 0, Fuchsia is still 0.
-  EXPECT_EQ(0u, ConvertToFuchsiaNodeId(2, 2));
-
-  // Regular AxNode is 0, Fuchsia can't be 0.
-  EXPECT_EQ(kInvalidIdRemappedForFuchsia, ConvertToFuchsiaNodeId(0, 2));
-
-  // Regular AxNode is not 0, Fuchsia is same value.
-  EXPECT_EQ(10u, ConvertToFuchsiaNodeId(10, 0));
-}
-
-TEST_F(AXTreeConverterTest, ConvertToAxNodeId) {
-  // Root Fuchsia id is 0, AxNode is also 0.
-  EXPECT_EQ(0, ConvertToAxNodeId(0u, 0));
-
-  // Root Fuchsia id is 0, AxNode is not 0.
-  EXPECT_EQ(2, ConvertToAxNodeId(0u, 2));
-
-  // Fuchsia node is the max value, AxNode should be 0.
-  EXPECT_EQ(0, ConvertToAxNodeId(kInvalidIdRemappedForFuchsia, 2));
-
-  // Regular Fuchsia node is not root, AxNode is same value.
-  EXPECT_EQ(10, ConvertToAxNodeId(10u, 0));
+TEST_F(AXTreeConverterTest, MapsNodeIDs) {
+  NodeIDMapper mapper;
+  const ui::AXTreeID tree_id_1 = ui::AXTreeID::CreateNewAXTreeID();
+  const ui::AXTreeID tree_id_2 = ui::AXTreeID::CreateNewAXTreeID();
+  const ui::AXTreeID tree_id_3 = ui::AXTreeID::CreateNewAXTreeID();
+
+  auto id = mapper.ToFuchsiaNodeID(tree_id_1, 1, false);
+  EXPECT_EQ(id, 1u);
+
+  id = mapper.ToFuchsiaNodeID(tree_id_2, 1, false);
+  EXPECT_EQ(id, 2u);
+
+  const auto result_1 = mapper.ToAXNodeID(1u);
+  EXPECT_TRUE(result_1);
+  EXPECT_EQ(result_1->first, tree_id_1);
+  EXPECT_EQ(result_1->second, 1);
+
+  const auto result_2 = mapper.ToAXNodeID(2u);
+  EXPECT_TRUE(result_2);
+  EXPECT_EQ(result_2->first, tree_id_2);
+  EXPECT_EQ(result_2->second, 1);
+
+  // Set the root.
+  id = mapper.ToFuchsiaNodeID(tree_id_1, 2, true);
+  EXPECT_EQ(id, 0u);
+
+  // Update the root. The old root should receive a new value.
+  id = mapper.ToFuchsiaNodeID(tree_id_1, 1, true);
+  EXPECT_EQ(id, 0u);
+  const auto result_3 = mapper.ToAXNodeID(3u);
+  EXPECT_TRUE(result_3);
+  EXPECT_EQ(result_3->first, tree_id_1);
+  EXPECT_EQ(result_3->second, 2);  // First root's ID.
+
+  mapper.UpdateAXTreeIDForCachedNodeIDs(tree_id_1, tree_id_3);
+  const auto result_4 = mapper.ToAXNodeID(3u);
+  EXPECT_TRUE(result_4);
+  EXPECT_EQ(result_4->first, tree_id_3);
+  EXPECT_EQ(result_4->second, 2);
 }
 
 TEST_F(AXTreeConverterTest, ConvertRoles) {
+  MockNodeIDMapper mapper;
   ui::AXNodeData node;
   node.id = 0;
   node.role = ax::mojom::Role::kButton;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::BUTTON,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 
   node.role = ax::mojom::Role::kCheckBox;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::CHECK_BOX,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 
   node.role = ax::mojom::Role::kHeader;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::HEADER,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 
   node.role = ax::mojom::Role::kImage;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::IMAGE,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 
   node.role = ax::mojom::Role::kLink;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::LINK,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 
   node.role = ax::mojom::Role::kRadioButton;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::RADIO_BUTTON,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 
   node.role = ax::mojom::Role::kSlider;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::SLIDER,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 
   node.role = ax::mojom::Role::kTextField;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::TEXT_FIELD,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 
   node.role = ax::mojom::Role::kStaticText;
   EXPECT_EQ(fuchsia::accessibility::semantics::Role::STATIC_TEXT,
-            AXNodeDataToSemanticNode(node, kRootId).role());
+            AXNodeDataToSemanticNode(node, ui::AXTreeID::CreateNewAXTreeID(),
+                                     false, &mapper)
+                .role());
 }
 
 }  // namespace