From 264601c5f73f5f388f50f6d91fee148c786061f5 Mon Sep 17 00:00:00 2001 From: Dennis Kobert Date: Tue, 11 Aug 2026 09:26:36 +0000 Subject: [PATCH] Return records from eval_batch through a RecordBatch view --- node-graph/libraries/core-types/src/node.rs | 141 ++++++++++++++++-- .../libraries/core-types/src/registry.rs | 2 +- node-graph/node-macro/src/codegen.rs | 2 +- node-graph/nodes/math/src/lib.rs | 6 +- 4 files changed, 136 insertions(+), 15 deletions(-) diff --git a/node-graph/libraries/core-types/src/node.rs b/node-graph/libraries/core-types/src/node.rs index fc4df24799..514cccb75b 100644 --- a/node-graph/libraries/core-types/src/node.rs +++ b/node-graph/libraries/core-types/src/node.rs @@ -6,8 +6,8 @@ use std::ops::Range; #[derive(Debug)] pub enum BatchStatus<'a, T> { - Lent(&'a [T], Finality), - Filled(FilledBatch<'a, T>, Finality), + Lent(RecordBatch<'a, T>, Finality), + Filled(RecordBatch<'a, T>, Finality), Pending, Error(GraphError), NeedBuffer, @@ -57,6 +57,121 @@ pub unsafe fn assume_init_prefix_mut(scratch: &mut [MaybeUninit], len: usi unsafe { std::slice::from_raw_parts_mut(scratch.as_mut_ptr().cast::(), len) } } +/// A borrow-for-scope view over a batch of records whose element type is `T`, +/// paired with their shared [`Layout`](crate::record::Layout). Row-major backed +/// today; the interface (`len`/`layout`/`get`/`for_each`) is storage-agnostic so +/// a columnar backing can replace it without touching consumers. +#[derive(Debug)] +pub struct RecordBatch<'a, T> { + lanes: LaneStore<'a, T>, + layout: &'a crate::record::Layout, +} + +#[derive(Debug)] +enum LaneStore<'a, T> { + /// Borrows resident storage (the `Lent` status): no drop obligation. + Borrowed(&'a [T]), + /// Owns the caller scratch's initialized prefix (the `Filled` status). + Owned(FilledBatch<'a, T>), +} + +impl<'a, T> RecordBatch<'a, T> { + pub fn lent(values: &'a [T], layout: &'a crate::record::Layout) -> Self { + Self { lanes: LaneStore::Borrowed(values), layout } + } + + pub fn filled(filled: FilledBatch<'a, T>, layout: &'a crate::record::Layout) -> Self { + Self { lanes: LaneStore::Owned(filled), layout } + } + + fn lanes(&self) -> &[T] { + match &self.lanes { + LaneStore::Borrowed(values) => values, + LaneStore::Owned(filled) => filled.values(), + } + } + + pub fn len(&self) -> usize { + self.lanes().len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn layout(&self) -> &crate::record::Layout { + self.layout + } + + /// Lends lane `lane`'s record to `f` for the callback's scope only. + pub fn get(&self, lane: usize, f: impl FnOnce(RecordLane<'_, T>) -> R) -> R { + f(RecordLane { value: &self.lanes()[lane], layout: self.layout }) + } + + /// Lends every lane's record in order, each for its callback's scope only. + pub fn for_each(&self, mut f: impl FnMut(usize, RecordLane<'_, T>)) { + for (lane, value) in self.lanes().iter().enumerate() { + f(lane, RecordLane { value, layout: self.layout }); + } + } + + /// Hands the owned scratch prefix back to the caller, cancelling the drop + /// obligation. Panics on a lent batch, which owns nothing to return. + pub fn into_values(self) -> &'a mut [T] { + match self.lanes { + LaneStore::Owned(filled) => filled.into_values(), + LaneStore::Borrowed(_) => panic!("into_values on a lent batch"), + } + } +} + +/// One lane's record, lent for a callback scope. Derefs to the raw lane value; +/// for record elements, [`rec`](RecordLane::rec) and [`attr`](RecordLane::attr) +/// read the record through its layout. +#[derive(Debug)] +pub struct RecordLane<'r, T> { + value: &'r T, + layout: &'r crate::record::Layout, +} + +impl std::ops::Deref for RecordLane<'_, T> { + type Target = T; + + fn deref(&self) -> &T { + self.value + } +} + +impl RecordLane<'_, T> { + pub fn layout(&self) -> &crate::record::Layout { + self.layout + } +} + +impl<'e> RecordLane<'_, crate::record::RecordValue<'e>> { + /// The record pointer, resolved through the layout. + pub fn rec(&self) -> crate::record::Rec { + self.layout.rec(self.value) + } + + /// The element at offset 0. + /// + /// # Safety + /// `U` must be the record's element type, proven at the consumer's wiring. + pub unsafe fn element(&self) -> U { + unsafe { self.rec().element::() } + } + + /// Attribute `A` at the record's top level, or its census default when the + /// layout does not carry it. + pub fn attr(&self) -> A::Value<'e> { + match self.layout.offset_of(A::NAME, 0) { + Some(offset) => unsafe { self.rec().read::>(offset) }, + None => A::default(), + } + } +} + pub trait Node { type Output; @@ -79,7 +194,7 @@ pub trait Node { crate::record::empty_layout() } - fn eval_batch<'a>(&self, input: &'a Input, range: Range, scratch: Option<&'a mut [MaybeUninit]>) -> BatchStatus<'a, Self::Output> + fn eval_batch<'a>(&'a self, input: &'a Input, range: Range, scratch: Option<&'a mut [MaybeUninit]>) -> BatchStatus<'a, Self::Output> where Input: InjectIndex + Copy, { @@ -119,7 +234,7 @@ pub trait Node { } } // SAFETY: all `len` lanes were written by the loop above. - BatchStatus::Filled(unsafe { FilledBatch::new(scratch, len) }, finality) + BatchStatus::Filled(RecordBatch::filled(unsafe { FilledBatch::new(scratch, len) }, self.layout()), finality) } } @@ -145,7 +260,7 @@ where (**self).layout() } - fn eval_batch<'a>(&self, input: &'a Input, range: Range, scratch: Option<&'a mut [MaybeUninit]>) -> BatchStatus<'a, Self::Output> + fn eval_batch<'a>(&'a self, input: &'a Input, range: Range, scratch: Option<&'a mut [MaybeUninit]>) -> BatchStatus<'a, Self::Output> where Input: InjectIndex + Copy, { @@ -175,7 +290,7 @@ where (**self).layout() } - fn eval_batch<'a>(&self, input: &'a Input, range: Range, scratch: Option<&'a mut [MaybeUninit]>) -> BatchStatus<'a, Self::Output> + fn eval_batch<'a>(&'a self, input: &'a Input, range: Range, scratch: Option<&'a mut [MaybeUninit]>) -> BatchStatus<'a, Self::Output> where Input: InjectIndex + Copy, { @@ -205,7 +320,7 @@ where (**self).layout() } - fn eval_batch<'a>(&self, input: &'a Input, range: Range, scratch: Option<&'a mut [MaybeUninit]>) -> BatchStatus<'a, Self::Output> + fn eval_batch<'a>(&'a self, input: &'a Input, range: Range, scratch: Option<&'a mut [MaybeUninit]>) -> BatchStatus<'a, Self::Output> where Input: InjectIndex + Copy, { @@ -337,10 +452,12 @@ mod tests { let input = TestInput { index: 0 }; let mut scratch = [const { MaybeUninit::uninit() }; 4]; let status = Double.eval_batch(&input, 2..6, Some(&mut scratch)); - let BatchStatus::Filled(lanes, finality) = status else { + let BatchStatus::Filled(batch, finality) = status else { panic!("expected filled, got {status:?}"); }; - assert_eq!(lanes.values(), &[4, 6, 8, 10]); + let mut got = Vec::new(); + batch.for_each(|_, lane| got.push(*lane)); + assert_eq!(got, vec![4, 6, 8, 10]); assert_eq!(finality, Finality::AllFinal); } @@ -399,10 +516,12 @@ mod tests { let input = TestInput { index: 0 }; let mut scratch = [const { MaybeUninit::uninit() }; 4]; let status = PartialAtThree.eval_batch(&input, 0..4, Some(&mut scratch)); - let BatchStatus::Filled(lanes, finality) = status else { + let BatchStatus::Filled(batch, finality) = status else { panic!("expected filled, got {status:?}"); }; - assert_eq!(lanes.values(), &[0, 1, 2, 3]); + let mut got = Vec::new(); + batch.for_each(|_, lane| got.push(*lane)); + assert_eq!(got, vec![0, 1, 2, 3]); assert_eq!(finality, Finality::Partial); } diff --git a/node-graph/libraries/core-types/src/registry.rs b/node-graph/libraries/core-types/src/registry.rs index a0ad8837b1..b25cff2c4b 100644 --- a/node-graph/libraries/core-types/src/registry.rs +++ b/node-graph/libraries/core-types/src/registry.rs @@ -180,7 +180,7 @@ where unsafe { self.ptr.as_ref() }.layout() } - fn eval_batch<'a>(&self, input: &'a Input, range: std::ops::Range, scratch: Option<&'a mut [std::mem::MaybeUninit]>) -> crate::node::BatchStatus<'a, Self::Output> + fn eval_batch<'a>(&'a self, input: &'a Input, range: std::ops::Range, scratch: Option<&'a mut [std::mem::MaybeUninit]>) -> crate::node::BatchStatus<'a, Self::Output> where Input: crate::context::InjectIndex + Copy, { diff --git a/node-graph/node-macro/src/codegen.rs b/node-graph/node-macro/src/codegen.rs index 5535a9c213..b5b6450578 100644 --- a/node-graph/node-macro/src/codegen.rs +++ b/node-graph/node-macro/src/codegen.rs @@ -1258,7 +1258,7 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn let batch_impl = match &parsed.attributes.batch { Some(path) => quote! { fn eval_batch<'__batch>( - &self, + &'__batch self, __input: &'__batch #ctx_ident, __range: ::std::ops::Range, __scratch: Option<&'__batch mut [::std::mem::MaybeUninit]>, diff --git a/node-graph/nodes/math/src/lib.rs b/node-graph/nodes/math/src/lib.rs index f485f44d7d..469e2bff5d 100644 --- a/node-graph/nodes/math/src/lib.rs +++ b/node-graph/nodes/math/src/lib.rs @@ -1107,10 +1107,12 @@ mod graphene_test { let erased: Box> = Box::new(AddNode::new(IndexNode, SourceNode(10.0f64))); let mut scratch = [const { MaybeUninit::uninit() }; 4]; let status = erased.eval_batch(&ctx, 2..6, Some(&mut scratch)); - let BatchStatus::Filled(lanes, finality) = status else { + let BatchStatus::Filled(batch, finality) = status else { panic!("expected filled, got {status:?}"); }; - assert_eq!(lanes.values(), &[12.0, 13.0, 14.0, 15.0]); + let mut got = Vec::new(); + batch.for_each(|_, lane| got.push(*lane)); + assert_eq!(got, vec![12.0, 13.0, 14.0, 15.0]); assert_eq!(finality, Finality::AllFinal); }