perf: Port FrozenOrderedSet to rust

Author: tobni

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,35 @@ 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: ...
+    def isdisjoint(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

src/rust/engine/src/externs/collection.rs

@@ -0,0 +1,119 @@
+// Copyright 2026 Pants project contributors (see CONTRIBUTORS.md).
+// Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+use std::fmt::Debug;
+use std::sync::OnceLock;
+
+use pyo3::prelude::*;
+use pyo3::types::{PyDict, PyIterator};
+
+pub trait HashCache: Debug + Send + Sync {
+    fn new_eager(hash: isize) -> Self;
+    fn new_lazy() -> Self;
+    fn get(
+        &self,
+        dict: &Bound<PyDict>,
+        compute: fn(&Bound<PyDict>) -> PyResult<isize>,
+    ) -> PyResult<isize>;
+}
+
+#[derive(Debug)]
+pub struct EagerHash(isize);
+
+impl HashCache for EagerHash {
+    fn new_eager(hash: isize) -> Self {
+        Self(hash)
+    }
+    fn new_lazy() -> Self {
+        panic!("EagerHash requires a value at construction")
+    }
+    fn get(
+        &self,
+        _dict: &Bound<PyDict>,
+        _compute: fn(&Bound<PyDict>) -> PyResult<isize>,
+    ) -> PyResult<isize> {
+        Ok(self.0)
+    }
+}
+
+#[derive(Debug)]
+pub struct LazyHash(OnceLock<isize>);
+
+impl HashCache for LazyHash {
+    fn new_eager(hash: isize) -> Self {
+        let lock = OnceLock::new();
+        let _ = lock.set(hash);
+        Self(lock)
+    }
+    fn new_lazy() -> Self {
+        Self(OnceLock::new())
+    }
+    fn get(
+        &self,
+        dict: &Bound<PyDict>,
+        compute: fn(&Bound<PyDict>) -> PyResult<isize>,
+    ) -> PyResult<isize> {
+        if let Some(&h) = self.0.get() {
+            return Ok(h);
+        }
+        let h = compute(dict)?;
+        let _ = self.0.set(h);
+        Ok(h)
+    }
+}
+
+#[derive(Debug)]
+pub struct FrozenCollectionData<H: HashCache = EagerHash> {
+    pub data: Py<PyDict>,
+    hash: H,
+}
+
+impl<H: HashCache> FrozenCollectionData<H> {
+    pub fn new(dict: Bound<PyDict>, hash: isize) -> Self {
+        Self {
+            data: dict.unbind(),
+            hash: H::new_eager(hash),
+        }
+    }
+
+    pub fn new_lazy(dict: Bound<PyDict>) -> Self {
+        Self {
+            data: dict.unbind(),
+            hash: H::new_lazy(),
+        }
+    }
+
+    pub fn get_hash(
+        &self,
+        py: Python,
+        compute: fn(&Bound<PyDict>) -> PyResult<isize>,
+    ) -> PyResult<isize> {
+        self.hash.get(&self.data.bind_borrowed(py), compute)
+    }
+
+    pub fn len(&self, py: Python) -> usize {
+        self.data.bind_borrowed(py).len()
+    }
+
+    pub fn contains(&self, key: &Bound<PyAny>) -> PyResult<bool> {
+        self.data.bind_borrowed(key.py()).contains(key)
+    }
+
+    pub fn iter<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyIterator>> {
+        self.data.as_any().bind_borrowed(py).try_iter()
+    }
+
+    pub fn reversed<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyIterator>> {
+        let keys = self.data.bind_borrowed(py).keys();
+        keys.reverse()?;
+        keys.try_iter()
+    }
+}
+
+pub fn xor_hash_keys(dict: &Bound<PyDict>) -> PyResult<isize> {
+    let mut h: isize = 0;
+    for key in dict.keys() {
+        h ^= key.hash()?;
+    }
+    Ok(h)
+}