use super::node::{DagNode, DagNodeId};
use bincode_next::{Decode, Encode};
#[derive(Debug, Clone, Encode, Decode)]
pub struct DagArena {
nodes: Vec<DagNode>,
}
impl DagArena {
#[must_use]
pub const fn new() -> Self {
Self { nodes: Vec::new() }
}
pub fn alloc(&mut self, node: DagNode) -> DagNodeId {
let index = self.nodes.len();
self.nodes.push(node);
#[allow(clippy::cast_possible_truncation)]
DagNodeId(index as u32)
}
#[must_use]
pub fn get(&self, id: DagNodeId) -> Option<&DagNode> {
if id.is_none() {
return None;
}
self.nodes.get(id.index())
}
pub fn get_mut(&mut self, id: DagNodeId) -> Option<&mut DagNode> {
if id.is_none() {
return None;
}
self.nodes.get_mut(id.index())
}
#[must_use]
pub const fn len(&self) -> usize {
self.nodes.len()
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
pub fn clear(&mut self) {
self.nodes.clear();
}
}
impl Default for DagArena {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dag::metadata::{NodeHash, NodeMetadata};
#[test]
fn arena_alloc_and_retrieve() {
let mut arena = DagArena::new();
assert!(arena.is_empty());
let meta = NodeMetadata::leaf(NodeHash(42));
let node = DagNode::constant(1.0, meta);
let id = arena.alloc(node);
assert_eq!(arena.len(), 1);
assert!(!arena.is_empty());
let retrieved = arena.get(id).unwrap();
assert!(
matches!(retrieved.kind, crate::dag::symbol::SymbolKind::Constant(v) if (v - 1.0).abs() < f64::EPSILON)
);
assert_eq!(retrieved.meta.hash, NodeHash(42));
let retrieved_mut = arena.get_mut(id).unwrap();
retrieved_mut.meta.hash = NodeHash(43);
assert_eq!(arena.get(id).unwrap().meta.hash, NodeHash(43));
}
}