use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::{Rc, Weak};
use crate::error::{StatorError, StatorResult};
use crate::objects::heap_object::HeapObject;
use crate::objects::property_map::PropertyMap;
use crate::objects::value::JsValue;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum TokenKey {
Address(usize),
Symbol(u64),
}
#[derive(Debug, Clone)]
struct Registration {
target: usize,
held_value: JsValue,
unregister_token: Option<TokenKey>,
}
#[derive(Debug, Clone)]
struct PlainRegistration {
target: Weak<RefCell<PropertyMap>>,
held_value: JsValue,
unregister_token: Option<TokenKey>,
}
#[derive(Debug, Default)]
pub struct JsFinalizationRegistry {
registrations: HashMap<u64, Registration>,
plain_registrations: HashMap<u64, PlainRegistration>,
next_id: u64,
cleanup_queue: Vec<JsValue>,
}
pub fn finalization_registry_new() -> JsFinalizationRegistry {
JsFinalizationRegistry::default()
}
pub fn finalization_registry_register(
registry: &mut JsFinalizationRegistry,
target: *mut HeapObject,
held_value: JsValue,
unregister_token: Option<*mut HeapObject>,
) -> StatorResult<()> {
if target.is_null() {
return Err(StatorError::TypeError(
"FinalizationRegistry target must be an object".into(),
));
}
if let Some(tok) = unregister_token
&& tok.is_null()
{
return Err(StatorError::TypeError(
"FinalizationRegistry unregister token must be an object".into(),
));
}
let id = registry.next_id;
registry.next_id += 1;
registry.registrations.insert(
id,
Registration {
target: target as usize,
held_value,
unregister_token: unregister_token.map(|t| TokenKey::Address(t as usize)),
},
);
Ok(())
}
pub fn finalization_registry_unregister(
registry: &mut JsFinalizationRegistry,
token: *mut HeapObject,
) -> StatorResult<bool> {
if token.is_null() {
return Err(StatorError::TypeError(
"FinalizationRegistry unregister token must be an object".into(),
));
}
let key = TokenKey::Address(token as usize);
let before = registry.registrations.len();
registry
.registrations
.retain(|_, reg| reg.unregister_token != Some(key.clone()));
Ok(registry.registrations.len() < before)
}
pub fn finalization_registry_register_with_token_key(
registry: &mut JsFinalizationRegistry,
target: *mut HeapObject,
held_value: JsValue,
unregister_token: Option<TokenKey>,
) -> StatorResult<()> {
if target.is_null() {
return Err(StatorError::TypeError(
"FinalizationRegistry target must be an object".into(),
));
}
let id = registry.next_id;
registry.next_id += 1;
registry.registrations.insert(
id,
Registration {
target: target as usize,
held_value,
unregister_token,
},
);
Ok(())
}
pub fn finalization_registry_register_plain_with_token_key(
registry: &mut JsFinalizationRegistry,
target: &Rc<RefCell<PropertyMap>>,
held_value: JsValue,
unregister_token: Option<TokenKey>,
) {
let id = registry.next_id;
registry.next_id += 1;
registry.plain_registrations.insert(
id,
PlainRegistration {
target: Rc::downgrade(target),
held_value,
unregister_token,
},
);
}
pub fn finalization_registry_unregister_by_key(
registry: &mut JsFinalizationRegistry,
token: TokenKey,
) -> bool {
let before = registry.registrations.len() + registry.plain_registrations.len();
registry
.registrations
.retain(|_, reg| reg.unregister_token != Some(token.clone()));
registry
.plain_registrations
.retain(|_, reg| reg.unregister_token != Some(token.clone()));
(registry.registrations.len() + registry.plain_registrations.len()) < before
}
pub fn finalization_registry_notify(
registry: &mut JsFinalizationRegistry,
target: *mut HeapObject,
) {
if target.is_null() {
return;
}
let addr = target as usize;
let mut remaining = HashMap::new();
for (id, reg) in registry.registrations.drain() {
if reg.target == addr {
registry.cleanup_queue.push(reg.held_value);
} else {
remaining.insert(id, reg);
}
}
registry.registrations = remaining;
}
pub fn finalization_registry_drain(registry: &mut JsFinalizationRegistry) -> Vec<JsValue> {
std::mem::take(&mut registry.cleanup_queue)
}
pub fn finalization_registry_register_plain(
registry: &mut JsFinalizationRegistry,
target: &Rc<RefCell<PropertyMap>>,
held_value: JsValue,
unregister_token: Option<&Rc<RefCell<PropertyMap>>>,
) {
let id = registry.next_id;
registry.next_id += 1;
registry.plain_registrations.insert(
id,
PlainRegistration {
target: Rc::downgrade(target),
held_value,
unregister_token: unregister_token.map(|t| TokenKey::Address(Rc::as_ptr(t) as usize)),
},
);
}
pub fn finalization_registry_unregister_plain(
registry: &mut JsFinalizationRegistry,
token: &Rc<RefCell<PropertyMap>>,
) -> bool {
let key = TokenKey::Address(Rc::as_ptr(token) as usize);
let before = registry.plain_registrations.len();
registry
.plain_registrations
.retain(|_, reg| reg.unregister_token != Some(key.clone()));
registry.plain_registrations.len() < before
}
pub fn finalization_registry_sweep_plain(registry: &mut JsFinalizationRegistry) {
let mut remaining = HashMap::new();
for (id, reg) in registry.plain_registrations.drain() {
if reg.target.strong_count() == 0 {
registry.cleanup_queue.push(reg.held_value);
} else {
remaining.insert(id, reg);
}
}
registry.plain_registrations = remaining;
}
pub fn finalization_registry_has_registrations(registry: &JsFinalizationRegistry) -> bool {
!registry.registrations.is_empty() || !registry.plain_registrations.is_empty()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_finalization_registry_new_is_empty() {
let fr = finalization_registry_new();
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_drain_empty_returns_empty() {
let mut fr = finalization_registry_new();
assert!(finalization_registry_drain(&mut fr).is_empty());
}
#[test]
fn test_finalization_registry_register_valid() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let result = finalization_registry_register(&mut fr, &raw mut obj, JsValue::Smi(1), None);
assert!(result.is_ok());
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_register_null_target_error() {
let mut fr = finalization_registry_new();
let result =
finalization_registry_register(&mut fr, std::ptr::null_mut(), JsValue::Smi(1), None);
assert!(result.is_err());
}
#[test]
fn test_finalization_registry_register_null_target_is_type_error() {
let mut fr = finalization_registry_new();
let err =
finalization_registry_register(&mut fr, std::ptr::null_mut(), JsValue::Smi(1), None)
.unwrap_err();
assert!(matches!(err, StatorError::TypeError(_)));
}
#[test]
fn test_finalization_registry_register_null_token_error() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let result = finalization_registry_register(
&mut fr,
&raw mut obj,
JsValue::Smi(1),
Some(std::ptr::null_mut()),
);
assert!(result.is_err());
}
#[test]
fn test_finalization_registry_register_with_token() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let mut tok = HeapObject::new_null();
let result = finalization_registry_register(
&mut fr,
&raw mut obj,
JsValue::Smi(1),
Some(&raw mut tok),
);
assert!(result.is_ok());
}
#[test]
fn test_finalization_registry_register_multiple() {
let mut fr = finalization_registry_new();
let mut a = HeapObject::new_null();
let mut b = HeapObject::new_null();
finalization_registry_register(&mut fr, &raw mut a, JsValue::Smi(1), None).unwrap();
finalization_registry_register(&mut fr, &raw mut b, JsValue::Smi(2), None).unwrap();
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_register_same_target_twice() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let ptr = &raw mut obj;
finalization_registry_register(&mut fr, ptr, JsValue::Smi(1), None).unwrap();
finalization_registry_register(&mut fr, ptr, JsValue::Smi(2), None).unwrap();
finalization_registry_notify(&mut fr, ptr);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held.len(), 2);
}
#[test]
fn test_finalization_registry_unregister_removes_matching() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let mut tok = HeapObject::new_null();
let target = &raw mut obj;
let token = &raw mut tok;
finalization_registry_register(&mut fr, target, JsValue::Smi(1), Some(token)).unwrap();
assert!(finalization_registry_unregister(&mut fr, token).unwrap());
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_unregister_returns_false_when_no_match() {
let mut fr = finalization_registry_new();
let mut tok = HeapObject::new_null();
assert!(!finalization_registry_unregister(&mut fr, &raw mut tok).unwrap());
}
#[test]
fn test_finalization_registry_unregister_null_token_error() {
let mut fr = finalization_registry_new();
let result = finalization_registry_unregister(&mut fr, std::ptr::null_mut());
assert!(result.is_err());
}
#[test]
fn test_finalization_registry_unregister_only_removes_matching_token() {
let mut fr = finalization_registry_new();
let mut obj1 = HeapObject::new_null();
let mut obj2 = HeapObject::new_null();
let mut tok1 = HeapObject::new_null();
let mut tok2 = HeapObject::new_null();
finalization_registry_register(
&mut fr,
&raw mut obj1,
JsValue::Smi(1),
Some(&raw mut tok1),
)
.unwrap();
finalization_registry_register(
&mut fr,
&raw mut obj2,
JsValue::Smi(2),
Some(&raw mut tok2),
)
.unwrap();
finalization_registry_unregister(&mut fr, &raw mut tok1).unwrap();
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_unregister_removes_multiple_with_same_token() {
let mut fr = finalization_registry_new();
let mut obj1 = HeapObject::new_null();
let mut obj2 = HeapObject::new_null();
let mut tok = HeapObject::new_null();
let token = &raw mut tok;
finalization_registry_register(&mut fr, &raw mut obj1, JsValue::Smi(1), Some(token))
.unwrap();
finalization_registry_register(&mut fr, &raw mut obj2, JsValue::Smi(2), Some(token))
.unwrap();
assert!(finalization_registry_unregister(&mut fr, token).unwrap());
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_unregister_does_not_affect_no_token_entries() {
let mut fr = finalization_registry_new();
let mut obj1 = HeapObject::new_null();
let mut obj2 = HeapObject::new_null();
let mut tok = HeapObject::new_null();
let token = &raw mut tok;
finalization_registry_register(&mut fr, &raw mut obj1, JsValue::Smi(1), Some(token))
.unwrap();
finalization_registry_register(&mut fr, &raw mut obj2, JsValue::Smi(2), None).unwrap();
finalization_registry_unregister(&mut fr, token).unwrap();
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_notify_moves_to_cleanup_queue() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let target = &raw mut obj;
finalization_registry_register(&mut fr, target, JsValue::Smi(42), None).unwrap();
finalization_registry_notify(&mut fr, target);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::Smi(42)]);
}
#[test]
fn test_finalization_registry_notify_removes_registration() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let target = &raw mut obj;
finalization_registry_register(&mut fr, target, JsValue::Smi(1), None).unwrap();
finalization_registry_notify(&mut fr, target);
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_notify_null_is_noop() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
finalization_registry_register(&mut fr, &raw mut obj, JsValue::Smi(1), None).unwrap();
finalization_registry_notify(&mut fr, std::ptr::null_mut());
assert!(finalization_registry_has_registrations(&fr));
assert!(finalization_registry_drain(&mut fr).is_empty());
}
#[test]
fn test_finalization_registry_notify_only_matching_target() {
let mut fr = finalization_registry_new();
let mut a = HeapObject::new_null();
let mut b = HeapObject::new_null();
let pa = &raw mut a;
let pb = &raw mut b;
finalization_registry_register(&mut fr, pa, JsValue::Smi(1), None).unwrap();
finalization_registry_register(&mut fr, pb, JsValue::Smi(2), None).unwrap();
finalization_registry_notify(&mut fr, pa);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::Smi(1)]);
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_notify_multiple_registrations_same_target() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let ptr = &raw mut obj;
finalization_registry_register(&mut fr, ptr, JsValue::Smi(10), None).unwrap();
finalization_registry_register(&mut fr, ptr, JsValue::Smi(20), None).unwrap();
finalization_registry_notify(&mut fr, ptr);
let mut held = finalization_registry_drain(&mut fr);
held.sort_by_key(|v| match v {
JsValue::Smi(n) => *n,
_ => 0,
});
assert_eq!(held, vec![JsValue::Smi(10), JsValue::Smi(20)]);
}
#[test]
fn test_finalization_registry_notify_already_unregistered_is_noop() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let mut tok = HeapObject::new_null();
let target = &raw mut obj;
let token = &raw mut tok;
finalization_registry_register(&mut fr, target, JsValue::Smi(1), Some(token)).unwrap();
finalization_registry_unregister(&mut fr, token).unwrap();
finalization_registry_notify(&mut fr, target);
assert!(finalization_registry_drain(&mut fr).is_empty());
}
#[test]
fn test_finalization_registry_drain_clears_queue() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let target = &raw mut obj;
finalization_registry_register(&mut fr, target, JsValue::Smi(1), None).unwrap();
finalization_registry_notify(&mut fr, target);
let first = finalization_registry_drain(&mut fr);
assert_eq!(first.len(), 1);
let second = finalization_registry_drain(&mut fr);
assert!(second.is_empty());
}
#[test]
fn test_finalization_registry_drain_preserves_order_per_notify() {
let mut fr = finalization_registry_new();
let mut a = HeapObject::new_null();
let mut b = HeapObject::new_null();
let pa = &raw mut a;
let pb = &raw mut b;
finalization_registry_register(&mut fr, pa, JsValue::Smi(1), None).unwrap();
finalization_registry_register(&mut fr, pb, JsValue::Smi(2), None).unwrap();
finalization_registry_notify(&mut fr, pa);
finalization_registry_notify(&mut fr, pb);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::Smi(1), JsValue::Smi(2)]);
}
#[test]
fn test_finalization_registry_register_then_notify_then_drain() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let target = &raw mut obj;
finalization_registry_register(&mut fr, target, JsValue::Boolean(true), None).unwrap();
finalization_registry_notify(&mut fr, target);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::Boolean(true)]);
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_held_value_undefined() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let target = &raw mut obj;
finalization_registry_register(&mut fr, target, JsValue::Undefined, None).unwrap();
finalization_registry_notify(&mut fr, target);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::Undefined]);
}
#[test]
fn test_finalization_registry_held_value_string() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let target = &raw mut obj;
finalization_registry_register(&mut fr, target, JsValue::String("cleanup".into()), None)
.unwrap();
finalization_registry_notify(&mut fr, target);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::String("cleanup".into())]);
}
#[test]
fn test_finalization_registry_no_notify_means_no_drain() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
finalization_registry_register(&mut fr, &raw mut obj, JsValue::Smi(1), None).unwrap();
assert!(finalization_registry_drain(&mut fr).is_empty());
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_notify_nonexistent_target_is_noop() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let mut other = HeapObject::new_null();
finalization_registry_register(&mut fr, &raw mut obj, JsValue::Smi(1), None).unwrap();
finalization_registry_notify(&mut fr, &raw mut other);
assert!(finalization_registry_drain(&mut fr).is_empty());
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_register_plain_has_registrations() {
let mut fr = finalization_registry_new();
let obj = Rc::new(RefCell::new(PropertyMap::new()));
finalization_registry_register_plain(&mut fr, &obj, JsValue::Smi(1), None);
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_sweep_plain_after_drop() {
let mut fr = finalization_registry_new();
let obj = Rc::new(RefCell::new(PropertyMap::new()));
finalization_registry_register_plain(&mut fr, &obj, JsValue::Smi(42), None);
drop(obj);
finalization_registry_sweep_plain(&mut fr);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::Smi(42)]);
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_sweep_plain_alive_target_not_drained() {
let mut fr = finalization_registry_new();
let obj = Rc::new(RefCell::new(PropertyMap::new()));
finalization_registry_register_plain(&mut fr, &obj, JsValue::Smi(1), None);
finalization_registry_sweep_plain(&mut fr);
assert!(finalization_registry_drain(&mut fr).is_empty());
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_unregister_plain_removes_matching() {
let mut fr = finalization_registry_new();
let obj = Rc::new(RefCell::new(PropertyMap::new()));
let tok = Rc::new(RefCell::new(PropertyMap::new()));
finalization_registry_register_plain(&mut fr, &obj, JsValue::Smi(1), Some(&tok));
assert!(finalization_registry_unregister_plain(&mut fr, &tok));
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_unregister_plain_returns_false_no_match() {
let mut fr = finalization_registry_new();
let tok = Rc::new(RefCell::new(PropertyMap::new()));
assert!(!finalization_registry_unregister_plain(&mut fr, &tok));
}
#[test]
fn test_finalization_registry_sweep_plain_multiple_targets() {
let mut fr = finalization_registry_new();
let a = Rc::new(RefCell::new(PropertyMap::new()));
let b = Rc::new(RefCell::new(PropertyMap::new()));
finalization_registry_register_plain(&mut fr, &a, JsValue::Smi(1), None);
finalization_registry_register_plain(&mut fr, &b, JsValue::Smi(2), None);
drop(a);
finalization_registry_sweep_plain(&mut fr);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::Smi(1)]);
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_register_with_symbol_token_key() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let sym_token = TokenKey::Symbol(42);
let result = finalization_registry_register_with_token_key(
&mut fr,
&raw mut obj,
JsValue::Smi(1),
Some(sym_token),
);
assert!(result.is_ok());
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_unregister_by_symbol_key() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
let sym_key = TokenKey::Symbol(99);
finalization_registry_register_with_token_key(
&mut fr,
&raw mut obj,
JsValue::Smi(1),
Some(sym_key.clone()),
)
.unwrap();
assert!(finalization_registry_unregister_by_key(
&mut fr,
sym_key.clone()
));
assert!(!finalization_registry_has_registrations(&fr));
assert!(!finalization_registry_unregister_by_key(&mut fr, sym_key));
}
#[test]
fn test_finalization_registry_symbol_token_does_not_match_address() {
let mut fr = finalization_registry_new();
let mut obj = HeapObject::new_null();
finalization_registry_register_with_token_key(
&mut fr,
&raw mut obj,
JsValue::Smi(1),
Some(TokenKey::Symbol(42)),
)
.unwrap();
assert!(!finalization_registry_unregister_by_key(
&mut fr,
TokenKey::Address(42)
));
assert!(finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_register_plain_with_symbol_key() {
let mut fr = finalization_registry_new();
let obj = Rc::new(RefCell::new(PropertyMap::new()));
let sym_key = TokenKey::Symbol(7);
finalization_registry_register_plain_with_token_key(
&mut fr,
&obj,
JsValue::Smi(77),
Some(sym_key.clone()),
);
assert!(finalization_registry_has_registrations(&fr));
assert!(finalization_registry_unregister_by_key(&mut fr, sym_key));
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_unregister_by_key_removes_across_heap_and_plain() {
let mut fr = finalization_registry_new();
let mut obj1 = HeapObject::new_null();
let obj2 = Rc::new(RefCell::new(PropertyMap::new()));
let sym_key = TokenKey::Symbol(55);
finalization_registry_register_with_token_key(
&mut fr,
&raw mut obj1,
JsValue::Smi(1),
Some(sym_key.clone()),
)
.unwrap();
finalization_registry_register_plain_with_token_key(
&mut fr,
&obj2,
JsValue::Smi(2),
Some(sym_key.clone()),
);
assert!(finalization_registry_unregister_by_key(&mut fr, sym_key));
assert!(!finalization_registry_has_registrations(&fr));
}
#[test]
fn test_finalization_registry_symbol_sweep_plain_still_works() {
let mut fr = finalization_registry_new();
let obj = Rc::new(RefCell::new(PropertyMap::new()));
finalization_registry_register_plain_with_token_key(
&mut fr,
&obj,
JsValue::Smi(100),
Some(TokenKey::Symbol(10)),
);
drop(obj);
finalization_registry_sweep_plain(&mut fr);
let held = finalization_registry_drain(&mut fr);
assert_eq!(held, vec![JsValue::Smi(100)]);
}
#[test]
fn test_finalization_registry_register_with_token_key_null_target() {
let mut fr = finalization_registry_new();
let result = finalization_registry_register_with_token_key(
&mut fr,
std::ptr::null_mut(),
JsValue::Smi(1),
None,
);
assert!(result.is_err());
}
}