use once_cell::sync::Lazy;
use std::collections::hash_map::{Entry, HashMap};
use std::sync::Mutex;
use libR_sys::{
R_NilValue, R_PreserveObject, R_ReleaseObject, R_xlen_t, Rf_allocVector, Rf_protect,
Rf_unprotect, LENGTH, SET_VECTOR_ELT, SEXP, VECSXP, VECTOR_ELT,
};
static OWNERSHIP: Lazy<Mutex<Ownership>> = Lazy::new(|| Mutex::new(Ownership::new()));
pub(crate) unsafe fn protect(sexp: SEXP) {
let mut own = OWNERSHIP.lock().expect("protect failed");
own.protect(sexp);
}
pub(crate) unsafe fn unprotect(sexp: SEXP) {
let mut own = OWNERSHIP.lock().expect("unprotect failed");
own.unprotect(sexp);
}
pub const INITIAL_PRESERVATION_SIZE: usize = 100000;
pub const EXTRA_PRESERVATION_SIZE: usize = 100000;
struct Object {
refcount: usize,
index: usize,
}
struct Ownership {
preservation: usize,
cur_index: usize,
max_index: usize,
objects: HashMap<usize, Object>,
}
impl Ownership {
fn new() -> Self {
unsafe {
let preservation = Rf_allocVector(VECSXP, INITIAL_PRESERVATION_SIZE as R_xlen_t);
R_PreserveObject(preservation);
Ownership {
preservation: preservation as usize,
cur_index: 0,
max_index: INITIAL_PRESERVATION_SIZE,
objects: HashMap::with_capacity(INITIAL_PRESERVATION_SIZE),
}
}
}
unsafe fn protect(&mut self, sexp: SEXP) {
Rf_protect(sexp);
if self.cur_index == self.max_index {
self.garbage_collect();
}
let sexp_usize = sexp as usize;
let Ownership {
ref mut preservation,
ref mut cur_index,
ref mut max_index,
ref mut objects,
} = *self;
let mut entry = objects.entry(sexp_usize);
let preservation_sexp = *preservation as SEXP;
match entry {
Entry::Occupied(ref mut occupied) => {
if occupied.get().refcount == 0 {
SET_VECTOR_ELT(preservation_sexp, occupied.get().index as R_xlen_t, sexp);
}
occupied.get_mut().refcount += 1;
}
Entry::Vacant(vacant) => {
let index = *cur_index;
SET_VECTOR_ELT(preservation_sexp, index as R_xlen_t, sexp);
*cur_index += 1;
assert!(index != *max_index);
let refcount = 1;
vacant.insert(Object { refcount, index });
}
}
Rf_unprotect(1);
}
pub unsafe fn unprotect(&mut self, sexp: SEXP) {
let sexp_usize = sexp as usize;
let Ownership {
preservation,
cur_index: _,
max_index: _,
ref mut objects,
} = *self;
let mut entry = objects.entry(sexp_usize);
match entry {
Entry::Occupied(ref mut occupied) => {
let object = occupied.get_mut();
if object.refcount == 0 {
panic!("Attempt to unprotect an already unprotected object.")
} else {
object.refcount -= 1;
if object.refcount == 0 {
let preservation_sexp = preservation as SEXP;
SET_VECTOR_ELT(preservation_sexp, object.index as R_xlen_t, R_NilValue);
}
}
}
Entry::Vacant(_) => {
panic!("Attempt to unprotect a never protected object.")
}
}
}
#[allow(dead_code)]
unsafe fn ref_count(&mut self, sexp: SEXP) -> usize {
let Ownership {
preservation: _,
cur_index: _,
max_index: _,
ref mut objects,
} = *self;
let sexp_usize = sexp as usize;
let mut entry = objects.entry(sexp_usize);
match entry {
Entry::Occupied(ref mut occupied) => occupied.get().refcount,
Entry::Vacant(_) => 0,
}
}
unsafe fn garbage_collect(&mut self) {
let new_size = self.cur_index * 2 + EXTRA_PRESERVATION_SIZE;
let new_sexp = Rf_allocVector(VECSXP, new_size as R_xlen_t);
R_PreserveObject(new_sexp);
let old_sexp = self.preservation as SEXP;
let mut new_objects = HashMap::with_capacity(new_size);
let mut j = 0;
for (addr, object) in self.objects.iter() {
if object.refcount != 0 {
SET_VECTOR_ELT(new_sexp, j as R_xlen_t, *addr as SEXP);
new_objects.insert(
*addr,
Object {
refcount: object.refcount,
index: j,
},
);
j += 1;
}
}
R_ReleaseObject(old_sexp);
self.preservation = new_sexp as usize;
self.cur_index = j;
self.max_index = new_size;
self.objects = new_objects;
}
#[allow(dead_code)]
unsafe fn check_objects(&mut self) {
let preservation_sexp = self.preservation as SEXP;
assert_eq!(self.max_index, LENGTH(preservation_sexp) as usize);
for (addr, object) in self.objects.iter() {
assert!(object.index < self.max_index);
let elt = VECTOR_ELT(preservation_sexp, object.index as R_xlen_t);
if object.refcount != 0 {
assert_eq!(elt, *addr as SEXP);
} else {
assert_eq!(elt, R_NilValue);
}
}
for i in 0..self.max_index {
let elt = VECTOR_ELT(preservation_sexp, i as R_xlen_t);
if elt == R_NilValue {
assert_eq!(self.ref_count(elt), 0);
} else {
assert!(self.ref_count(elt) != 0);
}
}
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::*;
use libR_sys::{Rf_ScalarInteger, Rf_protect, Rf_unprotect};
#[test]
fn basic_test() {
test! {
single_threaded(|| unsafe {
{
let mut own = OWNERSHIP.lock().expect("lock failed");
own.check_objects();
}
let sexp1 = Rf_protect(Rf_ScalarInteger(1));
let sexp2 = Rf_protect(Rf_ScalarInteger(2));
protect(sexp1);
{
let mut own = OWNERSHIP.lock().expect("lock failed");
own.check_objects();
assert_eq!(own.ref_count(sexp1), 1);
assert_eq!(own.ref_count(sexp2), 0);
}
protect(sexp1);
{
let mut own = OWNERSHIP.lock().expect("lock failed");
own.check_objects();
assert_eq!(own.ref_count(sexp1), 2);
assert_eq!(own.ref_count(sexp2), 0);
}
unprotect(sexp1);
{
let mut own = OWNERSHIP.lock().expect("lock failed");
own.check_objects();
assert_eq!(own.ref_count(sexp1), 1);
assert_eq!(own.ref_count(sexp2), 0);
}
unprotect(sexp1);
{
let mut own = OWNERSHIP.lock().expect("lock failed");
own.check_objects();
assert_eq!(own.ref_count(sexp1), 0);
assert_eq!(own.ref_count(sexp2), 0);
}
protect(sexp2);
{
let mut own = OWNERSHIP.lock().expect("lock failed");
own.check_objects();
assert_eq!(own.ref_count(sexp1), 0);
assert_eq!(own.ref_count(sexp2), 1);
}
protect(sexp1);
{
let mut own = OWNERSHIP.lock().expect("lock failed");
own.check_objects();
assert_eq!(own.ref_count(sexp1), 1);
assert_eq!(own.ref_count(sexp2), 1);
}
Rf_unprotect(2);
});
}
}
#[test]
fn collection_test() {
test! {
single_threaded(|| unsafe {
{
let mut own = OWNERSHIP.lock().expect("protect failed");
own.check_objects();
}
let test_size = INITIAL_PRESERVATION_SIZE + EXTRA_PRESERVATION_SIZE * 5;
let sexp_pres = Rf_allocVector(VECSXP, test_size as R_xlen_t);
Rf_protect(sexp_pres);
let sexps = (0..test_size).map(|i| {
let sexp = Rf_ScalarInteger(1);
SET_VECTOR_ELT(sexp_pres, i as R_xlen_t, sexp);
sexp
}).collect::<Vec<_>>();
for (i, sexp) in sexps.iter().enumerate() {
protect(*sexp);
if i % 2 == 0 {
unprotect(*sexp);
}
}
{
let mut own = OWNERSHIP.lock().expect("protect failed");
own.check_objects();
own.garbage_collect();
own.check_objects();
}
Rf_unprotect(1);
});
}
}
}