rustpython_vm/builtins/

set.rs

/*
 * Builtin set type with a sequence of unique items.
 */
use super::{
    IterStatus, PositionIterInternal, PyDict, PyDictRef, PyGenericAlias, PyTupleRef, PyType,
    PyTypeRef, builtins_iter, locked_step,
};
use crate::{
    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
    atomic_func,
    class::{PyClassDef, PyClassImpl},
    common::{
        ascii,
        hash::PyHash,
        lock::{LazyLock, PyMutex},
        rc::PyRc,
        wtf8::Wtf8Buf,
    },
    convert::ToPyResult,
    dict_inner::{self, DictSize},
    function::{
        ArgIterable, FuncArgs, NameOthers, OptionalArg, PosArgs, PyArithmeticValue,
        PyComparisonValue,
    },
    protocol::{PyIterReturn, PyNumberMethods, PySequenceMethods},
    recursion::ReprGuard,
    types::AsNumber,
    types::{
        AsSequence, Comparable, Constructor, DefaultConstructor, Hashable, Initializer, IterNext,
        Iterable, PyComparisonOp, Representable, SelfIter,
    },
    utils::collection_repr,
    vm::VirtualMachine,
};
use core::{borrow::Borrow, fmt};
use rustpython_common::{
    atomic::{Ordering, PyAtomic, Radium},
    hash,
};

pub(crate) type SetContentType = dict_inner::Dict<()>;

#[pyclass(module = false, name = "set", unhashable = true, traverse)]
#[derive(Default)]
pub struct PySet {
    pub(super) inner: PySetInner,
}

impl PySet {
    #[deprecated(note = "Use `PySet::default().into_ref(ctx)` instead")]
    pub fn new_ref(ctx: &Context) -> PyRef<Self> {
        Self::default().into_ref(ctx)
    }

    #[must_use]
    pub fn elements(&self) -> Vec<PyObjectRef> {
        self.inner.elements()
    }

    fn fold_op(
        &self,
        others: impl core::iter::Iterator<Item = ArgIterable>,
        op: fn(&PySetInner, ArgIterable, &VirtualMachine) -> PyResult<PySetInner>,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        Ok(Self {
            inner: self.inner.fold_op(others, op, vm)?,
        })
    }

    fn op(
        &self,
        other: AnySet,
        op: fn(&PySetInner, ArgIterable, &VirtualMachine) -> PyResult<PySetInner>,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        Ok(Self {
            inner: self
                .inner
                .fold_op(core::iter::once(other.into_iterable(vm)?), op, vm)?,
        })
    }
}

#[pyclass(module = false, name = "frozenset", unhashable = true)]
pub struct PyFrozenSet {
    inner: PySetInner,
    hash: PyAtomic<PyHash>,
}

impl Default for PyFrozenSet {
    fn default() -> Self {
        Self {
            inner: PySetInner::default(),
            hash: hash::SENTINEL.into(),
        }
    }
}

impl PyFrozenSet {
    // Also used by ssl.rs windows.
    pub fn from_iter(
        vm: &VirtualMachine,
        it: impl IntoIterator<Item = PyObjectRef>,
    ) -> PyResult<Self> {
        let inner = PySetInner::default();
        for elem in it {
            inner.add(&elem, vm)?;
        }
        // FIXME: empty set check
        Ok(Self {
            inner,
            ..Default::default()
        })
    }

    pub fn elements(&self) -> Vec<PyObjectRef> {
        self.inner.elements()
    }

    fn fold_op(
        &self,
        others: impl core::iter::Iterator<Item = ArgIterable>,
        op: fn(&PySetInner, ArgIterable, &VirtualMachine) -> PyResult<PySetInner>,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        Ok(Self {
            inner: self.inner.fold_op(others, op, vm)?,
            ..Default::default()
        })
    }

    fn op(
        &self,
        other: AnySet,
        op: fn(&PySetInner, ArgIterable, &VirtualMachine) -> PyResult<PySetInner>,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        Ok(Self {
            inner: self
                .inner
                .fold_op(core::iter::once(other.into_iterable(vm)?), op, vm)?,
            ..Default::default()
        })
    }
}

impl fmt::Debug for PySet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // TODO: implement more detailed, non-recursive Debug formatter
        f.write_str("set")
    }
}

impl fmt::Debug for PyFrozenSet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // TODO: implement more detailed, non-recursive Debug formatter
        f.write_str("PyFrozenSet ")?;
        f.debug_set().entries(self.elements().iter()).finish()
    }
}

impl PyPayload for PySet {
    #[inline]
    fn class(ctx: &Context) -> &'static Py<PyType> {
        ctx.types.set_type
    }
}

impl PyPayload for PyFrozenSet {
    #[inline]
    fn class(ctx: &Context) -> &'static Py<PyType> {
        ctx.types.frozenset_type
    }
}

#[derive(Default, Clone)]
pub(super) struct PySetInner {
    content: PyRc<SetContentType>,
}

unsafe impl crate::object::Traverse for PySetInner {
    fn traverse(&self, tracer_fn: &mut crate::object::TraverseFn<'_>) {
        // FIXME(discord9): Rc means shared ref, so should it be traced?
        self.content.traverse(tracer_fn)
    }
}

impl PySetInner {
    pub(super) fn from_iter<T>(iter: T, vm: &VirtualMachine) -> PyResult<Self>
    where
        T: IntoIterator<Item = PyResult<PyObjectRef>>,
    {
        let set = Self::default();
        for item in iter {
            let item = item?;
            set.add(&item, vm)?;
        }
        Ok(set)
    }

    /// Build a set from an arbitrary object, reusing stored hashes when the
    /// source is a set/frozenset/dict.
    fn from_object(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
        let set = Self::default();
        set.update_internal(iterable, vm)?;
        Ok(set)
    }

    /// Elements of `obj` with their stored hashes, or `None` if `obj` keeps
    /// none and must be iterated generically. Mirrors the `PyAnySet_Check` /
    /// `PyDict_CheckExact` fast paths in CPython's `set_update_internal`.
    fn cached_hashes(obj: &PyObject, vm: &VirtualMachine) -> Option<Vec<(PyObjectRef, PyHash)>> {
        if let Some(set) = extract_set(obj) {
            Some(set.content.keys_with_hashes())
        } else {
            obj.downcast_ref_if_exact::<PyDict>(vm)
                .map(|dict| dict._as_dict_inner().keys_with_hashes())
        }
    }

    fn fold_op<O>(
        &self,
        others: impl core::iter::Iterator<Item = O>,
        op: fn(&Self, O, &VirtualMachine) -> PyResult<Self>,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        let mut res = self.copy();
        for other in others {
            res = op(&res, other, vm)?;
        }
        Ok(res)
    }

    fn intersection_multi(
        &self,
        mut others: impl core::iter::Iterator<Item = ArgIterable>,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        let Some(other) = others.next() else {
            return Ok(self.copy());
        };
        let mut result = self.intersection(other, vm)?;
        for other in others {
            result = result.intersection(other, vm)?;
        }
        Ok(result)
    }

    fn difference_multi(
        &self,
        mut others: impl core::iter::Iterator<Item = ArgIterable>,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        let Some(other) = others.next() else {
            return Ok(self.copy());
        };
        let result = self.difference_new(other, vm)?;
        result.difference_update(others, vm)?;
        Ok(result)
    }

    fn len(&self) -> usize {
        self.content.len()
    }

    fn sizeof(&self) -> usize {
        self.content.sizeof()
    }

    fn copy(&self) -> Self {
        Self {
            content: PyRc::new((*self.content).clone()),
        }
    }

    fn contains(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
        let result = self
            .retry_op_with_frozenset(needle, vm, |needle, vm| self.content.contains(vm, needle));
        Self::wrap_unhashable_error(result, needle, vm)
    }

    /// Look up a key whose hash is already known, without a frozenset retry.
    fn contains_known_hash(
        &self,
        needle: &PyObject,
        hash: PyHash,
        vm: &VirtualMachine,
    ) -> PyResult<bool> {
        self.content.contains_known_hash(vm, needle, hash)
    }

    fn compare(&self, other: &Self, op: PyComparisonOp, vm: &VirtualMachine) -> PyResult<bool> {
        if op == PyComparisonOp::Ne {
            return self.compare(other, PyComparisonOp::Eq, vm).map(|eq| !eq);
        }
        if !op.eval_ord(self.len().cmp(&other.len())) {
            return Ok(false);
        }

        let (superset, subset) = match op {
            PyComparisonOp::Lt | PyComparisonOp::Le | PyComparisonOp::Eq => (other, self),
            _ => (self, other),
        };

        for (key, hash) in subset.content.keys_with_hashes() {
            if !superset.contains_known_hash(&key, hash, vm)? {
                return Ok(false);
            }
        }
        Ok(true)
    }

    pub(super) fn union(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
        let set = self.clone();
        if let Some(elements) = Self::cached_hashes(other.as_object(), vm) {
            for (item, hash) in elements {
                set.add_known_hash(&item, hash, vm)?;
            }
            return Ok(set);
        }
        for item in other.iter(vm)? {
            let item = item?;
            set.add(&item, vm)?;
        }

        Ok(set)
    }

    pub(super) fn intersection(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
        if let Some(other_set) = extract_set(other.as_object()) {
            return self.intersection_set(other_set, vm);
        }
        let set = Self::default();
        for item in other.iter(vm)? {
            let obj = item?;
            let hash = obj.hash(vm)?;
            if self.contains_known_hash(&obj, hash, vm)? {
                set.add_known_hash(&obj, hash, vm)?;
                if set.len() >= self.len() {
                    break;
                }
            }
        }
        Ok(set)
    }

    fn intersection_set(&self, other: &Self, vm: &VirtualMachine) -> PyResult<Self> {
        if PyRc::ptr_eq(&self.content, &other.content) {
            return Ok(self.copy());
        }
        let (target, source) = if self.len() < other.len() {
            (other, self)
        } else {
            (self, other)
        };
        let set = Self::default();
        for (obj, hash) in source.content.keys_with_hashes() {
            if target.contains_known_hash(&obj, hash, vm)? {
                set.add_known_hash(&obj, hash, vm)?;
            }
        }
        Ok(set)
    }

    fn difference_new(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
        // Scanning the left side avoids visiting and retaining a much larger
        // exclusion collection. Match CPython's comparison direction as well.
        if let Some(other_set) = extract_set(other.as_object()) {
            if self.len() >> 2 <= other_set.len() {
                return self
                    .difference_by(|key, hash| other_set.contains_known_hash(key, hash, vm), vm);
            }
        } else if let Some(dict) = other.as_object().downcast_ref_if_exact::<PyDict>(vm)
            && self.len() >> 2 <= dict._as_dict_inner().len()
        {
            return self.difference_by(
                |key, hash| dict._as_dict_inner().contains_known_hash(vm, key, hash),
                vm,
            );
        }
        self.copy().difference(other, vm)
    }

    // The dict-view caller supplies a private working table.
    pub(super) fn difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<Self> {
        self.difference_update(core::iter::once(other), vm)?;
        Ok(self.clone())
    }

    fn difference_by(
        &self,
        contains: impl Fn(&PyObject, PyHash) -> PyResult<bool>,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        let result = Self::default();
        for (key, hash) in self.content.keys_with_hashes() {
            if !contains(&key, hash)? {
                result.add_known_hash(&key, hash, vm)?;
            }
        }
        Ok(result)
    }

    pub(super) fn symmetric_difference(
        &self,
        other: ArgIterable,
        vm: &VirtualMachine,
    ) -> PyResult<Self> {
        let new_inner = self.clone();

        if let Some(elements) = Self::cached_hashes(other.as_object(), vm) {
            // the source is already duplicate-free
            for (item, hash) in elements {
                new_inner
                    .content
                    .delete_or_insert_known_hash(vm, &item, hash, ())?;
            }
            return Ok(new_inner);
        }

        // We want to remove duplicates in other
        let other_set = Self::from_iter(other.iter(vm)?, vm)?;

        for (item, hash) in other_set.content.keys_with_hashes() {
            new_inner
                .content
                .delete_or_insert_known_hash(vm, &item, hash, ())?;
        }

        Ok(new_inner)
    }

    fn issuperset(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        if let Some(other_set) = extract_set(other.as_object()) {
            return self.compare(other_set, PyComparisonOp::Ge, vm);
        }
        for item in other.iter(vm)? {
            if !self.contains(&*item?, vm)? {
                return Ok(false);
            }
        }
        Ok(true)
    }

    fn issubset(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        if let Some(other_set) = extract_set(other.as_object()) {
            return self.compare(other_set, PyComparisonOp::Le, vm);
        }
        Ok(self.intersection(other, vm)?.len() == self.len())
    }

    pub(super) fn isdisjoint(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        if let Some(other_set) = extract_set(other.as_object()) {
            if core::ptr::eq(self, other_set) {
                return Ok(self.len() == 0);
            }
            let other_type = other.as_object().class();
            if other_type.is(vm.ctx.types.set_type) || other_type.is(vm.ctx.types.frozenset_type) {
                let (target, source) = if self.len() < other_set.len() {
                    (other_set, self)
                } else {
                    (self, other_set)
                };
                for (key, hash) in source.content.keys_with_hashes() {
                    if target.contains_known_hash(&key, hash, vm)? {
                        return Ok(false);
                    }
                }
                return Ok(true);
            }
        }
        for item in other.iter(vm)? {
            if self.contains(&*item?, vm)? {
                return Ok(false);
            }
        }
        Ok(true)
    }

    fn repr(&self, class_name: Option<&str>, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
        let empty = format!("{}()", class_name.unwrap_or("set"));
        collection_repr(
            class_name,
            "{",
            "}",
            &empty,
            self.elements().iter().map(|o| &**o),
            vm,
        )
    }

    fn add(&self, item: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
        let result = self.content.insert(vm, item, ());
        Self::wrap_unhashable_error(result, item, vm)
    }

    /// [`Self::add`] with a known hash.
    fn add_known_hash(&self, item: &PyObject, hash: PyHash, vm: &VirtualMachine) -> PyResult<()> {
        let result = self.content.insert_known_hash(vm, item, hash, ());
        Self::wrap_unhashable_error(result, item, vm)
    }

    fn remove(&self, item: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
        let result =
            self.retry_op_with_frozenset(item, vm, |item, vm| self.content.delete(vm, item));
        Self::wrap_unhashable_error(result, item, vm)
    }

    fn discard(&self, item: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
        let result = self
            .retry_op_with_frozenset(item, vm, |item, vm| self.content.delete_if_exists(vm, item));
        Self::wrap_unhashable_error(result, item, vm)
    }

    fn clear(&self) {
        self.content.clear()
    }

    fn elements(&self) -> Vec<PyObjectRef> {
        self.content.keys()
    }

    fn pop(&self, vm: &VirtualMachine) -> PyResult {
        // TODO: should be pop_front, but that requires rearranging every index
        if let Some((key, _)) = self.content.pop_back() {
            Ok(key)
        } else {
            let err_msg = vm.ctx.new_str(ascii!("pop from an empty set")).into();
            Err(vm.new_key_error(err_msg))
        }
    }

    fn update_internal(&self, iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
        // check AnySet
        if let Ok(any_set) = AnySet::try_from_object(vm, iterable.to_owned()) {
            self.merge_set(any_set, vm)
        // check Dict
        } else if let Ok(dict) = iterable.to_owned().downcast_exact::<PyDict>(vm) {
            self.merge_dict(&dict, vm)
        } else {
            // add iterable that is not AnySet or Dict
            for item in iterable.try_into_value::<ArgIterable>(vm)?.iter(vm)? {
                let item = item?;
                self.add(&item, vm)?;
            }
            Ok(())
        }
    }

    fn merge_set(&self, any_set: AnySet, vm: &VirtualMachine) -> PyResult<()> {
        for (item, hash) in any_set.as_inner().content.keys_with_hashes() {
            self.add_known_hash(&item, hash, vm)?;
        }
        Ok(())
    }

    fn merge_dict(&self, dict: &Py<PyDict>, vm: &VirtualMachine) -> PyResult<()> {
        for (key, hash) in dict._as_dict_inner().keys_with_hashes() {
            self.add_known_hash(&key, hash, vm)?;
        }
        Ok(())
    }

    fn intersection_update(
        &self,
        others: impl core::iter::Iterator<Item = ArgIterable>,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        let temp_inner = self.intersection_multi(others, vm)?;
        let content = PyRc::try_unwrap(temp_inner.content).unwrap_or_else(|table| (*table).clone());
        self.content.replace_contents(content);
        Ok(())
    }

    fn difference_update(
        &self,
        others: impl core::iter::Iterator<Item = ArgIterable>,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        for iterable in others {
            let elements = if let Some(other_set) = extract_set(iterable.as_object()) {
                if PyRc::ptr_eq(&self.content, &other_set.content) {
                    self.clear();
                    continue;
                }
                // Build the intersection first: besides bounding the work by
                // our size, this preserves CPython's equality/error ordering.
                Some(if other_set.len() >> 3 > self.len() {
                    self.intersection_set(other_set, vm)?
                        .content
                        .keys_with_hashes()
                } else {
                    other_set.content.keys_with_hashes()
                })
            } else {
                Self::cached_hashes(iterable.as_object(), vm)
            };
            if let Some(elements) = elements {
                for (item, hash) in elements {
                    self.content.delete_if_exists_known_hash(vm, &*item, hash)?;
                }
                continue;
            }
            for item in iterable.iter(vm)? {
                self.content.delete_if_exists(vm, &*item?)?;
            }
        }
        Ok(())
    }

    fn symmetric_difference_update(
        &self,
        others: impl core::iter::Iterator<Item = ArgIterable>,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        for iterable in others {
            if let Some(elements) = Self::cached_hashes(iterable.as_object(), vm) {
                // the source is already duplicate-free
                for (item, hash) in elements {
                    self.content
                        .delete_or_insert_known_hash(vm, &item, hash, ())?;
                }
                continue;
            }
            // We want to remove duplicates in iterable
            let iterable_set = Self::from_iter(iterable.iter(vm)?, vm)?;
            for (item, hash) in iterable_set.content.keys_with_hashes() {
                self.content
                    .delete_or_insert_known_hash(vm, &item, hash, ())?;
            }
        }
        Ok(())
    }

    fn hash(&self) -> PyHash {
        let hasher = self.content.fold_hashes(
            hash::FrozenSetHash::new(self.len()),
            |mut hasher, element_hash| {
                hasher.add(element_hash);
                hasher
            },
        );
        hasher.finish()
    }

    // Run operation, on failure, if item is a set/set subclass, convert it
    // into a frozenset and try the operation again. Propagates original error
    // on failure to convert and restores item in KeyError on failure (remove).
    fn retry_op_with_frozenset<T, F>(
        &self,
        item: &PyObject,
        vm: &VirtualMachine,
        op: F,
    ) -> PyResult<T>
    where
        F: Fn(&PyObject, &VirtualMachine) -> PyResult<T>,
    {
        op(item, vm).or_else(|original_err| {
            item.downcast_ref::<PySet>()
                // Keep original error around.
                .ok_or(original_err)
                .and_then(|set| {
                    op(
                        &PyFrozenSet {
                            inner: set.inner.copy(),
                            ..Default::default()
                        }
                        .into_pyobject(vm),
                        vm,
                    )
                    // If operation raised KeyError, report original set (set.remove)
                    .map_err(|op_err| {
                        if op_err.fast_isinstance(vm.ctx.exceptions.key_error) {
                            vm.new_key_error(item.to_owned())
                        } else {
                            op_err
                        }
                    })
                })
        })
    }

    fn wrap_unhashable_error<T>(
        result: PyResult<T>,
        item: &PyObject,
        vm: &VirtualMachine,
    ) -> PyResult<T> {
        match result {
            Err(cause) if cause.fast_isinstance(vm.ctx.exceptions.type_error) => {
                let message = cause.as_object().str(vm)?;
                let err = vm.new_type_error(format!(
                    "cannot use '{}' as a set element ({message})",
                    item.class().name()
                ));
                err.set_cause(Some(cause));
                Err(err)
            }
            result => result,
        }
    }
}

fn extract_set(obj: &PyObject) -> Option<&PySetInner> {
    match_class!(match obj {
        ref set @ PySet => Some(&set.inner),
        ref frozen @ PyFrozenSet => Some(&frozen.inner),
        _ => None,
    })
}

/// Elements of `obj` with their stored hashes, or `None` unless `obj` is exactly
/// a `set` or `frozenset` — `PyAnySet_CheckExact`, where [`extract_set`] is the
/// subclass-inclusive `PyAnySet_Check`.
pub(super) fn exact_set_keys_with_hashes(
    obj: &PyObject,
    vm: &VirtualMachine,
) -> Option<Vec<(PyObjectRef, PyHash)>> {
    let inner = obj
        .downcast_ref_if_exact::<PySet>(vm)
        .map(|set| &set.inner)
        .or_else(|| {
            obj.downcast_ref_if_exact::<PyFrozenSet>(vm)
                .map(|frozen| &frozen.inner)
        })?;
    Some(inner.content.keys_with_hashes())
}

fn reduce_set(zelf: &PyObject, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
    (
        zelf.class().to_owned(),
        #[expect(clippy::or_fun_call, reason = "changing this won't compile")]
        vm.new_tuple((extract_set(zelf)
            .unwrap_or(&PySetInner::default())
            .elements(),)),
        zelf.dict(),
    )
}

impl PySet {
    fn __len__(&self) -> usize {
        self.inner.len()
    }

    pub fn contains(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
        self.inner.contains(needle, vm)
    }

    fn __or__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(self.op(
                other,
                PySetInner::union,
                vm,
            )?))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __and__(
        &self,
        other: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(Self {
                inner: self.inner.intersection(other.into_iterable(vm)?, vm)?,
            }))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __sub__(
        &self,
        other: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(Self {
                inner: self.inner.difference_new(other.into_iterable(vm)?, vm)?,
            }))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __rsub__(
        zelf: PyRef<Self>,
        other: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(Self {
                inner: other
                    .as_inner()
                    .difference_new(ArgIterable::try_from_object(vm, zelf.into())?, vm)?,
            }))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __xor__(
        &self,
        other: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(self.op(
                other,
                PySetInner::symmetric_difference,
                vm,
            )?))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __ior__(zelf: PyRef<Self>, set: AnySet, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
        zelf.inner.merge_set(set, vm)?;
        Ok(zelf)
    }

    fn __iand__(zelf: PyRef<Self>, set: AnySet, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
        if !set.is(zelf.as_object()) {
            zelf.inner
                .intersection_update(core::iter::once(set.into_iterable(vm)?), vm)?;
        }
        Ok(zelf)
    }

    fn __isub__(zelf: PyRef<Self>, set: AnySet, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
        if set.is(zelf.as_object()) {
            zelf.inner.clear();
        } else {
            zelf.inner
                .difference_update(set.into_iterable_iter(vm)?, vm)?;
        }
        Ok(zelf)
    }

    fn __ixor__(zelf: PyRef<Self>, set: AnySet, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
        if set.is(zelf.as_object()) {
            zelf.inner.clear();
        } else {
            zelf.inner
                .symmetric_difference_update(set.into_iterable_iter(vm)?, vm)?;
        }
        Ok(zelf)
    }

    pub fn add(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
        self.inner.add(&object, vm)
    }
}

#[pyclass(
    with(
        Constructor,
        Initializer,
        AsSequence,
        Comparable,
        Iterable,
        AsNumber,
        Representable
    ),
    flags(BASETYPE, _MATCH_SELF, HAS_WEAKREF)
)]
impl Py<PySet> {
    #[pymethod(coexist)]
    fn __contains__(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
        self.contains(&object, vm)
    }

    #[pymethod]
    fn __sizeof__(&self) -> usize {
        core::mem::size_of::<PySet>() + self.inner.sizeof()
    }

    #[pymethod]
    fn copy(&self) -> PySet {
        PySet {
            inner: self.inner.copy(),
        }
    }

    #[pymethod]
    fn union(
        &self,
        others: PosArgs<ArgIterable, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<PySet> {
        self.fold_op(others.into_iter(), PySetInner::union, vm)
    }

    #[pymethod]
    fn intersection(
        &self,
        others: PosArgs<ArgIterable, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<PySet> {
        Ok(PySet {
            inner: self.inner.intersection_multi(others.into_iter(), vm)?,
        })
    }

    #[pymethod]
    fn difference(
        &self,
        others: PosArgs<ArgIterable, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<PySet> {
        Ok(PySet {
            inner: self.inner.difference_multi(others.into_iter(), vm)?,
        })
    }

    #[pymethod]
    fn symmetric_difference(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<PySet> {
        self.fold_op(
            core::iter::once(other),
            PySetInner::symmetric_difference,
            vm,
        )
    }

    #[pymethod]
    fn issubset(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        self.inner.issubset(other, vm)
    }

    #[pymethod]
    fn issuperset(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        self.inner.issuperset(other, vm)
    }

    #[pymethod]
    fn isdisjoint(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        self.inner.isdisjoint(other, vm)
    }

    #[pymethod]
    pub fn add(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
        self.payload.add(object, vm)
    }

    #[pymethod]
    fn remove(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
        self.inner.remove(&object, vm)
    }

    #[pymethod]
    pub fn discard(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
        self.inner.discard(&object, vm).map(|_| ())
    }

    #[pymethod]
    pub fn clear(&self) {
        self.inner.clear()
    }

    #[pymethod]
    pub fn pop(&self, vm: &VirtualMachine) -> PyResult {
        self.inner.pop(vm)
    }

    #[pymethod]
    fn update(
        &self,
        others: PosArgs<PyObjectRef, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        for iterable in others {
            self.inner.update_internal(iterable, vm)?;
        }
        Ok(())
    }

    #[pymethod]
    fn intersection_update(
        &self,
        others: PosArgs<ArgIterable, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        self.inner.intersection_update(others.into_iter(), vm)?;
        Ok(())
    }

    #[pymethod]
    fn difference_update(
        &self,
        others: PosArgs<ArgIterable, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<()> {
        self.inner.difference_update(others.into_iter(), vm)
    }

    #[pymethod]
    fn symmetric_difference_update(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<()> {
        self.inner
            .symmetric_difference_update(core::iter::once(other), vm)
    }

    #[pymethod]
    fn __reduce__(
        zelf: PyRef<PySet>,
        vm: &VirtualMachine,
    ) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
        reduce_set(zelf.as_ref(), vm)
    }

    #[pyclassmethod]
    fn __class_getitem__(
        cls: PyTypeRef,
        object: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyGenericAlias> {
        PyGenericAlias::from_args(cls, object, vm)
    }
}

impl DefaultConstructor for PySet {}

impl Initializer for PySet {
    type Args = crate::function::PositionalIterable;

    fn init(zelf: &Py<Self>, args: Self::Args, vm: &VirtualMachine) -> PyResult<()> {
        zelf.clear();
        if let OptionalArg::Present(it) = args.iterable {
            zelf.update(PosArgs::<PyObjectRef, NameOthers>::named(vec![it]), vm)?;
        }
        Ok(())
    }
}

impl AsSequence for PySet {
    fn as_sequence() -> &'static PySequenceMethods {
        static AS_SEQUENCE: LazyLock<PySequenceMethods> = LazyLock::new(|| PySequenceMethods {
            length: atomic_func!(|seq, _vm| Ok(PySet::sequence_downcast(seq).__len__())),
            contains: atomic_func!(
                |seq, needle, vm| PySet::sequence_downcast(seq).contains(needle, vm)
            ),
            ..PySequenceMethods::NOT_IMPLEMENTED
        });
        &AS_SEQUENCE
    }
}

impl Comparable for PySet {
    fn cmp(
        zelf: &crate::Py<Self>,
        other: &PyObject,
        op: PyComparisonOp,
        vm: &VirtualMachine,
    ) -> PyResult<PyComparisonValue> {
        extract_set(other).map_or(Ok(PyComparisonValue::NotImplemented), |other| {
            Ok(zelf.inner.compare(other, op, vm)?.into())
        })
    }
}

impl Iterable for PySet {
    fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
        Ok(PySetIterator::new(AnySet {
            object: zelf.into(),
        })
        .into_pyobject(vm))
    }
}

impl AsNumber for PySet {
    fn as_number() -> &'static PyNumberMethods {
        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
            // Binary ops check both operands are sets (like CPython's set_sub, etc.)
            // This is needed because __rsub__ swaps operands: a.__rsub__(b) calls subtract(b, a)
            subtract: Some(|a, b, vm| {
                if !AnySet::check(a, vm) || !AnySet::check(b, vm) {
                    return Ok(vm.ctx.not_implemented());
                }
                if let Some(a) = a.downcast_ref::<PySet>() {
                    a.__sub__(b.to_owned(), vm).to_pyresult(vm)
                } else if let Some(a) = a.downcast_ref::<PyFrozenSet>() {
                    // When called via __rsub__, a might be PyFrozenSet
                    a.__sub__(b.to_owned(), vm)
                        .map(|r| r.map(|s| PySet { inner: s.inner }))
                        .to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            and: Some(|a, b, vm| {
                if !AnySet::check(a, vm) || !AnySet::check(b, vm) {
                    return Ok(vm.ctx.not_implemented());
                }
                if let Some(a) = a.downcast_ref::<PySet>() {
                    a.__and__(b.to_owned(), vm).to_pyresult(vm)
                } else if let Some(a) = a.downcast_ref::<PyFrozenSet>() {
                    a.__and__(b.to_owned(), vm)
                        .map(|r| r.map(|s| PySet { inner: s.inner }))
                        .to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            xor: Some(|a, b, vm| {
                if !AnySet::check(a, vm) || !AnySet::check(b, vm) {
                    return Ok(vm.ctx.not_implemented());
                }
                if let Some(a) = a.downcast_ref::<PySet>() {
                    a.__xor__(b.to_owned(), vm).to_pyresult(vm)
                } else if let Some(a) = a.downcast_ref::<PyFrozenSet>() {
                    a.__xor__(b.to_owned(), vm)
                        .map(|r| r.map(|s| PySet { inner: s.inner }))
                        .to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            or: Some(|a, b, vm| {
                if !AnySet::check(a, vm) || !AnySet::check(b, vm) {
                    return Ok(vm.ctx.not_implemented());
                }
                if let Some(a) = a.downcast_ref::<PySet>() {
                    a.__or__(b.to_owned(), vm).to_pyresult(vm)
                } else if let Some(a) = a.downcast_ref::<PyFrozenSet>() {
                    a.__or__(b.to_owned(), vm)
                        .map(|r| r.map(|s| PySet { inner: s.inner }))
                        .to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            inplace_subtract: Some(|a, b, vm| {
                if let Some(a) = a.downcast_ref::<PySet>() {
                    PySet::__isub__(a.to_owned(), AnySet::try_from_object(vm, b.to_owned())?, vm)
                        .to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            inplace_and: Some(|a, b, vm| {
                if let Some(a) = a.downcast_ref::<PySet>() {
                    PySet::__iand__(a.to_owned(), AnySet::try_from_object(vm, b.to_owned())?, vm)
                        .to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            inplace_xor: Some(|a, b, vm| {
                if let Some(a) = a.downcast_ref::<PySet>() {
                    PySet::__ixor__(a.to_owned(), AnySet::try_from_object(vm, b.to_owned())?, vm)
                        .to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            inplace_or: Some(|a, b, vm| {
                if let Some(a) = a.downcast_ref::<PySet>() {
                    PySet::__ior__(a.to_owned(), AnySet::try_from_object(vm, b.to_owned())?, vm)
                        .to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            ..PyNumberMethods::NOT_IMPLEMENTED
        };
        &AS_NUMBER
    }
}

impl Representable for PySet {
    #[inline]
    fn repr_wtf8(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
        let class = zelf.class();
        let borrowed_name = class.name();
        let class_name = &*borrowed_name;

        if zelf.inner.len() == 0 {
            return Ok(Wtf8Buf::from(format!("{class_name}()")));
        }

        if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
            let name = (class_name != "set").then_some(class_name);
            zelf.inner.repr(name, vm)
        } else {
            Ok(Wtf8Buf::from(format!("{class_name}(...)")))
        }
    }
}

impl Constructor for PyFrozenSet {
    type Args = crate::function::PositionalIterable;

    fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
        let is_exact_frozenset = cls.is(vm.ctx.types.frozenset_type);
        let is_frozenset_init = {
            let cls_init = cls
                .slots
                .init
                .load()
                .map(|init| crate::types::fn_addr(init));
            let frozenset_init = vm
                .ctx
                .types
                .frozenset_type
                .slots
                .init
                .load()
                .map(|init| crate::types::fn_addr(init));
            cls_init == frozenset_init
        };

        // Optimizations for exact frozenset type
        let iterable_opt = if is_exact_frozenset || is_frozenset_init {
            let iterable: crate::function::PositionalIterable = args.bind_for(vm, Self::NAME)?;
            let iterable = iterable.iterable;

            // Return exact frozenset as-is
            if is_exact_frozenset
                && let OptionalArg::Present(input) = &iterable
                && input.class().is(vm.ctx.types.frozenset_type)
            {
                return Ok(input.clone());
            }

            iterable
        } else {
            match &args.args[..] {
                [] => OptionalArg::Missing,
                [iterable] => OptionalArg::Present(iterable.clone()),
                slice => {
                    return Err(vm.new_arity_type_error(Self::NAME, 0..=1, slice.len()));
                }
            }
        };

        let payload = Self::py_new(
            &cls,
            Self::Args {
                iterable: iterable_opt,
            },
            vm,
        )?;

        // Return empty frozenset singleton
        if is_exact_frozenset && payload.inner.len() == 0 {
            return Ok(vm.ctx.empty_frozenset.clone().into());
        }

        payload.into_ref_with_type(vm, cls).map(Into::into)
    }

    fn py_new(_cls: &Py<PyType>, args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
        let inner = match args.iterable {
            OptionalArg::Present(iterable) => PySetInner::from_object(iterable, vm)?,
            OptionalArg::Missing => PySetInner::default(),
        };
        Ok(Self {
            inner,
            ..Default::default()
        })
    }
}

impl PyFrozenSet {
    fn __len__(&self) -> usize {
        self.inner.len()
    }

    pub fn contains(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
        self.inner.contains(needle, vm)
    }

    fn __or__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(set) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(self.op(
                set,
                PySetInner::union,
                vm,
            )?))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __and__(
        &self,
        other: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(Self {
                inner: self.inner.intersection(other.into_iterable(vm)?, vm)?,
                ..Default::default()
            }))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __sub__(
        &self,
        other: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(Self {
                inner: self.inner.difference_new(other.into_iterable(vm)?, vm)?,
                ..Default::default()
            }))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __rsub__(
        zelf: PyRef<Self>,
        other: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(Self {
                inner: other
                    .as_inner()
                    .difference_new(ArgIterable::try_from_object(vm, zelf.into())?, vm)?,
                ..Default::default()
            }))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }

    fn __xor__(
        &self,
        other: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyArithmeticValue<Self>> {
        if let Ok(other) = AnySet::try_from_object(vm, other) {
            Ok(PyArithmeticValue::Implemented(self.op(
                other,
                PySetInner::symmetric_difference,
                vm,
            )?))
        } else {
            Ok(PyArithmeticValue::NotImplemented)
        }
    }
}

#[pyclass(
    flags(BASETYPE, _MATCH_SELF, HAS_WEAKREF),
    with(
        Constructor,
        AsSequence,
        Hashable,
        Comparable,
        Iterable,
        AsNumber,
        Representable
    )
)]
impl Py<PyFrozenSet> {
    #[pymethod(coexist)]
    fn __contains__(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
        self.contains(&object, vm)
    }

    #[pymethod]
    fn __sizeof__(&self) -> usize {
        core::mem::size_of::<PyFrozenSet>() + self.inner.sizeof()
    }

    #[pymethod]
    fn copy(zelf: PyRef<PyFrozenSet>, vm: &VirtualMachine) -> PyRef<PyFrozenSet> {
        if zelf.class().is(vm.ctx.types.frozenset_type) {
            zelf
        } else {
            PyFrozenSet {
                inner: zelf.inner.copy(),
                ..Default::default()
            }
            .into_ref(&vm.ctx)
        }
    }

    #[pymethod]
    fn union(
        &self,
        others: PosArgs<ArgIterable, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<PyFrozenSet> {
        self.fold_op(others.into_iter(), PySetInner::union, vm)
    }

    #[pymethod]
    fn intersection(
        &self,
        others: PosArgs<ArgIterable, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<PyFrozenSet> {
        Ok(PyFrozenSet {
            inner: self.inner.intersection_multi(others.into_iter(), vm)?,
            ..Default::default()
        })
    }

    #[pymethod]
    fn difference(
        &self,
        others: PosArgs<ArgIterable, NameOthers>,
        vm: &VirtualMachine,
    ) -> PyResult<PyFrozenSet> {
        Ok(PyFrozenSet {
            inner: self.inner.difference_multi(others.into_iter(), vm)?,
            ..Default::default()
        })
    }

    #[pymethod]
    fn symmetric_difference(
        &self,
        other: ArgIterable,
        vm: &VirtualMachine,
    ) -> PyResult<PyFrozenSet> {
        self.fold_op(
            core::iter::once(other),
            PySetInner::symmetric_difference,
            vm,
        )
    }

    #[pymethod]
    fn issubset(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        self.inner.issubset(other, vm)
    }

    #[pymethod]
    fn issuperset(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        self.inner.issuperset(other, vm)
    }

    #[pymethod]
    fn isdisjoint(&self, other: ArgIterable, vm: &VirtualMachine) -> PyResult<bool> {
        self.inner.isdisjoint(other, vm)
    }

    #[pymethod]
    fn __reduce__(
        zelf: PyRef<PyFrozenSet>,
        vm: &VirtualMachine,
    ) -> (PyTypeRef, PyTupleRef, Option<PyDictRef>) {
        reduce_set(zelf.as_ref(), vm)
    }

    #[pyclassmethod]
    fn __class_getitem__(
        cls: PyTypeRef,
        object: PyObjectRef,
        vm: &VirtualMachine,
    ) -> PyResult<PyGenericAlias> {
        PyGenericAlias::from_args(cls, object, vm)
    }
}

impl AsSequence for PyFrozenSet {
    fn as_sequence() -> &'static PySequenceMethods {
        static AS_SEQUENCE: LazyLock<PySequenceMethods> = LazyLock::new(|| PySequenceMethods {
            length: atomic_func!(|seq, _vm| Ok(PyFrozenSet::sequence_downcast(seq).__len__())),
            contains: atomic_func!(
                |seq, needle, vm| PyFrozenSet::sequence_downcast(seq).contains(needle, vm)
            ),
            ..PySequenceMethods::NOT_IMPLEMENTED
        });
        &AS_SEQUENCE
    }
}

impl Hashable for PyFrozenSet {
    #[inline]
    fn hash(zelf: &crate::Py<Self>, _vm: &VirtualMachine) -> PyResult<PyHash> {
        let hash = match zelf.hash.load(Ordering::Relaxed) {
            hash::SENTINEL => {
                let hash = zelf.inner.hash();
                match Radium::compare_exchange(
                    &zelf.hash,
                    hash::SENTINEL,
                    hash::fix_sentinel(hash),
                    Ordering::Relaxed,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => hash,
                    Err(prev_stored) => prev_stored,
                }
            }
            hash => hash,
        };
        Ok(hash)
    }
}

impl Comparable for PyFrozenSet {
    fn cmp(
        zelf: &crate::Py<Self>,
        other: &PyObject,
        op: PyComparisonOp,
        vm: &VirtualMachine,
    ) -> PyResult<PyComparisonValue> {
        extract_set(other).map_or(Ok(PyComparisonValue::NotImplemented), |other| {
            Ok(zelf.inner.compare(other, op, vm)?.into())
        })
    }
}

impl Iterable for PyFrozenSet {
    fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
        Ok(PySetIterator::new(AnySet {
            object: zelf.into(),
        })
        .into_pyobject(vm))
    }
}

impl AsNumber for PyFrozenSet {
    fn as_number() -> &'static PyNumberMethods {
        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
            // Binary ops check both operands are sets (like CPython's set_sub, etc.)
            // __rsub__ swaps operands. Result type follows first operand's type.
            subtract: Some(|a, b, vm| {
                if !AnySet::check(a, vm) || !AnySet::check(b, vm) {
                    return Ok(vm.ctx.not_implemented());
                }
                if let Some(a) = a.downcast_ref::<PyFrozenSet>() {
                    a.__sub__(b.to_owned(), vm).to_pyresult(vm)
                } else if let Some(a) = a.downcast_ref::<PySet>() {
                    // When called via __rsub__, a might be PySet - return set (not frozenset)
                    a.__sub__(b.to_owned(), vm).to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            and: Some(|a, b, vm| {
                if !AnySet::check(a, vm) || !AnySet::check(b, vm) {
                    return Ok(vm.ctx.not_implemented());
                }
                if let Some(a) = a.downcast_ref::<PyFrozenSet>() {
                    a.__and__(b.to_owned(), vm).to_pyresult(vm)
                } else if let Some(a) = a.downcast_ref::<PySet>() {
                    a.__and__(b.to_owned(), vm).to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            xor: Some(|a, b, vm| {
                if !AnySet::check(a, vm) || !AnySet::check(b, vm) {
                    return Ok(vm.ctx.not_implemented());
                }
                if let Some(a) = a.downcast_ref::<PyFrozenSet>() {
                    a.__xor__(b.to_owned(), vm).to_pyresult(vm)
                } else if let Some(a) = a.downcast_ref::<PySet>() {
                    a.__xor__(b.to_owned(), vm).to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            or: Some(|a, b, vm| {
                if !AnySet::check(a, vm) || !AnySet::check(b, vm) {
                    return Ok(vm.ctx.not_implemented());
                }
                if let Some(a) = a.downcast_ref::<PyFrozenSet>() {
                    a.__or__(b.to_owned(), vm).to_pyresult(vm)
                } else if let Some(a) = a.downcast_ref::<PySet>() {
                    a.__or__(b.to_owned(), vm).to_pyresult(vm)
                } else {
                    Ok(vm.ctx.not_implemented())
                }
            }),
            ..PyNumberMethods::NOT_IMPLEMENTED
        };
        &AS_NUMBER
    }
}

impl Representable for PyFrozenSet {
    #[inline]
    fn repr_wtf8(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
        let inner = &zelf.inner;
        let class = zelf.class();
        let class_name = class.name();
        if inner.len() == 0 {
            return Ok(Wtf8Buf::from(format!("{class_name}()")));
        }
        if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
            inner.repr(Some(&class_name), vm)
        } else {
            Ok(Wtf8Buf::from(format!("{class_name}(...)")))
        }
    }
}

struct AnySet {
    object: PyObjectRef,
}

impl Borrow<PyObject> for AnySet {
    #[inline(always)]
    fn borrow(&self) -> &PyObject {
        &self.object
    }
}

impl AnySet {
    /// Check if object is a set or frozenset (including subclasses)
    /// Equivalent to CPython's PyAnySet_Check
    fn check(obj: &PyObject, vm: &VirtualMachine) -> bool {
        let ctx = &vm.ctx;
        obj.fast_isinstance(ctx.types.set_type) || obj.fast_isinstance(ctx.types.frozenset_type)
    }

    fn into_iterable(self, vm: &VirtualMachine) -> PyResult<ArgIterable> {
        self.object.try_into_value(vm)
    }

    fn into_iterable_iter(
        self,
        vm: &VirtualMachine,
    ) -> PyResult<impl core::iter::Iterator<Item = ArgIterable>> {
        Ok(core::iter::once(self.into_iterable(vm)?))
    }

    fn as_inner(&self) -> &PySetInner {
        match_class!(match self.object.as_object() {
            ref set @ PySet => &set.inner,
            ref frozen @ PyFrozenSet => &frozen.inner,
            _ => unreachable!("AnySet is always PySet or PyFrozenSet"), // should not be called.
        })
    }
}

impl TryFromObject for AnySet {
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        let class = obj.class();
        if class.fast_issubclass(vm.ctx.types.set_type)
            || class.fast_issubclass(vm.ctx.types.frozenset_type)
        {
            Ok(Self { object: obj })
        } else {
            Err(vm.new_type_error(format!("{class} is not a subtype of set or frozenset")))
        }
    }
}

#[pyclass(module = false, name = "set_iterator")]
pub(crate) struct PySetIterator {
    size: DictSize,
    /// Whether the set was found to have changed, which `setiter_iternext()`
    /// records by writing a size no set can have. Sticky: what it makes the
    /// iterator answer, it answers from then on.
    changed: PyAtomic<bool>,
    internal: PyMutex<PositionIterInternal<AnySet>>,
}

impl fmt::Debug for PySetIterator {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // TODO: implement more detailed, non-recursive Debug formatter
        f.write_str("set_iterator")
    }
}

impl PyPayload for PySetIterator {
    #[inline]
    fn class(ctx: &Context) -> &'static Py<PyType> {
        ctx.types.set_iterator_type
    }
}

impl PySetIterator {
    fn new(set: AnySet) -> Self {
        Self {
            size: set.as_inner().content.size(),
            changed: Radium::new(false),
            internal: PyMutex::new(PositionIterInternal::new(set, 0)),
        }
    }
}

#[pyclass(flags(DISALLOW_INSTANTIATION), with(IterNext, Iterable))]
impl Py<PySetIterator> {
    #[pymethod]
    fn __length_hint__(&self) -> usize {
        // `setiter_len()` answers for a set it can no longer walk with nothing,
        // comparing the size it captured against the set's own every time it is
        // asked.
        if self.changed.load(Ordering::Relaxed) {
            return 0;
        }
        self.internal.lock().length_hint(|set| {
            if set.as_inner().content.size() == self.size {
                self.size.entries_size
            } else {
                0
            }
        })
    }

    #[pymethod]
    fn __reduce__(
        zelf: PyRef<PySetIterator>,
        vm: &VirtualMachine,
    ) -> PyResult<(PyObjectRef, (PyObjectRef,))> {
        let internal = zelf.internal.lock();
        Ok((
            builtins_iter(vm)?,
            (vm.ctx
                .new_list(match &internal.status {
                    IterStatus::Exhausted => vec![],
                    IterStatus::Active(set) => set
                        .as_inner()
                        .content
                        .keys()
                        .into_iter()
                        .skip(internal.position)
                        .collect(),
                })
                .into(),),
        ))
    }
}

impl SelfIter for PySetIterator {}
impl IterNext for PySetIterator {
    fn next(zelf: &crate::Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
        locked_step(&zelf.internal, |internal| {
            let IterStatus::Active(set) = &internal.status else {
                return (Ok(PyIterReturn::StopIteration(None)), None);
            };
            let mutated = || vm.new_runtime_error("Set changed size during iteration");
            if zelf.changed.load(Ordering::Relaxed) {
                // The set is not looked at again once it has been found to
                // change: an iterator that has raised keeps raising.
                return (Err(mutated()), None);
            }
            let entry = set.as_inner().content.next_entry_checked(
                internal.position,
                &zelf.size,
                |key, ()| key.to_owned(),
            );
            match entry {
                Err(crate::dict_inner::DictChanged) => {
                    zelf.changed.store(true, Ordering::Relaxed);
                    (Err(mutated()), None)
                }
                Ok(Some((position, key))) => {
                    internal.position = position;
                    (Ok(PyIterReturn::Return(key)), None)
                }
                Ok(None) => (Ok(PyIterReturn::StopIteration(None)), internal.exhaust()),
            }
        })
    }
}

fn vectorcall_set(
    zelf_obj: &PyObject,
    args: Vec<PyObjectRef>,
    nargs: usize,
    kwnames: Option<&[PyObjectRef]>,
    vm: &VirtualMachine,
) -> PyResult {
    let zelf: &Py<PyType> = zelf_obj.downcast_ref().unwrap();
    let obj = PySet::default().into_ref_with_type(vm, zelf.to_owned())?;
    let func_args = FuncArgs::from_vectorcall_owned(args, nargs, kwnames);
    PySet::slot_init(obj.as_object(), func_args, vm)?;
    Ok(obj.into())
}

fn vectorcall_frozenset(
    zelf_obj: &PyObject,
    args: Vec<PyObjectRef>,
    nargs: usize,
    kwnames: Option<&[PyObjectRef]>,
    vm: &VirtualMachine,
) -> PyResult {
    let zelf: &Py<PyType> = zelf_obj.downcast_ref().unwrap();
    let func_args = FuncArgs::from_vectorcall_owned(args, nargs, kwnames);
    (zelf.slots.new.load().unwrap())(zelf.to_owned(), func_args, vm)
}

pub(crate) fn init(context: &'static Context) {
    PySet::extend_class(context, context.types.set_type);
    context
        .types
        .set_type
        .slots
        .vectorcall
        .store(Some(vectorcall_set));

    PyFrozenSet::extend_class(context, context.types.frozenset_type);
    context
        .types
        .frozenset_type
        .slots
        .vectorcall
        .store(Some(vectorcall_frozenset));

    PySetIterator::extend_class(context, context.types.set_iterator_type);
}