Files
endorphin-aster-vm/crates/objects/src/heap.rs
T

786 lines
26 KiB
Rust

use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::fmt;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::thread::ThreadId;
use crate::changes::MAX_PENDING_CHANGES;
use crate::{ChangeBatch, ChangeKey, ObjectChange};
use crate::{EndorObject, ObjectHeader, ObjectPayload, Value, VmError};
static NEXT_HEAP_ID: AtomicU64 = AtomicU64::new(1);
/// Checked, process-local reference. Copying it never retains an object.
/// The generation prevents a released handle from aliasing a reused slot.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ObjectRef {
heap_id: u64,
index: usize,
generation: u64,
}
impl fmt::Display for ObjectRef {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}:{}:{}", self.heap_id, self.index, self.generation)
}
}
/// A non-owning reference. Only the originating store can upgrade it; slot
/// reuse never makes a dead weak reference point at a different object.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct WeakObjectRef(pub(crate) ObjectRef);
#[derive(Debug)]
struct Slot {
generation: u64,
object: Option<Box<EndorObject>>,
}
#[derive(Debug)]
struct ReadOnly {
protected: BTreeSet<ObjectRef>,
reads: Mutex<BTreeSet<ObjectRef>>,
regions: Mutex<Vec<BTreeMap<ObjectRef, bool>>>,
dependencies: Mutex<BTreeMap<ReadKey, ReadDependency>>,
dependency_regions: Mutex<Vec<BTreeMap<ReadKey, ReadDependency>>>,
}
#[derive(Debug)]
pub(crate) struct Heap {
id: u64,
slots: Vec<Slot>,
free: Vec<usize>,
live: usize,
changes: BTreeMap<(ObjectRef, ChangeKey), u64>,
changes_overflow: bool,
read_only: Vec<ReadOnly>,
validation_reads: Mutex<HashMap<ThreadId, usize>>,
}
/// Borrowed read identity; content_read distinguishes ordinary reads from
/// annotation validation alone. Neither form retains the object.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ReadAccess {
pub object: ObjectRef,
pub content_read: bool,
}
/// Borrowed observations; Object handles never retain the observed graph.
/// Named values are immutable comparison snapshots, not managed owning edges.
#[derive(Debug, Clone, PartialEq)]
pub enum ReadDependency {
Object {
object: ObjectRef,
version: u64,
},
DictionaryKind {
object: ObjectRef,
},
Name {
object: ObjectRef,
name: String,
value: Option<crate::Value>,
},
}
impl ReadDependency {
pub fn object(&self) -> ObjectRef {
match self {
Self::Object { object, .. }
| Self::DictionaryKind { object }
| Self::Name { object, .. } => *object,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
enum ReadKey {
Object(ObjectRef),
DictionaryKind(ObjectRef),
Name(ObjectRef, String),
}
#[derive(Debug)]
pub struct TrackedReads {
pub accesses: Vec<ReadAccess>,
pub dependencies: Vec<ReadDependency>,
}
#[derive(Debug)]
pub struct ReadSet {
pub objects: Vec<ObjectRef>,
pub dependencies: Vec<ReadDependency>,
}
struct ValidationScope<'a> {
heap: &'a Heap,
thread: ThreadId,
}
impl Drop for ValidationScope<'_> {
fn drop(&mut self) {
let mut active = self
.heap
.validation_reads
.lock()
.expect("read tracker poisoned");
let depth = active
.get_mut(&self.thread)
.expect("active validation scope");
*depth -= 1;
if *depth == 0 {
active.remove(&self.thread);
}
}
}
impl Heap {
pub(crate) fn new() -> Result<Self, VmError> {
let id = NEXT_HEAP_ID
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
.map_err(|_| VmError::IdentityExhausted)?;
Ok(Self {
id,
slots: Vec::new(),
free: Vec::new(),
live: 0,
changes: BTreeMap::new(),
changes_overflow: false,
read_only: vec![],
validation_reads: Mutex::new(HashMap::new()),
})
}
pub(crate) fn live_references(&self) -> impl Iterator<Item = ObjectRef> + '_ {
self.slots.iter().enumerate().filter_map(|(index, slot)| {
slot.object.as_ref().map(|_| ObjectRef {
heap_id: self.id,
index,
generation: slot.generation,
})
})
}
pub(crate) fn begin_read_only(&mut self) -> usize {
let protected = self
.slots
.iter()
.enumerate()
.filter_map(|(index, slot)| {
slot.object.as_ref().map(|_| ObjectRef {
heap_id: self.id,
index,
generation: slot.generation,
})
})
.collect();
self.read_only.push(ReadOnly {
protected,
reads: Mutex::new(BTreeSet::new()),
regions: Mutex::new(Vec::new()),
dependencies: Mutex::new(BTreeMap::new()),
dependency_regions: Mutex::new(Vec::new()),
});
self.read_only.len()
}
pub(crate) fn end_read_only(&mut self, scope: usize) -> Result<Vec<ObjectRef>, VmError> {
Ok(self.end_read_only_precise(scope)?.objects)
}
pub(crate) fn end_read_only_precise(&mut self, scope: usize) -> Result<ReadSet, VmError> {
if scope == 0 || scope != self.read_only.len() {
return Err(VmError::InvalidArgument("read-only scope nesting mismatch"));
}
let scope = self.read_only.pop().unwrap();
Ok(ReadSet {
objects: scope
.reads
.into_inner()
.expect("read tracker poisoned")
.into_iter()
.collect(),
dependencies: scope
.dependencies
.into_inner()
.expect("read tracker poisoned")
.into_values()
.collect(),
})
}
pub(crate) fn begin_read_tracking(&self) -> Result<usize, VmError> {
let scope = self.read_only.last().ok_or(VmError::InvalidArgument(
"read tracking requires a read-only scope",
))?;
let mut regions = scope.regions.lock().expect("read tracker poisoned");
regions.push(BTreeMap::new());
scope
.dependency_regions
.lock()
.expect("read tracker poisoned")
.push(BTreeMap::new());
Ok(regions.len())
}
pub(crate) fn end_read_tracking(&self, token: usize) -> Result<Vec<ObjectRef>, VmError> {
Ok(self
.end_read_tracking_details(token)?
.into_iter()
.map(|read| read.object)
.collect())
}
pub(crate) fn end_read_tracking_details(
&self,
token: usize,
) -> Result<Vec<ReadAccess>, VmError> {
Ok(self.end_read_tracking_precise(token)?.accesses)
}
pub(crate) fn end_read_tracking_precise(&self, token: usize) -> Result<TrackedReads, VmError> {
let scope = self.read_only.last().ok_or(VmError::InvalidArgument(
"read tracking requires a read-only scope",
))?;
let mut regions = scope.regions.lock().expect("read tracker poisoned");
let mut precise = scope
.dependency_regions
.lock()
.expect("read tracker poisoned");
if token == 0 || token != regions.len() || token != precise.len() {
return Err(VmError::InvalidArgument("read tracking nesting mismatch"));
}
Ok(TrackedReads {
accesses: regions
.pop()
.unwrap()
.into_iter()
.map(|(object, content_read)| ReadAccess {
object,
content_read,
})
.collect(),
dependencies: precise.pop().unwrap().into_values().collect(),
})
}
pub(crate) fn is_read_only(&self) -> bool {
!self.read_only.is_empty()
}
pub(crate) fn with_validation_reads<T>(
&self,
compute: impl FnOnce() -> Result<T, VmError>,
) -> Result<T, VmError> {
if self.read_only.is_empty() {
return compute();
}
let thread = std::thread::current().id();
{
let mut active = self.validation_reads.lock().expect("read tracker poisoned");
let depth = active.entry(thread).or_default();
*depth = depth
.checked_add(1)
.ok_or(VmError::ExecutionLimit("validation scopes"))?;
}
let _scope = ValidationScope { heap: self, thread };
compute()
}
pub(crate) fn record_read(&self, reference: ObjectRef, version: u64) {
if self.read_only.is_empty() {
return;
}
self.record_observation(reference, ReadKey::Object(reference), || {
ReadDependency::Object {
object: reference,
version,
}
});
}
pub(crate) fn record_dictionary_kind(&self, reference: ObjectRef) {
self.record_observation(reference, ReadKey::DictionaryKind(reference), || {
ReadDependency::DictionaryKind { object: reference }
});
}
pub(crate) fn record_name(
&self,
reference: ObjectRef,
name: &str,
value: &Option<crate::Value>,
) {
if self.read_only.is_empty() {
return;
}
self.record_observation(reference, ReadKey::Name(reference, name.into()), || {
ReadDependency::Name {
object: reference,
name: name.into(),
value: value.clone(),
}
});
}
fn record_observation(
&self,
reference: ObjectRef,
key: ReadKey,
build: impl Fn() -> ReadDependency,
) {
if self.read_only.is_empty() {
return;
}
let content = !self
.validation_reads
.lock()
.expect("read tracker poisoned")
.contains_key(&std::thread::current().id());
for scope in &self.read_only {
if scope.protected.contains(&reference) {
scope
.dependencies
.lock()
.expect("read tracker poisoned")
.entry(key.clone())
.or_insert_with(&build);
for region in scope
.dependency_regions
.lock()
.expect("read tracker poisoned")
.iter_mut()
{
region.entry(key.clone()).or_insert_with(&build);
}
scope
.reads
.lock()
.expect("read tracker poisoned")
.insert(reference);
for region in scope
.regions
.lock()
.expect("read tracker poisoned")
.iter_mut()
{
region
.entry(reference)
.and_modify(|ordinary| *ordinary |= content)
.or_insert(content);
}
}
}
}
pub(crate) fn allocate(
&mut self,
payload: Box<dyn ObjectPayload>,
) -> Result<ObjectRef, VmError> {
let mut edges = Vec::new();
payload.trace(&mut |reference| edges.push(reference));
let increments = self.prepare_increments(&edges)?;
let index = self.free.pop().unwrap_or(self.slots.len());
if index == self.slots.len() {
self.slots.push(Slot {
generation: 0,
object: None,
});
}
let reference = ObjectRef {
heap_id: self.id,
index,
generation: self.slots[index].generation,
};
let object = Box::new(EndorObject {
header: ObjectHeader {
identity: reference,
strong_refs: 1,
external_refs: 1,
change_version: 0,
type_object: None,
},
payload,
attributes: BTreeMap::new(),
payload_edges: edges,
});
self.apply_increments(increments);
self.slots[index].object = Some(object);
self.live += 1;
Ok(reference)
}
pub(crate) fn get(&self, reference: ObjectRef) -> Result<&EndorObject, VmError> {
self.validate_owner(reference)?;
self.slots
.get(reference.index)
.filter(|slot| slot.generation == reference.generation)
.and_then(|slot| slot.object.as_deref())
.ok_or(VmError::InvalidHandle)
}
pub(crate) fn attach_type(
&mut self,
reference: ObjectRef,
type_object: ObjectRef,
) -> Result<(), VmError> {
if self.get(reference)?.header.type_object == Some(type_object) {
return Ok(());
}
if self.get(reference)?.header.type_object.is_some() {
return Err(VmError::InvalidArgument("object already has a type"));
}
if self.get(reference)?.attributes.contains_key("type") {
return Err(VmError::InvalidArgument(
"type attribute conflicts with managed type",
));
}
let increments = self.prepare_increments(&[type_object])?;
self.prepare_change(reference)?;
self.apply_increments(increments);
self.get_mut(reference)?.header.type_object = Some(type_object);
self.record_change(reference, ChangeKey::Contents);
Ok(())
}
/// Transactional copy-on-update for built-in containers. The EndorObject
/// box and identity stay stable; only its payload and edge snapshot change.
/// All new edges are validated/retained before any old edge is released.
pub(crate) fn replace_payload(
&mut self,
reference: ObjectRef,
payload: impl ObjectPayload,
) -> Result<(), VmError> {
self.get(reference)?;
let mut edges = Vec::new();
payload.trace(&mut |target| edges.push(target));
let increments = self.prepare_increments(&edges)?;
self.prepare_change(reference)?;
self.apply_increments(increments);
let object = self.get_mut(reference)?;
let old_edges = std::mem::replace(&mut object.payload_edges, edges);
let old_payload = std::mem::replace(&mut object.payload, Box::new(payload));
self.record_change(reference, ChangeKey::Contents);
self.release_edges(old_edges);
drop(old_payload);
Ok(())
}
fn get_mut(&mut self, reference: ObjectRef) -> Result<&mut EndorObject, VmError> {
self.get(reference)?;
Ok(self.slots[reference.index].object.as_deref_mut().unwrap())
}
fn validate_owner(&self, reference: ObjectRef) -> Result<(), VmError> {
if reference.heap_id != self.id {
return Err(VmError::ForeignHandle);
}
Ok(())
}
fn prepare_increments(
&self,
references: &[ObjectRef],
) -> Result<BTreeMap<ObjectRef, usize>, VmError> {
let mut counts = BTreeMap::<ObjectRef, usize>::new();
for reference in references {
self.get(*reference)?;
let count = counts.entry(*reference).or_default();
*count = count.checked_add(1).ok_or(VmError::RefCountOverflow)?;
}
for (reference, count) in &counts {
self.get(*reference)?
.header
.strong_refs
.checked_add(*count)
.ok_or(VmError::RefCountOverflow)?;
}
Ok(counts)
}
fn apply_increments(&mut self, increments: BTreeMap<ObjectRef, usize>) {
for (reference, count) in increments {
self.get_mut(reference).unwrap().header.strong_refs += count;
}
}
pub(crate) fn retain(&mut self, reference: ObjectRef) -> Result<(), VmError> {
let header = &mut self.get_mut(reference)?.header;
let strong = header
.strong_refs
.checked_add(1)
.ok_or(VmError::RefCountOverflow)?;
let external = header
.external_refs
.checked_add(1)
.ok_or(VmError::RefCountOverflow)?;
header.strong_refs = strong;
header.external_refs = external;
Ok(())
}
pub(crate) fn release(&mut self, reference: ObjectRef) -> Result<(), VmError> {
let header = &mut self.get_mut(reference)?.header;
if header.external_refs == 0 {
return Err(VmError::NoOwnedReference);
}
header.external_refs -= 1;
header.strong_refs -= 1;
let pending = if header.strong_refs == 0 {
vec![reference]
} else {
Vec::new()
};
self.drain_zero_counts(pending);
Ok(())
}
pub(crate) fn set_attr(
&mut self,
reference: ObjectRef,
name: String,
value: Value,
) -> Result<(), VmError> {
self.check_attributes(reference)?;
let edges: Vec<_> = match &value {
Value::Object(target) => vec![*target],
_ => Vec::new(),
};
let increments = self.prepare_increments(&edges)?;
self.prepare_change(reference)?;
self.apply_increments(increments);
let old = self
.get_mut(reference)?
.attributes
.insert(name.clone(), value);
self.record_change(reference, ChangeKey::Attribute(name));
if let Some(Value::Object(target)) = old {
self.release_edges(vec![target]);
}
Ok(())
}
pub(crate) fn get_attr(&self, reference: ObjectRef, name: &str) -> Result<Value, VmError> {
self.check_attributes(reference)?;
self.get(reference)?
.attributes
.get(name)
.cloned()
.ok_or_else(|| VmError::AttributeNotFound(name.into()))
}
pub(crate) fn del_attr(&mut self, reference: ObjectRef, name: &str) -> Result<(), VmError> {
self.check_attributes(reference)?;
self.prepare_change(reference)?;
let old = self
.get_mut(reference)?
.attributes
.remove(name)
.ok_or_else(|| VmError::AttributeNotFound(name.into()))?;
self.record_change(reference, ChangeKey::Attribute(name.into()));
if let Value::Object(target) = old {
self.release_edges(vec![target]);
}
Ok(())
}
pub(crate) fn prepare_change(&self, reference: ObjectRef) -> Result<(), VmError> {
self.get(reference)?;
if self
.read_only
.iter()
.any(|scope| scope.protected.contains(&reference))
{
return Err(VmError::UnsupportedOperation(
"mutation of existing object during read-only call",
));
}
self.get(reference)?
.header
.change_version
.checked_add(1)
.ok_or(VmError::ExecutionLimit("object change version"))?;
Ok(())
}
pub(crate) fn record_change(&mut self, reference: ObjectRef, key: ChangeKey) {
let object = self.get_mut(reference).expect("validated change target");
object.header.change_version += 1;
let version = object.header.change_version;
// Fresh result construction has no existing observers. Keep versions,
// but do not notify application dependencies about derived allocations.
if self.is_read_only() {
return;
}
if self.changes_overflow {
return;
}
let entry = (reference, key);
if self.changes.len() >= MAX_PENDING_CHANGES && !self.changes.contains_key(&entry) {
self.changes.clear();
self.changes_overflow = true;
} else {
self.changes.insert(entry, version);
}
}
pub(crate) fn take_changes(&mut self) -> ChangeBatch {
ChangeBatch {
changes: std::mem::take(&mut self.changes)
.into_iter()
.map(|((object, key), version)| ObjectChange {
object,
key,
version,
})
.collect(),
rescan_required: std::mem::take(&mut self.changes_overflow),
}
}
fn check_attributes(&self, reference: ObjectRef) -> Result<(), VmError> {
if !self.get(reference)?.payload.allows_attributes() {
return Err(VmError::UnsupportedOperation("attributes"));
}
Ok(())
}
fn release_edges(&mut self, edges: Vec<ObjectRef>) {
let mut pending = Vec::new();
for target in edges {
let header = &mut self.get_mut(target).unwrap().header;
header.strong_refs -= 1;
if header.strong_refs == 0 {
pending.push(target);
}
}
self.drain_zero_counts(pending);
}
/// Removes the object before running Drop. Reuse increments its generation;
/// an exhausted generation permanently retires that slot.
fn take(&mut self, reference: ObjectRef) -> Box<EndorObject> {
let slot = &mut self.slots[reference.index];
let object = slot.object.take().unwrap();
if let Some(next) = slot.generation.checked_add(1) {
slot.generation = next;
self.free.push(reference.index);
}
self.live -= 1;
object
}
fn drain_zero_counts(&mut self, mut pending: Vec<ObjectRef>) {
while let Some(reference) = pending.pop() {
let object = self.take(reference);
for target in object.edges() {
let header = &mut self.get_mut(target).unwrap().header;
header.strong_refs -= 1;
if header.strong_refs == 0 {
pending.push(target);
}
}
drop(object);
}
}
/// Full, explicit cycle collection. External references are roots; every
/// managed attribute and declared payload edge participates in tracing.
pub(crate) fn collect_cycles(&mut self) -> usize {
let mut marked = vec![false; self.slots.len()];
let mut pending = Vec::new();
for slot in &self.slots {
if let Some(object) = &slot.object
&& object.header.external_refs > 0
{
pending.push(object.header.identity);
}
}
while let Some(reference) = pending.pop() {
if marked[reference.index] {
continue;
}
marked[reference.index] = true;
pending.extend(self.get(reference).unwrap().edges());
}
let garbage: Vec<_> = self
.slots
.iter()
.enumerate()
.filter(|(index, _)| !marked[*index])
.filter_map(|(_, slot)| slot.object.as_ref().map(|object| object.header.identity))
.collect();
let objects: Vec<_> = garbage
.iter()
.map(|reference| self.take(*reference))
.collect();
// All garbage slots are detached before touching edges or running Drop.
for object in &objects {
for target in object.edges() {
if let Ok(survivor) = self.get_mut(target) {
survivor.header.strong_refs -= 1;
}
}
}
let collected = objects.len();
drop(objects);
collected
}
pub(crate) fn len(&self) -> usize {
self.live
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Record;
#[test]
fn invalid_slots_and_generations_are_rejected() {
let mut heap = Heap::new().unwrap();
let reference = heap.allocate(Box::new(Record)).unwrap();
let invalid = ObjectRef {
index: usize::MAX,
..reference
};
let stale = ObjectRef {
generation: 1,
..reference
};
assert!(matches!(heap.get(invalid), Err(VmError::InvalidHandle)));
assert!(matches!(heap.get_mut(stale), Err(VmError::InvalidHandle)));
}
#[test]
fn reused_slots_never_alias_and_exhausted_generations_retire() {
let mut heap = Heap::new().unwrap();
let first = heap.allocate(Box::new(Record)).unwrap();
heap.release(first).unwrap();
let second = heap.allocate(Box::new(Record)).unwrap();
assert_eq!(first.index, second.index);
assert_ne!(first.generation, second.generation);
assert!(matches!(heap.get(first), Err(VmError::InvalidHandle)));
heap.slots[second.index].generation = u64::MAX;
let last = ObjectRef {
generation: u64::MAX,
..second
};
heap.slots[last.index]
.object
.as_mut()
.unwrap()
.header
.identity = last;
heap.release(last).unwrap();
let next = heap.allocate(Box::new(Record)).unwrap();
assert_ne!(next.index, last.index);
}
#[test]
fn failed_reference_increments_leave_counts_and_attributes_unchanged() {
let mut heap = Heap::new().unwrap();
let owner = heap.allocate(Box::new(Record)).unwrap();
let target = heap.allocate(Box::new(Record)).unwrap();
heap.get_mut(target).unwrap().header.strong_refs = usize::MAX;
assert_eq!(heap.retain(target), Err(VmError::RefCountOverflow));
assert_eq!(heap.get(target).unwrap().header.external_refs, 1);
assert_eq!(
heap.set_attr(owner, "edge".into(), Value::Object(target)),
Err(VmError::RefCountOverflow)
);
assert!(heap.get(owner).unwrap().attributes.is_empty());
}
}