Memoize equality for CoarsenedTarget(s) to avoid exponential runtim…

…e in `check`. (#15277)

#15141 gave `CoarsenedTarget` new usecases, which unfortunately exposed a case where `check` would re-run in exponential runtime due to not having implemented a TODO about `__eq__` using a memoized graph walk.

[ci skip-rust]
[ci skip-build-wheels]

src/python/pants/engine/target.py

@@ -712,16 +712,31 @@ def coarsened_closure(
     def __hash__(self) -> int:
         return self._hashcode
 
-    def __eq__(self, other: Any) -> bool:
-        if not isinstance(other, CoarsenedTarget):
-            return NotImplemented
-        return (
+    def _eq_helper(self, other: CoarsenedTarget, equal_items: set[tuple[int, int]]) -> bool:
+        key = (id(self), id(other))
+        if key[0] == key[1] or key in equal_items:
+            return True
+
+        is_eq = (
             self._hashcode == other._hashcode
             and self.members == other.members
-            # TODO: Use a recursive memoized __eq__ if this ever shows up in profiles.
-            and self.dependencies == other.dependencies
+            and len(self.dependencies) == len(other.dependencies)
+            and all(
+                l._eq_helper(r, equal_items) for l, r in zip(self.dependencies, other.dependencies)
+            )
         )
 
+        # NB: We only track equal items because any non-equal item will cause the entire
+        # operation to shortcircuit.
+        if is_eq:
+            equal_items.add(key)
+        return is_eq
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, CoarsenedTarget):
+            return NotImplemented
+        return self._eq_helper(other, set())
+
     def __str__(self) -> str:
         if len(self.members) > 1:
             others = len(self.members) - 1
@@ -752,6 +767,16 @@ def coarsened_closure(self) -> Iterator[CoarsenedTarget]:
         visited: Set[CoarsenedTarget] = set()
         return (ct for root in self for ct in root.coarsened_closure(visited))
 
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, CoarsenedTargets):
+            return NotImplemented
+        equal_items: set[tuple[int, int]] = set()
+        return len(self) == len(other) and all(
+            l._eq_helper(r, equal_items) for l, r in zip(self, other)
+        )
+
+    __hash__ = Tuple.__hash__
+
 
 @frozen_after_init
 @dataclass(unsafe_hash=True)

src/python/pants/engine/target_test.py

@@ -1,6 +1,7 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+import string
 from collections import namedtuple
 from dataclasses import dataclass
 from enum import Enum
@@ -13,6 +14,7 @@
 from pants.engine.target import (
     AsyncFieldMixin,
     BoolField,
+    CoarsenedTarget,
     DictStringToStringField,
     DictStringToStringSequenceField,
     ExplicitlyProvidedDependencies,
@@ -442,6 +444,35 @@ def test_target_residence_dir() -> None:
     )
 
 
+# -----------------------------------------------------------------------------------------------
+# Test CoarsenedTarget
+# -----------------------------------------------------------------------------------------------
+
+
+def test_coarsened_target_equality() -> None:
+    a, b = (FortranTarget({}, Address(name)) for name in string.ascii_lowercase[:2])
+
+    def ct(members: List[Target], dependencies: List[CoarsenedTarget] = []):
+        return CoarsenedTarget(members, dependencies)
+
+    assert ct([]) == ct([])
+
+    assert ct([a]) == ct([a])
+    assert ct([a]) != ct([b])
+
+    # Unique instances.
+    assert ct([], [ct([a])]) == ct([], [ct([a])])
+    assert ct([], [ct([a])]) != ct([], [ct([b])])
+
+    # Create two root CTs (with unique `id`s), which contain some reused instances.
+    def nested():
+        ct_a = ct([a])
+        return ct([], [ct_a, ct([], [ct_a])])
+
+    assert id(nested()) != id(nested())
+    assert nested() == nested()
+
+
 # -----------------------------------------------------------------------------------------------
 # Test file-level target generation
 # -----------------------------------------------------------------------------------------------