dvcompute_utils 2.0.0

// Copyright (c) 2020-2022  David Sorokin <davsor@mail.ru>, based in Yoshkar-Ola, Russia
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

use std::cell::RefCell;
use std::ops::Deref;
use std::ops::DerefMut;
use std::ptr;
use std::hash::Hash;
use std::hash::Hasher;
use std::borrow::Borrow;
use std::borrow::BorrowMut;

/// A shared reference like the standard `Rc` but such one which can be
/// garbage collected after the simulation run is finished.
pub struct Grc<T> where T: 'static {

    /// The cell pointer.
    cell: *mut GrcCell<T>
}

impl<T> Grc<T> {

    /// Create a new shared reference.
    pub fn new(item: T) -> Self {
        let cell = Box::new(GrcCell {
            header: GrcHeader::new(),
            strong_count: 1,
            weak_count: 0,
            item: Some(item)
        });
        let cell = Box::into_raw(cell);

        unsafe {
            GrcStorage::add(cell);
        }

        Self {
            cell: cell
        }
    }

    /// Get the corresponding weak reference.
    #[inline]
    pub fn downgrade(r: &Grc<T>) -> Weak<T> {
        unsafe {
            (*r.cell).weak_count += 1;
        }

        Weak {
            cell: r.cell
        }
    }

    /// Compare reference pointers for equality.
    #[inline]
    pub fn ptr_eq(r1: &Grc<T>, r2: &Grc<T>) -> bool {
        ptr::eq(r1.cell, r2.cell)
    }
}

impl<T> Drop for Grc<T> {

    fn drop(&mut self) {
        unsafe {
            (*self.cell).strong_count -= 1;
            if (*self.cell).strong_count == 0 {
                if (*self.cell).weak_count == 0 {
                    GrcStorage::remove(self.cell);

                } else {
                    if (*self.cell).item.is_some() {
                        (*self.cell).strong_count +=1;
                        let _ = (*self.cell).item.take();
                        (*self.cell).strong_count -=1;

                        if (*self.cell).weak_count == 0 {
                            GrcStorage::remove(self.cell);
                        }
                    }
                }
            }
        }
    }
}

impl<T> Clone for Grc<T> {

    #[inline]
    fn clone(&self) -> Self {
        unsafe {
            (*self.cell).strong_count += 1;
        }

        Self {
            cell: self.cell
        }
    }
}

impl<T> Deref for Grc<T> {

    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        unsafe {
            match &((*self.cell).item) {
                Some(ref item) => item,
                None => panic!("Grc reference cannot be empty")
            }
        }
    }
}

impl<T> DerefMut for Grc<T> {

    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe {
            match &mut ((*self.cell).item) {
                Some(ref mut item) => item,
                None => panic!("Grc reference cannot be empty")
            }
        }
    }
}

impl<T> PartialEq for Grc<T> where T: PartialEq {

    #[inline]
    fn eq(&self, other: &Self) -> bool {
        (self.deref()).eq(other.deref())
    }
}

impl<T> Eq for Grc<T> where T: Eq {}

impl<T> Hash for Grc<T> where T: Hash {

    #[inline]
    fn hash<H>(&self, state: &mut H) 
        where H: Hasher
    {
        (self.deref()).hash(state)
    }
}

impl<T> Borrow<T> for Grc<T> {

    #[inline]
    fn borrow(&self) -> &T {
        self.deref()
    }
}

impl<T> BorrowMut<T> for Grc<T> {

    #[inline]
    fn borrow_mut(&mut self) -> &mut T {
        self.deref_mut()
    }
}

/// A weak reference like the standard `Weak` but such one which can be
/// garbage collected after the simulation run is finished.
pub struct Weak<T> where T: 'static {

    /// The cell pointer.
    cell: *mut GrcCell<T>
}

impl<T> Weak<T> {

    /// Try to get the corresponding strong [`Grc<T>`] reference
    /// if the referenced data are still available, i.e.
    /// if, at least, any strong reference still exists.
    pub fn upgrade(&self) -> Option<Grc<T>> {
        unsafe {
            match &((*self.cell).item) {
                None => None,
                Some(_) => {
                    (*self.cell).strong_count += 1;
                    Some(Grc {
                        cell: self.cell
                    })
                }
            }
        }
    }
}

impl<T> Drop for Weak<T> {

    fn drop(&mut self) {
        unsafe {
            (*self.cell).weak_count -= 1;
            if (*self.cell).weak_count == 0 {
                if (*self.cell).strong_count == 0 {
                    GrcStorage::remove(self.cell);
                }
            }
        }
    }
}

impl<T> Clone for Weak<T> {

    #[inline]
    fn clone(&self) -> Self {
        unsafe {
            (*self.cell).weak_count += 1;
        }

        Self {
            cell: self.cell
        }
    }
}

/// The header object.
struct GrcHeader {

    /// The object state.
    state: GrcState,

    /// The previous object.
    prev: Option<*mut dyn GrcObject>,

    /// The next object.
    next: Option<*mut dyn GrcObject>
}

impl GrcHeader {

    /// Create a new header.
    #[inline]
    fn new() -> Self {
        Self {
            state: GrcState::Active,
            prev: None,
            next: None
        }
    }
}

/// A trait to reprecent the reference object.
trait GrcObject {

    /// Get the header.
    fn header_mut(&mut self) -> &mut GrcHeader;

    /// Free the object.
    unsafe fn free(&mut self);

    /// Destroy the object.
    unsafe fn destroy(&mut self);
}

/// The cell that contains data.
struct GrcCell<T> {

    /// The header object.
    header: GrcHeader,

    /// The strong use count.
    strong_count: usize,

    /// The weak use count.
    weak_count: usize,

    /// The optional data.
    item: Option<T>
}

impl<T> GrcObject for GrcCell<T> {

    fn header_mut(&mut self) -> &mut GrcHeader {
        &mut self.header
    }

    unsafe fn free(&mut self) {
        let _ = self.item.take();
    }

    unsafe fn destroy(&mut self) {
        let _ = Box::from_raw(self);
    }
}

/// The [Grc] reference state.
#[derive(PartialEq, Eq)]
enum GrcState {

    /// Usual state.
    Active,

    /// When the objects are removed. 
    FreeObjects
}

/// A storage of [Grc] references.
pub struct GrcStorage {

    /// The storage state.
    state: GrcState,

    /// The head of active objects.
    head_active: Option<*mut dyn GrcObject>,

    /// The tail of active objects.
    tail_active: Option<*mut dyn GrcObject>
}

thread_local! {

    /// The thread local storage.
    static STORAGE: RefCell<GrcStorage> = RefCell::new(GrcStorage::new());
}

impl GrcStorage {

    /// Create a new storage.
    fn new() -> Self {
        Self {
            state: GrcState::Active,
            head_active: None,
            tail_active: None
        }
    }

    /// Add a new object pointer.
    unsafe fn add(p: *mut dyn GrcObject) {
        STORAGE.with(|cell| {
            let mut storage = cell.borrow_mut();
            storage.add_to_active(p)
        })
    }

    /// Remove the object pointer.
    unsafe fn remove(p: *mut dyn GrcObject) {
        STORAGE.with(|cell| {
            {
                let h = (*p).header_mut();
                match h.state {
                    GrcState::Active => {
                        // it was released earlier
                    },
                    GrcState::FreeObjects => {
                        let mut storage = cell.borrow_mut();
                        storage.remove_from_active(p);
                    }
                }
            }
            (*p).destroy();
        })
    }

    /// Add a new object pointer.
    unsafe fn add_to_active(&mut self, p: *mut dyn GrcObject) {
        let h = (*p).header_mut();
        debug_assert!(h.state == GrcState::Active);
        debug_assert!(h.prev.is_none());
        debug_assert!(h.next.is_none());

        h.state = GrcState::FreeObjects;
        h.prev = self.tail_active;
        h.next = None;

        if let Some(x) = self.tail_active {
            let tail_h = (*x).header_mut();
            tail_h.next = Some(p);

        } else if self.head_active.is_none() {
            self.head_active = Some(p);

        } else {
            panic!("Inconsistent state");
        }

        self.tail_active = Some(p);
    }

    /// Remove from the list of active objects.
    unsafe fn remove_from_active(&mut self, p: *mut dyn GrcObject) {
        let h = (*p).header_mut();
        debug_assert!(h.state == GrcState::FreeObjects);

        if let Some(x) = h.prev {
            let prev_h = (*x).header_mut();
            prev_h.next = h.next;

        } else if let Some(x) = self.head_active {
            if ptr::eq(x as *const u8, p as *const u8) {
                self.head_active = h.next;

            } else {
                panic!("Inconsistent state");
            }

        } else {
            panic!("Inconsistent state");
        }
        
        if let Some(x) = h.next {
            let next_h = (*x).header_mut();
            next_h.prev = h.prev;

        } else if let Some(x) = self.tail_active {
            if ptr::eq(x as *const u8, p as *const u8) {
                self.tail_active = h.prev;

            } else {
                panic!("Inconsistent state");
            }

        } else {
            panic!("Inconsistent state");
        }

        h.state = GrcState::Active;
        h.prev = None;
        h.next = None;
    }

    /// Free all [Grc] objects related to the current execution thread.
    pub fn free_thread_local() {
        loop {
            let mut idle = true;
            let mut x_to_free: Option<*mut dyn GrcObject> = None;

            STORAGE.with(|cell| {
                let mut storage = cell.borrow_mut();
                match storage.state {
                    GrcState::Active => {
                        storage.state = GrcState::FreeObjects;
                        idle = false; 
                    },
                    GrcState::FreeObjects => {
                        if let Some(x) = storage.head_active {
                            unsafe {
                                storage.remove_from_active(x);
                                x_to_free = Some(x);
                                idle = false;
                            }

                        } else {
                            storage.state = GrcState::Active;
                            idle = true;
                        }
                    }
                }
            });

            if idle {
                break;
            }

            if let Some(x) = x_to_free {
                unsafe {
                    (*x).free();
                }
            }
        }
    }
}

/// Execute some action and then release the local `Grc` storage.
pub fn with_grc_storage<F, R>(f: F) -> R
    where F: FnOnce() -> R
{
    let result = f();
    GrcStorage::free_thread_local();
    result
}