#![allow(unsafe_code)]
use bincode_next::de::{BorrowDecode, BorrowDecoder};
use bincode_next::enc::{Encode, Encoder};
use bincode_next::error::{DecodeError, EncodeError};
use crate::dag::arena::DagArena;
use crate::dag::metadata::{NodeFlags, NodeHash, NodeMetadata};
use crate::dag::node::{ChildList, DagNode, DagNodeId};
use crate::dag::symbol::{FnId, OpKind, SymbolId, SymbolKind};
use crate::zerocopy::{AlignedBytes, BorrowedSlice, Pod, decode_zerocopy, encode_zerocopy};
use bincode_next::enc::write::Writer as BincodeWriter;
extern crate alloc;
use alloc::vec::Vec;
pub mod kind_tag {
pub const VARIABLE: u8 = 0;
pub const CONSTANT: u8 = 1;
pub const OPERATOR: u8 = 2;
pub const FUNCTION: u8 = 3;
}
pub mod op_tag {
pub const ADD: u8 = 0;
pub const SUB: u8 = 1;
pub const MUL: u8 = 2;
pub const DIV: u8 = 3;
pub const POW: u8 = 4;
pub const NEG: u8 = 5;
pub const MOD: u8 = 6;
}
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PackedDagNode {
pub hash: u64,
pub coefficient: f64,
pub child0: u32,
pub child1: u32,
pub kind_payload: u32,
pub arity: u8,
pub kind_tag: u8,
pub op_tag: u8,
pub flags: u8,
}
const _: () = assert!(
core::mem::size_of::<PackedDagNode>() == 32,
"PackedDagNode must be exactly 32 bytes"
);
const _: () = assert!(
core::mem::align_of::<PackedDagNode>() <= 8,
"PackedDagNode align must be ≤ 8 bytes"
);
unsafe impl Pod for PackedDagNode {}
impl PackedDagNode {
#[must_use]
pub const fn kind(&self) -> Option<SymbolKind> {
match self.kind_tag {
kind_tag::VARIABLE => Some(SymbolKind::Variable(SymbolId(self.kind_payload))),
kind_tag::CONSTANT => Some(SymbolKind::Constant(self.coefficient)),
kind_tag::OPERATOR => match self.op_tag {
op_tag::ADD => Some(SymbolKind::Operator(OpKind::Add)),
op_tag::SUB => Some(SymbolKind::Operator(OpKind::Sub)),
op_tag::MUL => Some(SymbolKind::Operator(OpKind::Mul)),
op_tag::DIV => Some(SymbolKind::Operator(OpKind::Div)),
op_tag::POW => Some(SymbolKind::Operator(OpKind::Pow)),
op_tag::NEG => Some(SymbolKind::Operator(OpKind::Neg)),
op_tag::MOD => Some(SymbolKind::Operator(OpKind::Mod)),
_ => None,
},
kind_tag::FUNCTION => Some(SymbolKind::Function(FnId(self.kind_payload))),
_ => None,
}
}
#[must_use]
pub const fn value(&self) -> Option<f64> {
if self.kind_tag == kind_tag::CONSTANT {
Some(self.coefficient)
} else {
None
}
}
#[must_use]
pub const fn meta(&self) -> NodeMetadata {
NodeMetadata {
hash: NodeHash(self.hash),
coefficient: if self.kind_tag == kind_tag::CONSTANT {
1.0
} else {
self.coefficient
},
flags: NodeFlags::from_bits(self.flags),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct PackedArenaImage {
nodes: Vec<PackedDagNode>,
children_pool: Vec<u32>,
pub rule_fingerprint: u64,
}
impl PackedArenaImage {
#[must_use]
#[allow(clippy::cast_possible_truncation)]
pub fn from_arena(arena: &DagArena) -> Self {
let n = arena.len();
let mut nodes: Vec<PackedDagNode> = Vec::with_capacity(n);
let mut children_pool: Vec<u32> = Vec::new();
for i in 0..n {
let id = DagNodeId::new(i as u32);
let node = arena
.get(id)
.map_or_else(empty_placeholder, pack_one_node_ref);
let packed = match node {
PackOne::Inline(p) => p,
PackOne::WithPool(mut p, extra) => {
#[allow(clippy::cast_possible_truncation)]
{
p.child0 = children_pool.len() as u32;
}
if p.arity == 255 {
#[allow(clippy::cast_possible_truncation)]
children_pool.push(extra.len() as u32);
}
children_pool.extend(extra.iter().map(|id| id.value()));
p
}
};
nodes.push(packed);
}
Self {
nodes,
children_pool,
rule_fingerprint: 0,
}
}
#[must_use]
pub const fn len(&self) -> usize {
self.nodes.len()
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
#[must_use]
pub const fn nodes(&self) -> &[PackedDagNode] {
self.nodes.as_slice()
}
#[must_use]
pub const fn children_pool(&self) -> &[u32] {
self.children_pool.as_slice()
}
pub fn encode(&self) -> Result<AlignedBytes, EncodeError> {
let view = SerializableView {
nodes: BorrowedSlice::new(&self.nodes),
children_pool: BorrowedSlice::new(&self.children_pool),
rule_fingerprint: self.rule_fingerprint,
};
encode_zerocopy(view)
}
pub fn write_to<W: std::io::Write>(&self, writer: W) -> Result<(), EncodeError> {
let view = SerializableView {
nodes: BorrowedSlice::new(&self.nodes),
children_pool: BorrowedSlice::new(&self.children_pool),
rule_fingerprint: self.rule_fingerprint,
};
struct IoWriter<W: std::io::Write>(W);
impl<W: std::io::Write> BincodeWriter for IoWriter<W> {
fn write(&mut self, bytes: &[u8]) -> Result<(), EncodeError> {
self.0
.write_all(bytes)
.map_err(|e| EncodeError::OtherString(alloc::format!("I/O error: {e}")))
}
}
bincode_next::encode_into_writer(
view,
&mut IoWriter(writer),
crate::zerocopy::zerocopy_config(),
)
}
}
enum PackOne {
Inline(PackedDagNode),
WithPool(PackedDagNode, Vec<DagNodeId>),
}
fn pack_one_node_ref(node: &DagNode) -> PackOne {
let mut packed = PackedDagNode {
hash: node.meta.hash.0,
coefficient: if let SymbolKind::Constant(val) = node.kind {
val
} else {
node.meta.coefficient
},
child0: u32::MAX,
child1: u32::MAX,
kind_payload: u32::MAX,
arity: 0,
kind_tag: 0,
op_tag: 0,
flags: node.meta.flags.bits(),
};
match node.kind {
SymbolKind::Variable(sym) => {
packed.kind_tag = kind_tag::VARIABLE;
packed.kind_payload = sym.0;
}
SymbolKind::Constant(_) => {
packed.kind_tag = kind_tag::CONSTANT;
}
SymbolKind::Operator(op) => {
packed.kind_tag = kind_tag::OPERATOR;
packed.op_tag = match op {
OpKind::Add => op_tag::ADD,
OpKind::Sub => op_tag::SUB,
OpKind::Mul => op_tag::MUL,
OpKind::Div => op_tag::DIV,
OpKind::Pow => op_tag::POW,
OpKind::Neg => op_tag::NEG,
OpKind::Mod => op_tag::MOD,
};
}
SymbolKind::Function(fn_id) => {
packed.kind_tag = kind_tag::FUNCTION;
packed.kind_payload = fn_id.0;
}
SymbolKind::ControlFlow(_) => {
packed.kind_tag = kind_tag::FUNCTION;
packed.kind_payload = u32::MAX;
}
}
let children = node.children.as_slice();
let arity = children.len();
#[allow(clippy::cast_possible_truncation)]
{
packed.arity = if arity > 255 { 255 } else { arity as u8 };
}
match arity {
0 => PackOne::Inline(packed),
1 => {
packed.child0 = children[0].value();
PackOne::Inline(packed)
}
2 => {
packed.child0 = children[0].value();
packed.child1 = children[1].value();
PackOne::Inline(packed)
}
_ => PackOne::WithPool(packed, children.to_vec()),
}
}
const fn empty_placeholder() -> PackOne {
PackOne::Inline(PackedDagNode {
hash: 0,
coefficient: 0.0,
child0: u32::MAX,
child1: u32::MAX,
kind_payload: u32::MAX,
arity: 0,
kind_tag: kind_tag::CONSTANT,
op_tag: 0,
flags: 0,
})
}
struct SerializableView<'a> {
nodes: BorrowedSlice<'a, PackedDagNode>,
children_pool: BorrowedSlice<'a, u32>,
rule_fingerprint: u64,
}
impl Encode for SerializableView<'_> {
fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
self.rule_fingerprint.encode(encoder)?;
self.nodes.encode(encoder)?;
self.children_pool.encode(encoder)?;
Ok(())
}
}
#[derive(Debug, Clone, Copy)]
pub struct BorrowedArenaView<'a> {
nodes: BorrowedSlice<'a, PackedDagNode>,
children_pool: BorrowedSlice<'a, u32>,
pub rule_fingerprint: u64,
}
impl<'a> BorrowedArenaView<'a> {
pub fn decode(bytes: &'a AlignedBytes) -> Result<Self, DecodeError> {
decode_zerocopy::<Self>(bytes)
}
#[must_use]
pub const fn nodes(&self) -> &'a [PackedDagNode] {
self.nodes.as_slice()
}
#[must_use]
pub fn get(&self, id: DagNodeId) -> Option<&'a PackedDagNode> {
if id.is_none() {
return None;
}
self.nodes.as_slice().get(id.index())
}
#[must_use]
pub fn children(&self, node: &PackedDagNode) -> Children<'a> {
let arity = node.arity as usize;
match arity {
0 => Children::Inline {
ids: [DagNodeId::NONE; 2],
len: 0,
},
1 => Children::Inline {
ids: [DagNodeId::new(node.child0), DagNodeId::NONE],
len: 1,
},
2 => Children::Inline {
ids: [DagNodeId::new(node.child0), DagNodeId::new(node.child1)],
len: 2,
},
_ => {
let pool = self.children_pool.as_slice();
let start = node.child0 as usize;
if arity == 255 {
let Some(&count) = pool.get(start) else {
return Children::Pool(&[]);
};
let data_start = start + 1;
let data_end = data_start + count as usize;
let slice = pool.get(data_start..data_end).unwrap_or(&[]);
Children::Pool(slice)
} else {
let slice = pool.get(start..start + arity).unwrap_or(&[]);
Children::Pool(slice)
}
}
}
}
#[must_use]
pub const fn len(&self) -> usize {
self.nodes.len()
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
#[must_use]
pub fn to_owned_arena(&self) -> DagArena {
let mut arena = DagArena::new();
for packed in self.nodes() {
let kind = packed.kind().unwrap_or(SymbolKind::Constant(0.0));
let meta = NodeMetadata {
hash: NodeHash(packed.hash),
coefficient: if matches!(kind, SymbolKind::Constant(_)) {
1.0
} else {
packed.coefficient
},
flags: NodeFlags::from_bits(packed.flags),
};
let kids: Vec<DagNodeId> = self.children(packed).iter().collect();
let children = ChildList::from_slice(&kids);
let node = match kind {
SymbolKind::Constant(val) => DagNode::constant(val, meta),
SymbolKind::Variable(_) => DagNode::variable(kind, meta),
SymbolKind::Operator(_) | SymbolKind::Function(_) | SymbolKind::ControlFlow(_) => {
DagNode::operator(kind, meta, children)
}
};
arena.alloc(node);
}
arena
}
}
impl<'de> BorrowDecode<'de, ()> for BorrowedArenaView<'de> {
fn borrow_decode<D: BorrowDecoder<'de, Context = ()>>(
decoder: &mut D,
) -> Result<Self, DecodeError> {
let rule_fingerprint = u64::borrow_decode(decoder)?;
let nodes = BorrowedSlice::<PackedDagNode>::borrow_decode(decoder)?;
let children_pool = BorrowedSlice::<u32>::borrow_decode(decoder)?;
Ok(Self {
nodes,
children_pool,
rule_fingerprint,
})
}
}
#[derive(Debug, Clone, Copy)]
pub enum Children<'a> {
Inline {
ids: [DagNodeId; 2],
len: usize,
},
Pool(&'a [u32]),
}
impl<'a> Children<'a> {
#[must_use]
pub const fn len(&self) -> usize {
match self {
Self::Inline { len, .. } => *len,
Self::Pool(s) => s.len(),
}
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn iter(&self) -> ChildIter<'a> {
match self {
Self::Inline { ids, len } => ChildIter::Inline {
ids: *ids,
len: *len,
idx: 0,
},
Self::Pool(slice) => ChildIter::Pool(slice.iter()),
}
}
}
impl<'a> IntoIterator for &Children<'a> {
type Item = DagNodeId;
type IntoIter = ChildIter<'a>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
pub enum ChildIter<'a> {
Inline {
ids: [DagNodeId; 2],
len: usize,
idx: usize,
},
Pool(core::slice::Iter<'a, u32>),
}
impl Iterator for ChildIter<'_> {
type Item = DagNodeId;
fn next(&mut self) -> Option<DagNodeId> {
match self {
Self::Inline { ids, len, idx } => {
if *idx >= *len {
return None;
}
let id = ids[*idx];
*idx += 1;
Some(id)
}
Self::Pool(it) => it.next().copied().map(DagNodeId::new),
}
}
}
#[must_use]
pub fn children_to_child_list(children: Children<'_>) -> ChildList {
let ids: Vec<DagNodeId> = children.iter().collect();
ChildList::from_slice(&ids)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dag::builder::DagBuilder;
#[test]
fn packed_node_is_32_bytes() {
assert_eq!(core::mem::size_of::<PackedDagNode>(), 32);
assert_eq!(core::mem::align_of::<PackedDagNode>(), 8);
}
#[test]
fn pack_then_view_roundtrip() {
let mut builder = DagBuilder::new();
let a = builder.variable("a");
let b = builder.variable("b");
let sum = builder.add(a, b);
let coeff = builder.constant(2.5);
let _root = builder.mul(sum, coeff);
let image = PackedArenaImage::from_arena(builder.arena());
assert_eq!(image.len(), builder.arena().len());
let bytes = image.encode().expect("encode");
let view = BorrowedArenaView::decode(&bytes).expect("decode");
assert_eq!(view.len(), image.len());
for i in 0..view.len() {
let id = DagNodeId::new(i as u32);
let original = builder.arena().get(id).expect("rich node");
let packed = view.get(id).expect("packed node");
assert_eq!(packed.kind(), Some(original.kind), "kind mismatch at {i}");
let orig_value = if let SymbolKind::Constant(v) = original.kind {
Some(v)
} else {
None
};
assert_eq!(packed.value(), orig_value, "value mismatch at {i}");
assert_eq!(packed.meta().hash, original.meta.hash);
assert_eq!(packed.arity as usize, original.children.len());
let kids: Vec<DagNodeId> = view.children(packed).iter().collect();
assert_eq!(kids.as_slice(), original.children.as_slice());
}
}
#[test]
fn variadic_children_spill_to_pool() {
let mut builder = DagBuilder::new();
let a = builder.variable("a");
let b = builder.variable("b");
let c = builder.variable("c");
let d = builder.variable("d");
let e = builder.variable("e");
let big = builder.operator(
SymbolKind::Function(FnId(0)),
&[a, b, c, d, e],
NodeFlags::EMPTY,
);
let image = PackedArenaImage::from_arena(builder.arena());
let bytes = image.encode().expect("encode");
let view = BorrowedArenaView::decode(&bytes).expect("decode");
let node = view.get(big).expect("get big");
assert_eq!(node.arity, 5);
let kids: Vec<DagNodeId> = view.children(node).iter().collect();
assert_eq!(kids, alloc::vec![a, b, c, d, e]);
assert_ne!(node.child0, u32::MAX);
}
#[test]
fn decoded_nodes_borrow_from_buffer() {
let mut builder = DagBuilder::new();
let x = builder.variable("x");
let _ = builder.mul(x, x);
let image = PackedArenaImage::from_arena(builder.arena());
let bytes = image.encode().expect("encode");
let view = BorrowedArenaView::decode(&bytes).expect("decode");
let buf_start = bytes.as_bytes().as_ptr() as usize;
let buf_end = buf_start + bytes.as_bytes().len();
let nodes_start = view.nodes().as_ptr() as usize;
assert!(
nodes_start >= buf_start && nodes_start < buf_end,
"BorrowedArenaView.nodes() does not point into the input \
buffer — zero-copy invariant violated"
);
}
}