diff --git a/node-graph/libraries/core-types/src/extent.rs b/node-graph/libraries/core-types/src/extent.rs index 39b2117b9e..40d06fba11 100644 --- a/node-graph/libraries/core-types/src/extent.rs +++ b/node-graph/libraries/core-types/src/extent.rs @@ -75,4 +75,10 @@ impl LevelIn { pub fn pushed(&self) -> bool { self.level + 1 == self.depth } + + /// Whether the query targets the topmost level; `pushed` under the name a + /// non-creator (concat, remap) reads naturally. + pub fn top(&self) -> bool { + self.pushed() + } } diff --git a/node-graph/libraries/core-types/src/gpoll.rs b/node-graph/libraries/core-types/src/gpoll.rs index 071010fb16..fdd6692699 100644 --- a/node-graph/libraries/core-types/src/gpoll.rs +++ b/node-graph/libraries/core-types/src/gpoll.rs @@ -185,6 +185,16 @@ impl Extent { _ => Extent::Free, }) } + + /// The sum of two extents, used to concatenate a level; a free operand + /// counts as one lane, so a scalar edge joins a concat as a single item. + pub fn sum(a: GPoll, b: GPoll) -> GPoll { + let lanes = |extent| match extent { + Extent::Exactly(count) => count, + Extent::Free => 1, + }; + a.zip(b).map(|(a, b)| Extent::Exactly(lanes(a) + lanes(b))) + } } /// A query over a node's nesting levels: one level, the product below or above diff --git a/node-graph/libraries/core-types/src/node.rs b/node-graph/libraries/core-types/src/node.rs index 3c289cbd64..f9fe7b8b3e 100644 --- a/node-graph/libraries/core-types/src/node.rs +++ b/node-graph/libraries/core-types/src/node.rs @@ -490,6 +490,15 @@ impl<'a, N> LazyInput<'a, N> { { self.cell.eval_input(self.input_index, self.node, ctx) } + + /// The edge's composite extent, for kernels that split or shift indices + /// over their sources. + pub fn extent(&self, ctx: &Input, at: Level) -> GPoll + where + N: Node, + { + self.node.extent(ctx, at) + } } #[cfg(test)] diff --git a/node-graph/nodes/gcore/src/record.rs b/node-graph/nodes/gcore/src/record.rs index f5d9a60d73..4ede80920c 100644 --- a/node-graph/nodes/gcore/src/record.rs +++ b/node-graph/nodes/gcore/src/record.rs @@ -8,7 +8,7 @@ use core_types::attribute::{Attr, Opacity, RemoveAttr}; use core_types::context::{DeriveCtx, ExtractIndex, IndexLink, InjectIndex}; use core_types::extent::{ExtentIn, LevelIn, ValueIn}; -use core_types::gpoll::{ErrorKind, Extent, GPoll, GraphError, Interrupt}; +use core_types::gpoll::{ErrorKind, Extent, GPoll, GraphError, Interrupt, Level}; use core_types::Ctx; core_types::attribute! { @@ -132,6 +132,39 @@ fn repeat_faded_extent(content: ExtentIn<'_>, count: ValueIn<'_, u32>, level: Le } } +/// Rank-model Extend: the output's top level is `base`'s lanes followed by +/// `new`'s, each side evaluated within its own index range. +#[node_macro::node(category("Test"), extent(extend_extent))] +fn extend( + ctx: impl Ctx + ExtractIndex + InjectIndex + Copy, + base: impl Node, Output = T>, + new: impl Node, Output = T>, +) -> Result { + let split = match base.extent(ctx, Level::Total) { + GPoll::Final(Extent::Exactly(count)) => count as u64, + GPoll::Pending => return Err(Interrupt::Pending), + _ => return Err(GraphError::new("extend over a non-exact base extent").into()), + }; + let lane = ctx.innermost_index(); + match lane < split { + true => base.eval(ctx), + false => { + let mut shifted = *ctx; + shifted.set_index(lane - split); + new.eval(&shifted) + } + } +} + +/// The top level sums both sides; inner levels forward the base's, which the +/// new side must match (rectangular). +fn extend_extent(base: ExtentIn<'_>, new: ExtentIn<'_>, level: LevelIn) -> GPoll { + match level.top() { + true => Extent::sum(base.at(level), new.at(level)), + false => base.at(level), + } +} + #[node_macro::node(category("Test"))] fn source_opacity(_: impl Ctx, _: (), element: f64, opacity: f64) -> (f64, Attr) { (element, Attr(opacity)) @@ -264,6 +297,37 @@ mod tests { } } + struct LeveledSourceNode { + layout: Layout, + elements: Vec, + field: Option<(usize, f64)>, + } + + impl<'e> Node> for LeveledSourceNode { + type Output = RecordValue<'e>; + + fn eval(&self, input: &ContextImpl<'e>) -> GPoll> { + let element = self.elements[input.innermost_index() as usize % self.elements.len()]; + let dst = stack::push(self.layout.frame_bytes()); + unsafe { + dst.cast::().write(element); + if let Some((offset, value)) = self.field { + dst.add(offset).cast::().write(value); + } + } + stack::pop(dst); + GPoll::Final(RecordValue::spilled(unsafe { Rec::new(dst.cast_const()) })) + } + + fn extent_at(&self, _input: &ContextImpl<'e>, _level: u8) -> GPoll { + GPoll::Final(Extent::Exactly(self.elements.len())) + } + + fn layout(&self) -> &Layout { + &self.layout + } + } + struct IndexSourceNode { layout: Layout, } @@ -597,6 +661,82 @@ mod tests { } } + #[test] + fn extend_concatenates_the_top_level_and_fills_the_union() { + let arena = Arena::new(1024).unwrap(); + let generations = []; + let scope = scope_fixture(&generations, &arena); + let ctx = ContextImpl::root(&scope); + + let leveled_f64_layout = |names: &[&'static str]| { + let writes: Vec = names + .iter() + .map(|name| core_types::record::FieldWrite { + name, + level: 0, + size: 8, + align: 8, + read_erased: ::read_erased, + repark: None, + }) + .collect(); + Layout::default().with_writes(1, core_types::record::element_write::(), &writes) + }; + let base_layout = leveled_f64_layout(&[Opacity::NAME]); + let new_layout = leveled_f64_layout(&[Length::NAME]); + let union = Layout::union(&[&base_layout, &new_layout]); + reserve_for(&[&base_layout, &new_layout, &union]); + + let base = LeveledSourceNode { + layout: base_layout.clone(), + elements: vec![10., 11.], + field: Some((base_layout.offset_of(Opacity::NAME, 0).unwrap(), 0.5)), + }; + let new = LeveledSourceNode { + layout: new_layout.clone(), + elements: vec![100., 101., 102.], + field: Some((new_layout.offset_of(Length::NAME, 0).unwrap(), 7.)), + }; + let meta = core_types::record::LayoutMeta { + sources: vec![0, 1], + reads: vec![], + element: core_types::record::ElementSpec::Carried, + writes: vec![], + removes: vec![], + level_delta: 0, + }; + let node = install( + ExtendNode::new(RecordSource::new(base, &base_layout, &union), RecordSource::new(new, &new_layout, &union), &union), + meta, + &[Some(&base_layout), Some(&new_layout)], + ); + let out = Node::::layout(&node).clone(); + assert_eq!(out.depth, 1); + assert_eq!(node.extent_at(&ctx, 0), GPoll::Final(Extent::Exactly(5)), "the top level sums both sides"); + + let head = ctx.index_head(); + let expected = [10., 11., 100., 101., 102.]; + for (lane, &element) in expected.iter().enumerate() { + let mark = stack::sp(); + let scoped = ctx.promoted(&head, lane as u64); + let GPoll::Final(value) = node.eval(&scoped) else { + panic!("expected a final record"); + }; + let rec = out.rec(&value); + assert_eq!(unsafe { rec.element::() }, element); + let opacity = unsafe { rec.read::(out.offset_of(Opacity::NAME, 0).unwrap()) }; + let length = unsafe { rec.read::(out.offset_of(Length::NAME, 0).unwrap()) }; + match lane < 2 { + // The base side wrote its opacity; length fills from the census. + true => assert_eq!((opacity, length), (0.5, 0.)), + // The new side wrote its length; opacity fills from the census. + false => assert_eq!((opacity, length), (1., 7.)), + } + // SAFETY: the element and attrs were read out above, so no borrow into this lane's frames remains. + unsafe { stack::rewind(mark) }; + } + } + #[test] fn reducer_folds_a_repeated_level() { let arena = Arena::new(1024).unwrap();