perf: Port FrozenOrderedSet to rust

Author: tobni

bench_frozen_ordered_set.py

@@ -0,0 +1,179 @@
+"""Benchmark: Rust FrozenOrderedSet vs Python FrozenOrderedSet."""
+
+import sys
+import timeit
+from collections.abc import Hashable, Iterable, Iterator
+from typing import AbstractSet, Any, TypeVar
+
+sys.path.insert(0, "src/python")
+
+from pants.engine.internals.native_engine import FrozenOrderedSet as RustFrozenOrderedSet
+
+T = TypeVar("T")
+
+
+class PyFrozenOrderedSet(AbstractSet[T], Hashable):
+    """The old pure-Python FrozenOrderedSet (pre-port)."""
+
+    def __init__(self, iterable=None):
+        self._items = dict.fromkeys(iterable) if iterable else {}
+        self._hash = None
+
+    def __len__(self):
+        return len(self._items)
+
+    def __contains__(self, key):
+        return key in self._items
+
+    def __iter__(self) -> Iterator:
+        return iter(self._items)
+
+    def __reversed__(self):
+        return reversed(tuple(self._items.keys()))
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+        return len(self._items) == len(other._items) and all(
+            x == y for x, y in zip(self._items, other._items)
+        )
+
+    def __hash__(self):
+        if self._hash is None:
+            self._hash = 0
+            for item in self._items.keys():
+                self._hash ^= hash(item)
+        return self._hash
+
+    def __repr__(self):
+        return f"PyFrozenOrderedSet({list(self)!r})"
+
+    def __bool__(self):
+        return bool(self._items)
+
+    def union(self, other):
+        return self.__class__(list(self) + [x for x in other if x not in self._items])
+
+    def intersection(self, other):
+        s = set(other)
+        return self.__class__(x for x in self if x in s)
+
+    def difference(self, other):
+        s = set(other)
+        return self.__class__(x for x in self if x not in s)
+
+    def issubset(self, other):
+        if len(self) > len(other):
+            return False
+        return all(item in other for item in self)
+
+
+WARMUP = 1000
+
+def measure(stmt, number, globs):
+    timeit.timeit(stmt, number=WARMUP, globals=globs)
+    t = timeit.timeit(stmt, number=number, globals=globs)
+    return t / number * 1_000_000
+
+
+BENCHMARKS = [
+    ("Construction", "Cls(data)", lambda data, py, rs, **_: [
+        {"Cls": PyFrozenOrderedSet, "data": data},
+        {"Cls": RustFrozenOrderedSet, "data": data},
+    ]),
+    ("hash()", "hash(fd)", lambda py, rs, **_: [
+        {"fd": py},
+        {"fd": rs},
+    ]),
+    ("__contains__", "k in fd", lambda py, rs, mid, **_: [
+        {"fd": py, "k": mid},
+        {"fd": rs, "k": mid},
+    ]),
+    ("__contains__ miss", "k in fd", lambda py, rs, **_: [
+        {"fd": py, "k": "MISSING"},
+        {"fd": rs, "k": "MISSING"},
+    ]),
+    ("__eq__", "fd == fd2", lambda py, rs, py2, rs2, **_: [
+        {"fd": py, "fd2": py2},
+        {"fd": rs, "fd2": rs2},
+    ]),
+    ("iteration", "list(fd)", lambda py, rs, **_: [
+        {"fd": py},
+        {"fd": rs},
+    ]),
+    ("union", "fd.union(other)", lambda py, rs, py_other, rs_other, **_: [
+        {"fd": py, "other": py_other},
+        {"fd": rs, "other": rs_other},
+    ]),
+    ("intersection", "fd.intersection(other)", lambda py, rs, py_other, rs_other, **_: [
+        {"fd": py, "other": py_other},
+        {"fd": rs, "other": rs_other},
+    ]),
+    ("difference", "fd.difference(other)", lambda py, rs, py_other, rs_other, **_: [
+        {"fd": py, "other": py_other},
+        {"fd": rs, "other": rs_other},
+    ]),
+    ("issubset", "small.issubset(fd)", lambda py, rs, py_small, rs_small, **_: [
+        {"small": py_small, "fd": py},
+        {"small": rs_small, "fd": rs},
+    ]),
+    ("dict key", "d[fd]", lambda py, rs, **_: [
+        {"fd": py, "d": {py: 1}},
+        {"fd": rs, "d": {rs: 1}},
+    ]),
+]
+
+SMALL = list(range(5))
+MEDIUM = list(range(20))
+LARGE = list(range(200))
+
+DATASETS = [("small (5)", SMALL), ("medium (20)", MEDIUM), ("large (200)", LARGE)]
+
+all_results: dict[str, dict[str, tuple[float, float]]] = {}
+
+for ds_name, data in DATASETS:
+    print(f"\n{'=' * 60}")
+    print(f"  Dataset: {ds_name}")
+    print(f"{'=' * 60}")
+
+    py = PyFrozenOrderedSet(data)
+    rs = RustFrozenOrderedSet(data)
+    py2 = PyFrozenOrderedSet(data)
+    rs2 = RustFrozenOrderedSet(data)
+    half = data[:len(data) // 2]
+    py_other = PyFrozenOrderedSet(half + list(range(1000, 1000 + len(half))))
+    rs_other = RustFrozenOrderedSet(half + list(range(1000, 1000 + len(half))))
+    py_small = PyFrozenOrderedSet(data[:3])
+    rs_small = RustFrozenOrderedSet(data[:3])
+    # Warm up lazy hashes
+    for obj in (py, py2, py_other, py_small):
+        hash(obj)
+    n = 500_000 if len(data) <= 20 else 50_000
+    mid = data[len(data) // 2]
+
+    ctx = dict(data=data, py=py, rs=rs, py2=py2, rs2=rs2, mid=mid,
+               py_other=py_other, rs_other=rs_other, py_small=py_small, rs_small=rs_small)
+
+    for bench_name, stmt, make_globs in BENCHMARKS:
+        py_globs, rs_globs = make_globs(**ctx)
+        py_us = measure(stmt, n, py_globs)
+        rs_us = measure(stmt, n, rs_globs)
+        print(f"  {bench_name:.<20s} Python {py_us:8.3f} µs  Rust {rs_us:8.3f} µs  ({py_us / rs_us:.1f}x)")
+        all_results.setdefault(bench_name, {})[ds_name] = (py_us, rs_us)
+
+ds_names = [name for name, _ in DATASETS]
+header = f"  {'Operation':<20s}" + "".join(f" | {name:>12s}" for name in ds_names)
+sep = f"  {'-'*20}" + "".join(f"-+-{'-'*12}" for _ in ds_names)
+
+print(f"\n{'=' * 60}")
+print("  Summary (Python / Rust speedup)")
+print(f"{'=' * 60}")
+print(header)
+print(sep)
+for bench_name, _, _ in BENCHMARKS:
+    row = f"  {bench_name:<20s}"
+    for ds_name in ds_names:
+        py_us, rs_us = all_results[bench_name][ds_name]
+        ratio = py_us / rs_us
+        row += f" | {ratio:11.1f}x"
+    print(row)

src/python/pants/backend/python/util_rules/interpreter_constraints.py

@@ -89,14 +89,14 @@ def for_fixed_python_version(
     ) -> InterpreterConstraints:
         return cls([f"{interpreter_type}=={python_version_str}"])
 
-    def __init__(self, constraints: Iterable[str | Requirement] = ()) -> None:
+    def __new__(cls, constraints: Iterable[str | Requirement] = ()) -> InterpreterConstraints:
         # #12578 `parse_constraint` will sort the requirement's component constraints into a stable form.
         # We need to sort the component constraints for each requirement _before_ sorting the entire list
         # for the ordering to be correct.
         parsed_constraints = (
             i if isinstance(i, Requirement) else parse_constraint(i) for i in constraints
         )
-        super().__init__(sorted(parsed_constraints, key=lambda c: str(c)))
+        return super().__new__(cls, sorted(parsed_constraints, key=lambda c: str(c)))
 
     def __str__(self) -> str:
         return " OR ".join(str(constraint) for constraint in self)

src/python/pants/backend/python/util_rules/pex.py

@@ -133,7 +133,6 @@ class CompletePlatforms(DeduplicatedCollection[str]):
     sort_input = True
 
     def __init__(self, iterable: Iterable[str] = (), *, digest: Digest = EMPTY_DIGEST):
-        super().__init__(iterable)
         self._digest = digest
 
     @classmethod

src/python/pants/engine/collection.py

@@ -79,10 +79,11 @@ class Examples(DeduplicatedCollection[Example]):
 
     sort_input: ClassVar[bool] = False
 
-    def __init__(self, iterable: Iterable[T] = ()) -> None:
-        super().__init__(
-            iterable if not self.sort_input else sorted(iterable)  # type: ignore[type-var]
+    def __new__(cls, iterable: Iterable[T] = (), **_kwargs: object) -> DeduplicatedCollection[T]:
+        return super().__new__(
+            cls,
+            iterable if not cls.sort_input else sorted(iterable),  # type: ignore[type-var]
         )
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({list(self._items)})"
+        return f"{self.__class__.__name__}({list(self)})"

src/python/pants/engine/internals/native_engine.pyi

@@ -6,11 +6,11 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
 from datetime import datetime
 from io import RawIOBase
 from pathlib import Path
-from typing import Any, ClassVar, Protocol, Self, TextIO, TypeVar, overload
+from typing import AbstractSet, Any, ClassVar, Protocol, Self, TextIO, TypeVar, overload
 
 from pants.engine.fs import (
     CreateDigest,
@@ -81,6 +81,34 @@ class FrozenDict(Mapping[K, V]):
     def __hash__(self) -> int: ...
     def __repr__(self) -> str: ...
 
+T_co = TypeVar("T_co", covariant=True)
+
+class FrozenOrderedSet(AbstractSet[T_co], Hashable):
+    """A frozen (i.e. immutable) ordered set backed by Rust.
+
+    This is safe to use with the V2 engine.
+    """
+
+    def __new__(cls, iterable: Iterable[T_co] | None = None) -> Self: ...
+    def __len__(self) -> int: ...
+    def __contains__(self, key: Any) -> bool: ...
+    def __iter__(self) -> Iterator[T_co]: ...
+    def __reversed__(self) -> Iterator[T_co]: ...
+    def __hash__(self) -> int: ...
+    def __eq__(self, other: Any) -> bool: ...
+    def __or__(self, other: Iterable[T_co]) -> FrozenOrderedSet[T_co]: ...  # type: ignore[override]  # widens from AbstractSet
+    def __and__(self, other: Iterable[T_co]) -> FrozenOrderedSet[T_co]: ...
+    def __sub__(self, other: Iterable[T_co]) -> FrozenOrderedSet[T_co]: ...
+    def __xor__(self, other: Iterable[T_co]) -> FrozenOrderedSet[T_co]: ...  # type: ignore[override]  # widens from AbstractSet
+    def __bool__(self) -> bool: ...
+    def __repr__(self) -> str: ...
+    def union(self, *others: Iterable[T_co]) -> FrozenOrderedSet[T_co]: ...
+    def intersection(self, *others: Iterable[T_co]) -> FrozenOrderedSet[T_co]: ...
+    def difference(self, *others: Iterable[T_co]) -> FrozenOrderedSet[T_co]: ...
+    def symmetric_difference(self, other: Iterable[T_co]) -> FrozenOrderedSet[T_co]: ...
+    def issubset(self, other: Iterable[T_co]) -> bool: ...
+    def issuperset(self, other: Iterable[T_co]) -> bool: ...
+
 # ------------------------------------------------------------------------------
 # Address
 # ------------------------------------------------------------------------------

src/python/pants/util/ordered_set.py

@@ -15,9 +15,11 @@
 from __future__ import annotations
 
 import itertools
-from collections.abc import Hashable, Iterable, Iterator, MutableSet
+from collections.abc import Iterable, Iterator, MutableSet
 from typing import AbstractSet, Any, TypeVar, cast
 
+from pants.engine.internals.native_engine import FrozenOrderedSet as FrozenOrderedSet  # noqa: F401
+
 T = TypeVar("T")
 T_co = TypeVar("T_co", covariant=True)
 _TAbstractOrderedSet = TypeVar("_TAbstractOrderedSet", bound="_AbstractOrderedSet")
@@ -195,21 +197,3 @@ def symmetric_difference_update(self, other: Iterable[T]) -> None:
         self._items = {item: None for item in self._items.keys() if item not in items_to_remove}
         for item in items_to_add:
             self._items[item] = None
-
-
-class FrozenOrderedSet(_AbstractOrderedSet[T_co], Hashable):  # type: ignore[type-var]
-    """A frozen (i.e. immutable) set that retains its order.
-
-    This is safe to use with the V2 engine.
-    """
-
-    def __init__(self, iterable: Iterable[T_co] | None = None) -> None:
-        super().__init__(iterable)
-        self.__hash: int | None = None
-
-    def __hash__(self) -> int:
-        if self.__hash is None:
-            self.__hash = 0
-            for item in self._items.keys():
-                self.__hash ^= hash(item)
-        return self.__hash