ax_ui: make AXPosition keeping an anchor as an AXNode instance

Change-Id: I7ae44cbcd789683715fc3fd32e832ed100845e2b
Bug: 1362839
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3989063
Reviewed-by: Aaron Leventhal <aleventhal@chromium.org>
Commit-Queue: Alexander Surkkov <asurkov@igalia.com>
Cr-Commit-Position: refs/heads/main@{#1076372}

ui/accessibility/ax_node.h

@@ -17,6 +17,7 @@
 #include "base/containers/queue.h"
 #include "base/containers/stack.h"
 #include "base/memory/raw_ptr.h"
+#include "base/memory/weak_ptr.h"
 #include "build/build_config.h"
 #include "third_party/abseil-cpp/absl/types/optional.h"
 #include "third_party/skia/include/core/SkColor.h"
@@ -38,7 +39,7 @@ struct AXLanguageInfo;
 class AXTree;
 
 // This class is used to represent a node in an accessibility tree (`AXTree`).
-class AX_EXPORT AXNode final {
+class AX_EXPORT AXNode final : public base::SupportsWeakPtr<AXNode> {
  public:
   // Replacement character used to represent an embedded (or, additionally for
   // text navigation, an empty) object. Part of the Unicode Standard.

ui/accessibility/ax_position.h

@@ -237,8 +237,8 @@ class AXPosition {
 
   static AXPositionInstance CreateNullPosition() {
     AXPositionInstance new_position(new AXPositionType());
-    new_position->Initialize(AXPositionKind::NULL_POSITION, AXTreeIDUnknown(),
-                             kInvalidAXNodeID, INVALID_INDEX, INVALID_OFFSET,
+    new_position->Initialize(AXPositionKind::NULL_POSITION, nullptr,
+                             INVALID_INDEX, INVALID_OFFSET,
                              ax::mojom::TextAffinity::kDownstream);
     return new_position;
   }
@@ -250,8 +250,7 @@ class AXPosition {
     DCHECK_NE(anchor.id(), kInvalidAXNodeID);
 
     AXPositionInstance new_position(new AXPositionType());
-    new_position->Initialize(AXPositionKind::TREE_POSITION,
-                             anchor.tree()->GetAXTreeID(), anchor.id(),
+    new_position->Initialize(AXPositionKind::TREE_POSITION, &anchor,
                              child_index, INVALID_OFFSET,
                              ax::mojom::TextAffinity::kDownstream);
     return new_position;
@@ -282,8 +281,7 @@ class AXPosition {
     DCHECK_NE(anchor.id(), kInvalidAXNodeID);
 
     AXPositionInstance new_position(new AXPositionType());
-    new_position->Initialize(AXPositionKind::TEXT_POSITION,
-                             anchor.tree()->GetAXTreeID(), anchor.id(),
+    new_position->Initialize(AXPositionKind::TEXT_POSITION, &anchor,
                              INVALID_INDEX, text_offset, affinity);
     return new_position;
   }
@@ -343,12 +341,12 @@ class AXPosition {
     result.kind = kind_;
 
     // A tree ID can be serialized as a 32-byte string.
-    std::string tree_id_string = tree_id_.ToString();
+    std::string tree_id_string = GetTreeID().ToString();
     DCHECK_LE(tree_id_string.size(), 32U);
     strncpy(result.tree_id, tree_id_string.c_str(), 32);
     result.tree_id[32] = 0;
 
-    result.anchor_id = anchor_id_;
+    result.anchor_id = GetAnchorID();
     result.child_index = child_index_;
     result.text_offset = text_offset_;
     result.affinity = affinity_;
@@ -381,8 +379,8 @@ class AXPosition {
         } else {
           str_child_index = base::NumberToString(child_index_);
         }
-        str = "TreePosition tree_id=" + tree_id_.ToString() +
-              " anchor_id=" + base::NumberToString(anchor_id_) +
+        str = "TreePosition tree_id=" + GetTreeID().ToString() +
+              " anchor_id=" + base::NumberToString(GetAnchorID()) +
               " child_index=" + str_child_index;
         break;
       }
@@ -393,7 +391,7 @@ class AXPosition {
         } else {
           str_text_offset = base::NumberToString(text_offset_);
         }
-        str = "TextPosition anchor_id=" + base::NumberToString(anchor_id_) +
+        str = "TextPosition anchor_id=" + base::NumberToString(GetAnchorID()) +
               " text_offset=" + str_text_offset + " affinity=" +
               ui::ToString(static_cast<ax::mojom::TextAffinity>(affinity_));
         break;
@@ -443,10 +441,11 @@ class AXPosition {
   }
 
   AXPositionKind kind() const { return kind_; }
-  AXTreeID tree_id() const { return tree_id_; }
-  AXNodeID anchor_id() const { return anchor_id_; }
 
-  AXTreeManager* GetManager() const { return AXTreeManager::FromID(tree_id()); }
+  // Deprecated.
+  // TODO(crbug.com/1362839): replace on GetTreeID()/GetAnchorID().
+  AXTreeID tree_id() const { return GetTreeID(); }
+  AXNodeID anchor_id() const { return GetAnchorID(); }
 
   // Returns true if this position is within an "empty object", i.e. within a
   // node that should contribute no text to the accessibility tree's text
@@ -521,15 +520,18 @@ class AXPosition {
     return !node.GetChildCountCrossingTreeBoundary() || IsEmptyObject(node);
   }
 
-  AXNode* GetAnchor() const {
-    if (tree_id_ == AXTreeIDUnknown() || anchor_id_ == kInvalidAXNodeID)
-      return nullptr;
-
-    const AXTreeManager* manager = GetManager();
-    if (manager)
-      return manager->GetNode(anchor_id());
+  AXNode* GetAnchor() const { return anchor_.get(); }
+  AXTree* GetTree() const { return anchor_ ? anchor_->tree() : nullptr; }
+  AXTreeManager* GetManager() const {
+    return anchor_ ? anchor_->GetManager() : nullptr;
+  }
 
-    return nullptr;
+  AXNodeID GetAnchorID() const {
+    return anchor_ ? anchor_->id() : kInvalidAXNodeID;
+  }
+  AXTreeID GetTreeID() const {
+    DCHECK(!anchor_ || anchor_->tree());
+    return anchor_ ? anchor_->tree()->GetAXTreeID() : AXTreeIDUnknown();
   }
 
   int GetAnchorSiblingCount() const {
@@ -650,9 +652,7 @@ class AXPosition {
   bool IsValid() const {
     switch (kind_) {
       case AXPositionKind::NULL_POSITION:
-        return tree_id_ == AXTreeIDUnknown() &&
-               anchor_id_ == kInvalidAXNodeID &&
-               child_index_ == INVALID_INDEX &&
+        return !GetAnchor() && child_index_ == INVALID_INDEX &&
                text_offset_ == INVALID_OFFSET &&
                affinity_ == ax::mojom::TextAffinity::kDownstream;
       case AXPositionKind::TREE_POSITION:
@@ -2338,7 +2338,7 @@ class AXPosition {
             ->CreatePositionAtStartOfAnchor();
     if (IsTextPosition())
       root_position = root_position->AsTextPosition();
-    DCHECK_EQ(root_position->tree_id_, tree_id_)
+    DCHECK_EQ(root_position->GetTree(), GetTree())
         << "`CreatePositionAtStartOfAXTree` should not cross any tree "
            "boundaries, neither return the null position.";
     return root_position;
@@ -4257,8 +4257,7 @@ class AXPosition {
 
   void swap(AXPosition& other) {
     std::swap(kind_, other.kind_);
-    std::swap(tree_id_, other.tree_id_);
-    std::swap(anchor_id_, other.anchor_id_);
+    std::swap(anchor_, other.anchor_);
     std::swap(child_index_, other.child_index_);
     std::swap(text_offset_, other.text_offset_);
     std::swap(affinity_, other.affinity_);
@@ -4401,18 +4400,15 @@ class AXPosition {
 
  protected:
   AXPosition()
-      : kind_(AXPositionKind::NULL_POSITION),
-        tree_id_(AXTreeIDUnknown()),
-        anchor_id_(kInvalidAXNodeID),
+      : anchor_(nullptr),
         child_index_(INVALID_INDEX),
         text_offset_(INVALID_OFFSET),
         affinity_(ax::mojom::TextAffinity::kDownstream) {}
 
   // We explicitly don't copy any cached members.
   AXPosition(const AXPosition& other)
       : kind_(other.kind_),
-        tree_id_(other.tree_id_),
-        anchor_id_(other.anchor_id_),
+        anchor_(other.anchor_),
         child_index_(other.child_index_),
         text_offset_(other.text_offset_),
         affinity_(other.affinity_),
@@ -4482,46 +4478,44 @@ class AXPosition {
                                    int text_offset,
                                    ax::mojom::TextAffinity affinity) {
     kind_ = kind;
-    tree_id_ = tree_id;
-    anchor_id_ = anchor_id;
     child_index_ = child_index;
     text_offset_ = text_offset;
     affinity_ = affinity;
 
+    AXTreeManager* manager = AXTreeManager::FromID(tree_id);
+    if (manager) {
+      AXNode* anchor = manager->GetNode(anchor_id);
+      if (anchor)
+        anchor_ = anchor->AsWeakPtr();
+    }
+
     if (!IsValid()) {
       // Reset to the null position.
       kind_ = AXPositionKind::NULL_POSITION;
-      tree_id_ = AXTreeIDUnknown();
-      anchor_id_ = kInvalidAXNodeID;
+      anchor_ = nullptr;
       child_index_ = INVALID_INDEX;
       text_offset_ = INVALID_OFFSET;
       affinity_ = ax::mojom::TextAffinity::kDownstream;
     }
   }
 
   void Initialize(AXPositionKind kind,
-                  AXTreeID tree_id,
-                  AXNodeID anchor_id,
+                  const AXNode* anchor,
                   int child_index,
                   int text_offset,
                   ax::mojom::TextAffinity affinity) {
     kind_ = kind;
-    tree_id_ = tree_id;
-    anchor_id_ = anchor_id;
     child_index_ = child_index;
     text_offset_ = text_offset;
     affinity_ = affinity;
 
-    // TODO(accessibility) Consider using WeakPtr<AXTree> instead of an
-    // AXTreeID, which would be both faster and easier to use in combination
-    // with AXTreeSnapshotter, which does not use AXTreeManager to cache
-    // AXTreeIDs in a map.
-    SANITIZER_CHECK(GetManager() || kind_ == AXPositionKind::NULL_POSITION)
-        << "Tree manager required, tree_id = " << tree_id.ToString()
-        << "  is unknown = " << (tree_id == AXTreeIDUnknown());
-    SANITIZER_CHECK(GetAnchor() || kind_ == AXPositionKind::NULL_POSITION)
-        << "Creating a position without an anchor is disallowed:\n"
-        << ToDebugString();
+    if (kind_ != AXPositionKind::NULL_POSITION) {
+      DCHECK(anchor);
+      SANITIZER_CHECK(anchor->GetManager())
+          << "Creating a position without a tree manager is disallowed:\n"
+          << ToDebugString();
+      anchor_ = const_cast<AXNode*>(anchor)->AsWeakPtr();
+    }
 
     // TODO(accessibility) Remove this line and let the below IsValid()
     // assertion get triggered instead. We shouldn't be creating test positions
@@ -5574,12 +5568,8 @@ class AXPosition {
     return text_position;
   }
 
-  AXPositionKind kind_;
-  // TODO(crbug.com/1362839): use weak pointers for the AXTree, so that
-  // AXPosition can be used without AXTreeManager support (and also faster than
-  // the slow AXTreeID).
-  AXTreeID tree_id_;
-  AXNodeID anchor_id_;
+  AXPositionKind kind_ = AXPositionKind::NULL_POSITION;
+  base::WeakPtr<AXNode> anchor_;
 
   // For text positions, |child_index_| is initially set to |-1| and only
   // computed on demand. The same with tree positions and |text_offset_|.