diff --git a/crates/vm/src/builtins/mod.rs b/crates/vm/src/builtins/mod.rs index c40f09bde..e9ca1f8b4 100644 --- a/crates/vm/src/builtins/mod.rs +++ b/crates/vm/src/builtins/mod.rs @@ -92,7 +92,7 @@ pub(crate) mod zip; pub use zip::PyZip; #[path = "union.rs"] pub(crate) mod union_; -pub use union_::PyUnion; +pub use union_::{PyUnion, make_union}; pub(crate) mod descriptor; pub use float::try_to_bigint as try_f64_to_bigint; diff --git a/crates/vm/src/stdlib/_abc.rs b/crates/vm/src/stdlib/_abc.rs new file mode 100644 index 000000000..06b403206 --- /dev/null +++ b/crates/vm/src/stdlib/_abc.rs @@ -0,0 +1,498 @@ +//! Implementation of the `_abc` module. +//! +//! This module provides the C implementation of Abstract Base Classes (ABCs) +//! as defined in PEP 3119. + +pub(crate) use _abc::make_module; + +#[pymodule] +mod _abc { + use crate::{ + AsObject, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, + builtins::{PyFrozenSet, PyList, PySet, PyStr, PyTupleRef, PyTypeRef, PyWeak}, + common::lock::PyRwLock, + convert::ToPyObject, + protocol::PyIterReturn, + types::Constructor, + }; + use std::sync::atomic::{AtomicU64, Ordering}; + + // Global invalidation counter + static ABC_INVALIDATION_COUNTER: AtomicU64 = AtomicU64::new(0); + + fn get_invalidation_counter() -> u64 { + ABC_INVALIDATION_COUNTER.load(Ordering::SeqCst) + } + + fn increment_invalidation_counter() { + ABC_INVALIDATION_COUNTER.fetch_add(1, Ordering::SeqCst); + } + + /// Internal state held by ABC machinery. + #[pyattr] + #[pyclass(name = "_abc_data", module = "_abc")] + #[derive(Debug, PyPayload)] + struct AbcData { + // WeakRef sets for registry and caches + registry: PyRwLock>>, + cache: PyRwLock>>, + negative_cache: PyRwLock>>, + negative_cache_version: AtomicU64, + } + + #[pyclass(with(Constructor))] + impl AbcData { + fn new() -> Self { + AbcData { + registry: PyRwLock::new(None), + cache: PyRwLock::new(None), + negative_cache: PyRwLock::new(None), + negative_cache_version: AtomicU64::new(get_invalidation_counter()), + } + } + + fn get_cache_version(&self) -> u64 { + self.negative_cache_version.load(Ordering::SeqCst) + } + + fn set_cache_version(&self, version: u64) { + self.negative_cache_version.store(version, Ordering::SeqCst); + } + } + + impl Constructor for AbcData { + type Args = (); + + fn py_new( + _cls: &crate::Py, + _args: Self::Args, + _vm: &VirtualMachine, + ) -> PyResult { + Ok(AbcData::new()) + } + } + + /// Get the _abc_impl attribute from an ABC class + fn get_impl(cls: &PyObject, vm: &VirtualMachine) -> PyResult> { + let impl_obj = cls.get_attr("_abc_impl", vm)?; + impl_obj + .downcast::() + .map_err(|_| vm.new_type_error("_abc_impl is set to a wrong type".to_owned())) + } + + /// Check if obj is in the weak set + fn in_weak_set( + set_lock: &PyRwLock>>, + obj: &PyObject, + vm: &VirtualMachine, + ) -> PyResult { + let set_opt = set_lock.read(); + let set = match &*set_opt { + Some(s) if !s.elements().is_empty() => s.clone(), + _ => return Ok(false), + }; + drop(set_opt); + + // Create a weak reference to the object + let weak_ref = match obj.downgrade(None, vm) { + Ok(w) => w, + Err(e) => { + // If we can't create a weakref (e.g., TypeError), the object can't be in the set + if e.class().is(vm.ctx.exceptions.type_error) { + return Ok(false); + } + return Err(e); + } + }; + + // Use vm.call_method to call __contains__ + let weak_ref_obj: PyObjectRef = weak_ref.into(); + vm.call_method(set.as_ref(), "__contains__", (weak_ref_obj,))? + .try_to_bool(vm) + } + + /// Add obj to the weak set + fn add_to_weak_set( + set_lock: &PyRwLock>>, + obj: &PyObject, + vm: &VirtualMachine, + ) -> PyResult<()> { + let mut set_opt = set_lock.write(); + let set = match &*set_opt { + Some(s) => s.clone(), + None => { + let new_set = PySet::default().into_ref(&vm.ctx); + *set_opt = Some(new_set.clone()); + new_set + } + }; + drop(set_opt); + + // Create a weak reference to the object + let weak_ref = obj.downgrade(None, vm)?; + set.add(weak_ref.into(), vm)?; + Ok(()) + } + + /// Returns the current ABC cache token. + #[pyfunction] + fn get_cache_token() -> u64 { + get_invalidation_counter() + } + + /// Compute set of abstract method names. + fn compute_abstract_methods(cls: &PyObject, vm: &VirtualMachine) -> PyResult<()> { + let mut abstracts = Vec::new(); + + // Stage 1: direct abstract methods + let ns = cls.get_attr("__dict__", vm)?; + let items = vm.call_method(&ns, "items", ())?; + let iter = items.get_iter(vm)?; + + loop { + match iter.next(vm)? { + PyIterReturn::Return(item) => { + let tuple: PyTupleRef = item + .downcast() + .map_err(|_| vm.new_type_error("items() returned non-tuple".to_owned()))?; + let elements = tuple.as_slice(); + if elements.len() != 2 { + return Err(vm.new_type_error( + "items() returned item which size is not 2".to_owned(), + )); + } + let key = &elements[0]; + let value = &elements[1]; + + // Check if value has __isabstractmethod__ = True + if let Ok(is_abstract) = value.get_attr("__isabstractmethod__", vm) { + if is_abstract.try_to_bool(vm)? { + abstracts.push(key.clone()); + } + } + } + PyIterReturn::StopIteration(_) => break, + } + } + + // Stage 2: inherited abstract methods + let bases: PyTupleRef = cls + .get_attr("__bases__", vm)? + .downcast() + .map_err(|_| vm.new_type_error("__bases__ is not a tuple".to_owned()))?; + + for base in bases.iter() { + if let Ok(base_abstracts) = base.get_attr("__abstractmethods__", vm) { + let iter = base_abstracts.get_iter(vm)?; + loop { + match iter.next(vm)? { + PyIterReturn::Return(key) => { + // Try to get the attribute from cls - key should be a string + if let Some(key_str) = key.downcast_ref::() { + if let Some(value) = + vm.get_attribute_opt(cls.to_owned(), key_str)? + { + if let Ok(is_abstract) = + value.get_attr("__isabstractmethod__", vm) + { + if is_abstract.try_to_bool(vm)? { + abstracts.push(key); + } + } + } + } + } + PyIterReturn::StopIteration(_) => break, + } + } + } + } + + // Set __abstractmethods__ + let abstracts_set = PyFrozenSet::from_iter(vm, abstracts.into_iter())?; + cls.set_attr("__abstractmethods__", abstracts_set.into_pyobject(vm), vm)?; + + Ok(()) + } + + /// Internal ABC helper for class set-up. Should be never used outside abc module. + #[pyfunction] + fn _abc_init(cls: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + compute_abstract_methods(&cls, vm)?; + + // Set up inheritance registry + let data = AbcData::new(); + cls.set_attr("_abc_impl", data.to_pyobject(vm), vm)?; + + Ok(()) + } + + /// Internal ABC helper for subclass registration. Should be never used outside abc module. + #[pyfunction] + fn _abc_register( + cls: PyObjectRef, + subclass: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + // Type check + if !subclass.class().fast_issubclass(vm.ctx.types.type_type) { + return Err(vm.new_type_error("Can only register classes".to_owned())); + } + + // Check if already a subclass + if subclass.is_subclass(&cls, vm)? { + return Ok(subclass); + } + + // Check for cycles + if cls.is_subclass(&subclass, vm)? { + return Err(vm.new_runtime_error("Refusing to create an inheritance cycle".to_owned())); + } + + // Add to registry + let impl_data = get_impl(&cls, vm)?; + add_to_weak_set(&impl_data.registry, &subclass, vm)?; + + // Invalidate negative cache + increment_invalidation_counter(); + + Ok(subclass) + } + + /// Internal ABC helper for instance checks. Should be never used outside abc module. + #[pyfunction] + fn _abc_instancecheck( + cls: PyObjectRef, + instance: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + let impl_data = get_impl(&cls, vm)?; + + // Get instance.__class__ + let subclass = instance.get_attr("__class__", vm)?; + + // Check cache + if in_weak_set(&impl_data.cache, &subclass, vm)? { + return Ok(vm.ctx.true_value.clone().into()); + } + + let subtype: PyObjectRef = instance.class().to_owned().into(); + if subtype.is(&subclass) { + let invalidation_counter = get_invalidation_counter(); + if impl_data.get_cache_version() == invalidation_counter { + if in_weak_set(&impl_data.negative_cache, &subclass, vm)? { + return Ok(vm.ctx.false_value.clone().into()); + } + } + // Fall back to __subclasscheck__ + return vm.call_method(&cls, "__subclasscheck__", (subclass,)); + } + + // Call __subclasscheck__ on subclass + let result = vm.call_method(&cls, "__subclasscheck__", (subclass.clone(),))?; + + match result.clone().try_to_bool(vm) { + Ok(true) => Ok(result), + Ok(false) => { + // Also try with subtype + vm.call_method(&cls, "__subclasscheck__", (subtype,)) + } + Err(e) => Err(e), + } + } + + /// Check if subclass is in registry (recursive) + fn subclasscheck_check_registry( + impl_data: &AbcData, + subclass: &PyObject, + vm: &VirtualMachine, + ) -> PyResult> { + // Fast path: check if subclass is in weakref directly + if in_weak_set(&impl_data.registry, subclass, vm)? { + return Ok(Some(true)); + } + + let registry_opt = impl_data.registry.read(); + let registry = match &*registry_opt { + Some(s) => s.clone(), + None => return Ok(None), + }; + drop(registry_opt); + + // Make a local copy to protect against concurrent modifications + let registry_copy = PyFrozenSet::from_iter(vm, registry.elements().into_iter())?; + + for weak_ref_obj in registry_copy.elements() { + if let Ok(weak_ref) = weak_ref_obj.downcast::() { + if let Some(rkey) = weak_ref.upgrade() { + if subclass.to_owned().is_subclass(&rkey, vm)? { + add_to_weak_set(&impl_data.cache, subclass, vm)?; + return Ok(Some(true)); + } + } + } + } + + Ok(None) + } + + /// Internal ABC helper for subclass checks. Should be never used outside abc module. + #[pyfunction] + fn _abc_subclasscheck( + cls: PyObjectRef, + subclass: PyObjectRef, + vm: &VirtualMachine, + ) -> PyResult { + // Type check + if !subclass.class().fast_issubclass(vm.ctx.types.type_type) { + return Err(vm.new_type_error("issubclass() arg 1 must be a class".to_owned())); + } + + let impl_data = get_impl(&cls, vm)?; + + // 1. Check cache + if in_weak_set(&impl_data.cache, &subclass, vm)? { + return Ok(true); + } + + // 2. Check negative cache; may have to invalidate + let invalidation_counter = get_invalidation_counter(); + if impl_data.get_cache_version() < invalidation_counter { + // Invalidate the negative cache + // Clone set ref and drop lock before calling into VM to avoid reentrancy + let set = impl_data.negative_cache.read().clone(); + if let Some(ref set) = set { + vm.call_method(set.as_ref(), "clear", ())?; + } + impl_data.set_cache_version(invalidation_counter); + } else if in_weak_set(&impl_data.negative_cache, &subclass, vm)? { + return Ok(false); + } + + // 3. Check the subclass hook + let ok = vm.call_method(&cls, "__subclasshook__", (subclass.clone(),))?; + if ok.is(&vm.ctx.true_value) { + add_to_weak_set(&impl_data.cache, &subclass, vm)?; + return Ok(true); + } + if ok.is(&vm.ctx.false_value) { + add_to_weak_set(&impl_data.negative_cache, &subclass, vm)?; + return Ok(false); + } + if !ok.is(&vm.ctx.not_implemented) { + return Err(vm.new_exception_msg( + vm.ctx.exceptions.assertion_error.to_owned(), + "__subclasshook__ must return either False, True, or NotImplemented".to_owned(), + )); + } + + // 4. Check if it's a direct subclass + let subclass_type: PyTypeRef = subclass + .clone() + .downcast() + .map_err(|_| vm.new_type_error("expected a type object".to_owned()))?; + let cls_type: PyTypeRef = cls + .clone() + .downcast() + .map_err(|_| vm.new_type_error("expected a type object".to_owned()))?; + if subclass_type.fast_issubclass(&cls_type) { + add_to_weak_set(&impl_data.cache, &subclass, vm)?; + return Ok(true); + } + + // 5. Check if it's a subclass of a registered class (recursive) + if let Some(result) = subclasscheck_check_registry(&impl_data, &subclass, vm)? { + return Ok(result); + } + + // 6. Check if it's a subclass of a subclass (recursive) + let subclasses: PyRef = vm + .call_method(&cls, "__subclasses__", ())? + .downcast() + .map_err(|_| vm.new_type_error("__subclasses__() must return a list".to_owned()))?; + + for scls in subclasses.borrow_vec().iter() { + if subclass.is_subclass(scls, vm)? { + add_to_weak_set(&impl_data.cache, &subclass, vm)?; + return Ok(true); + } + } + + // No dice; update negative cache + add_to_weak_set(&impl_data.negative_cache, &subclass, vm)?; + Ok(false) + } + + /// Internal ABC helper for cache and registry debugging. + #[pyfunction] + fn _get_dump(cls: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let impl_data = get_impl(&cls, vm)?; + + let registry = { + let r = impl_data.registry.read(); + match &*r { + Some(s) => { + // Use copy method to get a shallow copy + vm.call_method(s.as_ref(), "copy", ())? + } + None => PySet::default().to_pyobject(vm), + } + }; + + let cache = { + let c = impl_data.cache.read(); + match &*c { + Some(s) => vm.call_method(s.as_ref(), "copy", ())?, + None => PySet::default().to_pyobject(vm), + } + }; + + let negative_cache = { + let nc = impl_data.negative_cache.read(); + match &*nc { + Some(s) => vm.call_method(s.as_ref(), "copy", ())?, + None => PySet::default().to_pyobject(vm), + } + }; + + let version = impl_data.get_cache_version(); + + Ok(vm.ctx.new_tuple(vec![ + registry, + cache, + negative_cache, + vm.ctx.new_int(version).into(), + ])) + } + + /// Internal ABC helper to reset registry of a given class. + #[pyfunction] + fn _reset_registry(cls: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let impl_data = get_impl(&cls, vm)?; + // Clone set ref and drop lock before calling into VM to avoid reentrancy + let set = impl_data.registry.read().clone(); + if let Some(ref set) = set { + vm.call_method(set.as_ref(), "clear", ())?; + } + Ok(()) + } + + /// Internal ABC helper to reset both caches of a given class. + #[pyfunction] + fn _reset_caches(cls: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let impl_data = get_impl(&cls, vm)?; + + // Clone set refs and drop locks before calling into VM to avoid reentrancy + let cache = impl_data.cache.read().clone(); + if let Some(ref set) = cache { + vm.call_method(set.as_ref(), "clear", ())?; + } + + let negative_cache = impl_data.negative_cache.read().clone(); + if let Some(ref set) = negative_cache { + vm.call_method(set.as_ref(), "clear", ())?; + } + + Ok(()) + } +} diff --git a/crates/vm/src/stdlib/mod.rs b/crates/vm/src/stdlib/mod.rs index e46f333a2..85c28983d 100644 --- a/crates/vm/src/stdlib/mod.rs +++ b/crates/vm/src/stdlib/mod.rs @@ -1,3 +1,4 @@ +mod _abc; #[cfg(feature = "ast")] pub(crate) mod ast; pub mod atexit; @@ -87,6 +88,7 @@ pub fn get_module_inits() -> StdlibMap { modules! { #[cfg(all())] { + "_abc" => _abc::make_module, "atexit" => atexit::make_module, "_codecs" => codecs::make_module, "_collections" => collections::make_module, diff --git a/crates/vm/src/stdlib/typevar.rs b/crates/vm/src/stdlib/typevar.rs index 65249bfd0..5d5b1fc7b 100644 --- a/crates/vm/src/stdlib/typevar.rs +++ b/crates/vm/src/stdlib/typevar.rs @@ -1,7 +1,7 @@ // spell-checker:ignore typevarobject funcobj use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, - builtins::{PyTupleRef, PyType, PyTypeRef, pystr::AsPyStr}, + builtins::{PyTuple, PyTupleRef, PyType, PyTypeRef, make_union, pystr::AsPyStr}, common::lock::PyMutex, function::{FuncArgs, IntoFuncArgs, PyComparisonValue}, protocol::PyNumberMethods, @@ -250,7 +250,8 @@ impl AsNumber for TypeVar { fn as_number() -> &'static PyNumberMethods { static AS_NUMBER: PyNumberMethods = PyNumberMethods { or: Some(|a, b, vm| { - _call_typing_func_object(vm, "_make_union", (a.to_owned(), b.to_owned())) + let args = PyTuple::new_ref(vec![a.to_owned(), b.to_owned()], &vm.ctx); + Ok(make_union(&args, vm)) }), ..PyNumberMethods::NOT_IMPLEMENTED }; @@ -525,7 +526,8 @@ impl AsNumber for ParamSpec { fn as_number() -> &'static PyNumberMethods { static AS_NUMBER: PyNumberMethods = PyNumberMethods { or: Some(|a, b, vm| { - _call_typing_func_object(vm, "_make_union", (a.to_owned(), b.to_owned())) + let args = PyTuple::new_ref(vec![a.to_owned(), b.to_owned()], &vm.ctx); + Ok(make_union(&args, vm)) }), ..PyNumberMethods::NOT_IMPLEMENTED };