rustpython-vm 0.5.0

use super::{
    core::{Py, PyObject, PyObjectRef, PyRef},
    payload::PyPayload,
};
use crate::common::{
    atomic::{Ordering, PyAtomic, Radium},
    lock::PyRwLockReadGuard,
};
use crate::{
    VirtualMachine,
    builtins::{PyBaseExceptionRef, PyStrInterned, PyType},
    convert::{IntoPyException, ToPyObject, ToPyResult, TryFromObject},
    vm::Context,
};
use alloc::fmt;

use core::{
    borrow::Borrow,
    marker::PhantomData,
    ops::Deref,
    ptr::{NonNull, null_mut},
};

/* Python objects and references.

Okay, so each python object itself is an class itself (PyObject). Each
python object can have several references to it (PyObjectRef). These
references are Rc (reference counting) rust smart pointers. So when
all references are destroyed, the object itself also can be cleaned up.
Basically reference counting, but then done by rust.

*/

/*
 * Good reference: https://github.com/ProgVal/pythonvm-rust/blob/master/src/objects/mod.rs
 */

/// Use this type for functions which return a python object or an exception.
/// Both the python object and the python exception are `PyObjectRef` types
/// since exceptions are also python objects.
pub type PyResult<T = PyObjectRef> = Result<T, PyBaseExceptionRef>; // A valid value, or an exception

impl<T: fmt::Display> fmt::Display for PyRef<T>
where
    T: PyPayload + fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&**self, f)
    }
}

impl<T: fmt::Display> fmt::Display for Py<T>
where
    T: PyPayload + fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&**self, f)
    }
}

#[repr(transparent)]
pub struct PyExact<T> {
    inner: Py<T>,
}

impl<T: PyPayload> PyExact<T> {
    /// # Safety
    /// Given reference must be exact type of payload T
    #[inline(always)]
    pub const unsafe fn ref_unchecked(r: &Py<T>) -> &Self {
        unsafe { &*(r as *const _ as *const Self) }
    }
}

impl<T: PyPayload> Deref for PyExact<T> {
    type Target = Py<T>;

    #[inline(always)]
    fn deref(&self) -> &Py<T> {
        &self.inner
    }
}

impl<T: PyPayload> Borrow<PyObject> for PyExact<T> {
    #[inline(always)]
    fn borrow(&self) -> &PyObject {
        self.inner.borrow()
    }
}

impl<T: PyPayload> AsRef<PyObject> for PyExact<T> {
    #[inline(always)]
    fn as_ref(&self) -> &PyObject {
        self.inner.as_ref()
    }
}

impl<T: PyPayload> Borrow<Py<T>> for PyExact<T> {
    #[inline(always)]
    fn borrow(&self) -> &Py<T> {
        &self.inner
    }
}

impl<T: PyPayload> AsRef<Py<T>> for PyExact<T> {
    #[inline(always)]
    fn as_ref(&self) -> &Py<T> {
        &self.inner
    }
}

impl<T: PyPayload> alloc::borrow::ToOwned for PyExact<T> {
    type Owned = PyRefExact<T>;

    fn to_owned(&self) -> Self::Owned {
        let owned = self.inner.to_owned();
        unsafe { PyRefExact::new_unchecked(owned) }
    }
}

impl<T: PyPayload> PyRef<T> {
    pub fn into_exact_or(
        self,
        ctx: &Context,
        f: impl FnOnce(Self) -> PyRefExact<T>,
    ) -> PyRefExact<T> {
        if self.class().is(T::class(ctx)) {
            unsafe { PyRefExact::new_unchecked(self) }
        } else {
            f(self)
        }
    }
}

/// PyRef but guaranteed not to be a subtype instance
#[derive(Debug)]
#[repr(transparent)]
pub struct PyRefExact<T: PyPayload> {
    inner: PyRef<T>,
}

impl<T: PyPayload> PyRefExact<T> {
    /// # Safety
    /// obj must have exact type for the payload
    pub const unsafe fn new_unchecked(obj: PyRef<T>) -> Self {
        Self { inner: obj }
    }

    pub fn into_pyref(self) -> PyRef<T> {
        self.inner
    }
}

impl<T: PyPayload> Clone for PyRefExact<T> {
    fn clone(&self) -> Self {
        let inner = self.inner.clone();
        Self { inner }
    }
}

impl<T: PyPayload> TryFromObject for PyRefExact<T> {
    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
        let target_cls = T::class(&vm.ctx);
        let cls = obj.class();
        if cls.is(target_cls) {
            let obj = obj
                .downcast()
                .map_err(|obj| vm.new_downcast_runtime_error(target_cls, &obj))?;
            Ok(Self { inner: obj })
        } else if cls.fast_issubclass(target_cls) {
            Err(vm.new_type_error(format!(
                "Expected an exact instance of '{}', not a subclass '{}'",
                target_cls.name(),
                cls.name(),
            )))
        } else {
            Err(vm.new_type_error(format!(
                "Expected type '{}', not '{}'",
                target_cls.name(),
                cls.name(),
            )))
        }
    }
}

impl<T: PyPayload> Deref for PyRefExact<T> {
    type Target = PyExact<T>;

    #[inline(always)]
    fn deref(&self) -> &PyExact<T> {
        unsafe { PyExact::ref_unchecked(self.inner.deref()) }
    }
}

impl<T: PyPayload> Borrow<PyObject> for PyRefExact<T> {
    #[inline(always)]
    fn borrow(&self) -> &PyObject {
        self.inner.borrow()
    }
}

impl<T: PyPayload> AsRef<PyObject> for PyRefExact<T> {
    #[inline(always)]
    fn as_ref(&self) -> &PyObject {
        self.inner.as_ref()
    }
}

impl<T: PyPayload> Borrow<Py<T>> for PyRefExact<T> {
    #[inline(always)]
    fn borrow(&self) -> &Py<T> {
        self.inner.borrow()
    }
}

impl<T: PyPayload> AsRef<Py<T>> for PyRefExact<T> {
    #[inline(always)]
    fn as_ref(&self) -> &Py<T> {
        self.inner.as_ref()
    }
}

impl<T: PyPayload> Borrow<PyExact<T>> for PyRefExact<T> {
    #[inline(always)]
    fn borrow(&self) -> &PyExact<T> {
        self
    }
}

impl<T: PyPayload> AsRef<PyExact<T>> for PyRefExact<T> {
    #[inline(always)]
    fn as_ref(&self) -> &PyExact<T> {
        self
    }
}

impl<T: PyPayload> ToPyObject for PyRefExact<T> {
    #[inline(always)]
    fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
        self.inner.into()
    }
}

pub struct PyAtomicRef<T> {
    inner: PyAtomic<*mut u8>,
    _phantom: PhantomData<T>,
}

impl<T> Drop for PyAtomicRef<T> {
    fn drop(&mut self) {
        // SAFETY: We are dropping the atomic reference, so we can safely
        // release the pointer.
        unsafe {
            let ptr = Radium::swap(&self.inner, null_mut(), Ordering::Relaxed);
            if let Some(ptr) = NonNull::<PyObject>::new(ptr.cast()) {
                let _: PyObjectRef = PyObjectRef::from_raw(ptr);
            }
        }
    }
}

cfg_if::cfg_if! {
    if #[cfg(feature = "threading")] {
        unsafe impl<T: Send + PyPayload> Send for PyAtomicRef<T> {}
        unsafe impl<T: Sync + PyPayload> Sync for PyAtomicRef<T> {}
        unsafe impl<T: Send + PyPayload> Send for PyAtomicRef<Option<T>> {}
        unsafe impl<T: Sync + PyPayload> Sync for PyAtomicRef<Option<T>> {}
        unsafe impl Send for PyAtomicRef<PyObject> {}
        unsafe impl Sync for PyAtomicRef<PyObject> {}
        unsafe impl Send for PyAtomicRef<Option<PyObject>> {}
        unsafe impl Sync for PyAtomicRef<Option<PyObject>> {}
    }
}

impl<T: fmt::Debug> fmt::Debug for PyAtomicRef<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "PyAtomicRef(")?;
        unsafe {
            self.inner
                .load(Ordering::Relaxed)
                .cast::<T>()
                .as_ref()
                .fmt(f)
        }?;
        write!(f, ")")
    }
}

impl<T: PyPayload> From<PyRef<T>> for PyAtomicRef<T> {
    fn from(pyref: PyRef<T>) -> Self {
        let py = PyRef::leak(pyref);
        let ptr = py as *const _ as *mut u8;
        // Expose provenance so we can re-derive via with_exposed_provenance
        // without Stacked Borrows tag restrictions during bootstrap
        ptr.expose_provenance();
        Self {
            inner: Radium::new(ptr),
            _phantom: Default::default(),
        }
    }
}

impl<T: PyPayload> Deref for PyAtomicRef<T> {
    type Target = Py<T>;

    fn deref(&self) -> &Self::Target {
        unsafe {
            self.inner
                .load(Ordering::Relaxed)
                .cast::<Py<T>>()
                .as_ref()
                .unwrap_unchecked()
        }
    }
}

impl<T: PyPayload> PyAtomicRef<T> {
    /// Load the raw pointer without creating a reference.
    /// Avoids Stacked Borrows retag, safe for use during bootstrap
    /// when type objects have self-referential pointers being mutated.
    #[inline(always)]
    pub(super) fn load_raw(&self) -> *const Py<T> {
        self.inner.load(Ordering::Relaxed).cast::<Py<T>>()
    }

    /// # Safety
    /// The caller is responsible to keep the returned PyRef alive
    /// until no more reference can be used via PyAtomicRef::deref()
    #[must_use]
    pub unsafe fn swap(&self, pyref: PyRef<T>) -> PyRef<T> {
        let py = PyRef::leak(pyref) as *const Py<T> as *mut _;
        let old = Radium::swap(&self.inner, py, Ordering::AcqRel);
        unsafe { PyRef::from_raw(old.cast()) }
    }

    pub fn swap_to_temporary_refs(&self, pyref: PyRef<T>, vm: &VirtualMachine) {
        let old = unsafe { self.swap(pyref) };
        if let Some(frame) = vm.current_frame() {
            frame.temporary_refs.lock().push(old.into());
        }
    }
}

impl<T: PyPayload> From<Option<PyRef<T>>> for PyAtomicRef<Option<T>> {
    fn from(opt_ref: Option<PyRef<T>>) -> Self {
        let val = opt_ref
            .map(|x| PyRef::leak(x) as *const Py<T> as *mut _)
            .unwrap_or(null_mut());
        Self {
            inner: Radium::new(val),
            _phantom: Default::default(),
        }
    }
}

impl<T: PyPayload> PyAtomicRef<Option<T>> {
    pub fn deref(&self) -> Option<&Py<T>> {
        self.deref_ordering(Ordering::Relaxed)
    }

    pub fn deref_ordering(&self, ordering: Ordering) -> Option<&Py<T>> {
        unsafe { self.inner.load(ordering).cast::<Py<T>>().as_ref() }
    }

    pub fn to_owned(&self) -> Option<PyRef<T>> {
        self.to_owned_ordering(Ordering::Relaxed)
    }

    pub fn to_owned_ordering(&self, ordering: Ordering) -> Option<PyRef<T>> {
        self.deref_ordering(ordering).map(|x| x.to_owned())
    }

    /// # Safety
    /// The caller is responsible to keep the returned PyRef alive
    /// until no more reference can be used via PyAtomicRef::deref()
    #[must_use]
    pub unsafe fn swap(&self, opt_ref: Option<PyRef<T>>) -> Option<PyRef<T>> {
        let val = opt_ref
            .map(|x| PyRef::leak(x) as *const Py<T> as *mut _)
            .unwrap_or(null_mut());
        let old = Radium::swap(&self.inner, val, Ordering::AcqRel);
        unsafe { old.cast::<Py<T>>().as_ref().map(|x| PyRef::from_raw(x)) }
    }

    pub fn swap_to_temporary_refs(&self, opt_ref: Option<PyRef<T>>, vm: &VirtualMachine) {
        let Some(old) = (unsafe { self.swap(opt_ref) }) else {
            return;
        };
        if let Some(frame) = vm.current_frame() {
            frame.temporary_refs.lock().push(old.into());
        }
    }
}

impl From<PyObjectRef> for PyAtomicRef<PyObject> {
    fn from(obj: PyObjectRef) -> Self {
        let obj = obj.into_raw();
        Self {
            inner: Radium::new(obj.cast().as_ptr()),
            _phantom: Default::default(),
        }
    }
}

impl Deref for PyAtomicRef<PyObject> {
    type Target = PyObject;

    fn deref(&self) -> &Self::Target {
        unsafe {
            self.inner
                .load(Ordering::Relaxed)
                .cast::<PyObject>()
                .as_ref()
                .unwrap_unchecked()
        }
    }
}

impl PyAtomicRef<PyObject> {
    /// # Safety
    /// The caller is responsible to keep the returned PyRef alive
    /// until no more reference can be used via PyAtomicRef::deref()
    #[must_use]
    pub unsafe fn swap(&self, obj: PyObjectRef) -> PyObjectRef {
        let obj = obj.into_raw();
        let old = Radium::swap(&self.inner, obj.cast().as_ptr(), Ordering::AcqRel);
        unsafe { PyObjectRef::from_raw(NonNull::new_unchecked(old.cast())) }
    }

    pub fn swap_to_temporary_refs(&self, obj: PyObjectRef, vm: &VirtualMachine) {
        let old = unsafe { self.swap(obj) };
        if let Some(frame) = vm.current_frame() {
            frame.temporary_refs.lock().push(old);
        }
    }
}

impl From<Option<PyObjectRef>> for PyAtomicRef<Option<PyObject>> {
    fn from(obj: Option<PyObjectRef>) -> Self {
        let val = obj
            .map(|x| x.into_raw().as_ptr().cast())
            .unwrap_or(null_mut());
        Self {
            inner: Radium::new(val),
            _phantom: Default::default(),
        }
    }
}

impl PyAtomicRef<Option<PyObject>> {
    pub fn deref(&self) -> Option<&PyObject> {
        self.deref_ordering(Ordering::Relaxed)
    }

    pub fn deref_ordering(&self, ordering: Ordering) -> Option<&PyObject> {
        unsafe { self.inner.load(ordering).cast::<PyObject>().as_ref() }
    }

    pub fn to_owned(&self) -> Option<PyObjectRef> {
        self.to_owned_ordering(Ordering::Relaxed)
    }

    pub fn to_owned_ordering(&self, ordering: Ordering) -> Option<PyObjectRef> {
        self.deref_ordering(ordering).map(|x| x.to_owned())
    }

    /// # Safety
    /// The caller is responsible to keep the returned PyRef alive
    /// until no more reference can be used via PyAtomicRef::deref()
    #[must_use]
    pub unsafe fn swap(&self, obj: Option<PyObjectRef>) -> Option<PyObjectRef> {
        let val = obj
            .map(|x| x.into_raw().as_ptr().cast())
            .unwrap_or(null_mut());
        let old = Radium::swap(&self.inner, val, Ordering::AcqRel);
        unsafe { NonNull::new(old.cast::<PyObject>()).map(|x| PyObjectRef::from_raw(x)) }
    }

    pub fn swap_to_temporary_refs(&self, obj: Option<PyObjectRef>, vm: &VirtualMachine) {
        let Some(old) = (unsafe { self.swap(obj) }) else {
            return;
        };
        if let Some(frame) = vm.current_frame() {
            frame.temporary_refs.lock().push(old);
        }
    }
}

/// Atomic borrowed (non-ref-counted) optional reference to a Python object.
/// Unlike `PyAtomicRef`, this does NOT own the reference.
/// The pointed-to object must outlive this reference.
pub struct PyAtomicBorrow {
    inner: PyAtomic<*mut u8>,
}

// Safety: Access patterns ensure the pointed-to object outlives this reference.
// The owner (generator/coroutine) clears this in its Drop impl before deallocation.
unsafe impl Send for PyAtomicBorrow {}
unsafe impl Sync for PyAtomicBorrow {}

impl PyAtomicBorrow {
    pub fn new() -> Self {
        Self {
            inner: Radium::new(null_mut()),
        }
    }

    pub fn store(&self, obj: &PyObject) {
        let ptr = obj as *const PyObject as *mut u8;
        Radium::store(&self.inner, ptr, Ordering::Relaxed);
    }

    pub fn load(&self) -> Option<&PyObject> {
        let ptr = Radium::load(&self.inner, Ordering::Relaxed);
        if ptr.is_null() {
            None
        } else {
            Some(unsafe { &*(ptr as *const PyObject) })
        }
    }

    pub fn clear(&self) {
        Radium::store(&self.inner, null_mut(), Ordering::Relaxed);
    }

    pub fn to_owned(&self) -> Option<PyObjectRef> {
        self.load().map(|obj| obj.to_owned())
    }
}

impl Default for PyAtomicBorrow {
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Debug for PyAtomicBorrow {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "PyAtomicBorrow({:?})",
            Radium::load(&self.inner, Ordering::Relaxed)
        )
    }
}

pub trait AsObject
where
    Self: Borrow<PyObject>,
{
    #[inline(always)]
    fn as_object(&self) -> &PyObject {
        self.borrow()
    }

    #[inline(always)]
    fn get_id(&self) -> usize {
        self.as_object().unique_id()
    }

    #[inline(always)]
    fn is<T>(&self, other: &T) -> bool
    where
        T: AsObject,
    {
        self.get_id() == other.get_id()
    }

    #[inline(always)]
    fn class(&self) -> &Py<PyType> {
        self.as_object().class()
    }

    fn get_class_attr(&self, attr_name: &'static PyStrInterned) -> Option<PyObjectRef> {
        self.class().get_attr(attr_name)
    }

    /// Determines if `obj` actually an instance of `cls`, this doesn't call __instancecheck__, so only
    /// use this if `cls` is known to have not overridden the base __instancecheck__ magic method.
    #[inline]
    fn fast_isinstance(&self, cls: &Py<PyType>) -> bool {
        self.class().fast_issubclass(cls)
    }
}

impl<T> AsObject for T where T: Borrow<PyObject> {}

impl PyObject {
    #[inline(always)]
    fn unique_id(&self) -> usize {
        self as *const Self as usize
    }
}

// impl<T: ?Sized> Borrow<PyObject> for PyRc<T> {
//     #[inline(always)]
//     fn borrow(&self) -> &PyObject {
//         unsafe { &*(&**self as *const T as *const PyObject) }
//     }
// }

/// A borrow of a reference to a Python object. This avoids having clone the `PyRef<T>`/
/// `PyObjectRef`, which isn't that cheap as that increments the atomic reference counter.
// TODO: check if we still need this
#[allow(dead_code)]
pub struct PyLease<'a, T: PyPayload> {
    inner: PyRwLockReadGuard<'a, PyRef<T>>,
}

impl<T: PyPayload> PyLease<'_, T> {
    #[inline(always)]
    pub fn into_owned(self) -> PyRef<T> {
        self.inner.clone()
    }
}

impl<T: PyPayload> Borrow<PyObject> for PyLease<'_, T> {
    #[inline(always)]
    fn borrow(&self) -> &PyObject {
        self.inner.as_ref()
    }
}

impl<T: PyPayload> Deref for PyLease<'_, T> {
    type Target = PyRef<T>;
    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<T> fmt::Display for PyLease<'_, T>
where
    T: PyPayload + fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&**self, f)
    }
}

impl<T: PyPayload> ToPyObject for PyRef<T> {
    #[inline(always)]
    fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
        self.into()
    }
}

impl ToPyObject for PyObjectRef {
    #[inline(always)]
    fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
        self
    }
}

impl ToPyObject for &PyObject {
    #[inline(always)]
    fn to_pyobject(self, _vm: &VirtualMachine) -> PyObjectRef {
        self.to_owned()
    }
}

// Allows a built-in function to return any built-in object payload without
// explicitly implementing `ToPyObject`.
impl<T> ToPyObject for T
where
    T: PyPayload + core::fmt::Debug + Sized,
{
    #[inline(always)]
    fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef {
        PyPayload::into_pyobject(self, vm)
    }
}

impl<T> ToPyResult for T
where
    T: ToPyObject,
{
    #[inline(always)]
    fn to_pyresult(self, vm: &VirtualMachine) -> PyResult {
        Ok(self.to_pyobject(vm))
    }
}

impl<T, E> ToPyResult for Result<T, E>
where
    T: ToPyObject,
    E: IntoPyException,
{
    #[inline(always)]
    fn to_pyresult(self, vm: &VirtualMachine) -> PyResult {
        self.map(|res| T::to_pyobject(res, vm))
            .map_err(|e| E::into_pyexception(e, vm))
    }
}

impl IntoPyException for PyBaseExceptionRef {
    #[inline(always)]
    fn into_pyexception(self, _vm: &VirtualMachine) -> PyBaseExceptionRef {
        self
    }
}