From 951059aa3a52442e0510a34151571b899c3f1038 Mon Sep 17 00:00:00 2001 From: Dennis Kobert Date: Fri, 21 Aug 2026 12:36:41 +0000 Subject: [PATCH] Drain lower-bound reducer subjects by guess-and-double --- node-graph/node-macro/src/codegen.rs | 63 ++++++++++++++++++++++------ node-graph/nodes/gcore/src/record.rs | 54 ++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 13 deletions(-) diff --git a/node-graph/node-macro/src/codegen.rs b/node-graph/node-macro/src/codegen.rs index 0e58879a2b..8f32919d56 100644 --- a/node-graph/node-macro/src/codegen.rs +++ b/node-graph/node-macro/src/codegen.rs @@ -1124,12 +1124,18 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn let non_exact = fail(quote!(#core_types::gpoll::GraphError::new("reduce over a non-exact extent"))); let batch_error = fail(quote!(__error)); let batch_failed = fail(quote!(#core_types::gpoll::GraphError::new("reduce batch failed"))); - // A reducer's subject on a deeper wire folds the flat span - // of the consumer's lane: the flat convention encodes the - // outer coordinate in the range, not the chain. + // A reducer's subject on a deeper wire folds the flat span of + // the consumer's lane: the flat convention encodes the outer + // coordinate in the range, not the chain. The span offset + // needs the exact inner count, so a lower bound only drains + // when the fold covers the whole wire. + let deeper = match derived_materialized(index) { + true => quote!(#core_types::node::Node::<#ctx_ident>::layout(&self.#name).depth > #levels), + false => quote!(false), + }; let start = match derived_materialized(index) { true => quote! { - match #core_types::node::Node::<#ctx_ident>::layout(&self.#name).depth > #levels { + match __deeper { true => #core_types::context::ExtractIndex::innermost_index(__input) * __count as u64, false => 0, } @@ -1138,18 +1144,49 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn }; quote! { let __arena = #core_types::context::ExtractArena::arena(__input); - let __count = match #core_types::node::Node::extent(&self.#name, __input, #core_types::gpoll::Level::Below(#levels)) { - #core_types::gpoll::GPoll::Final(#core_types::gpoll::Extent::Exactly(__count)) => __count, + let __deeper = #deeper; + let __sized = match #core_types::node::Node::extent(&self.#name, __input, #core_types::gpoll::Level::Below(#levels)) { + #core_types::gpoll::GPoll::Final(#core_types::gpoll::Extent::Exactly(__count)) => ::core::result::Result::Ok(__count), + #core_types::gpoll::GPoll::Final(#core_types::gpoll::Extent::AtLeast(__bound)) if !__deeper => ::core::result::Result::Err(__bound), #core_types::gpoll::GPoll::Pending => #pending, _ => #non_exact, }; - let __start: u64 = #start; - let __batch = match #core_types::record::materialize_batch(&self.#name, __input, __start..__start + __count as u64, __arena) { - #core_types::node::BatchStatus::Lent(__batch, ..) => __batch, - #core_types::node::BatchStatus::Filled(__batch, ..) => __batch.into_shared(), - #core_types::node::BatchStatus::Pending => #pending, - #core_types::node::BatchStatus::Error(__error) => #batch_error, - _ => #batch_failed, + let __batch = match __sized { + ::core::result::Result::Ok(__count) => { + let __start: u64 = #start; + match #core_types::record::materialize_batch(&self.#name, __input, __start..__start + __count as u64, __arena) { + #core_types::node::BatchStatus::Lent(__batch, ..) => __batch, + #core_types::node::BatchStatus::Filled(__batch, ..) => __batch.into_shared(), + #core_types::node::BatchStatus::Pending => #pending, + #core_types::node::BatchStatus::Error(__error) => #batch_error, + _ => #batch_failed, + } + } + // The count is a lower bound: drain by guess-and-double + // until a short fill, each reply's hint seeding the next + // guess. + ::core::result::Result::Err(__bound) => { + let mut __guess = __bound.max(16); + loop { + let (__batch, __hint) = match #core_types::record::materialize_batch(&self.#name, __input, 0..__guess as u64, __arena) { + #core_types::node::BatchStatus::Lent(__batch, _, __hint) => (__batch, __hint), + #core_types::node::BatchStatus::Filled(__batch, _, __hint) => (__batch.into_shared(), __hint), + #core_types::node::BatchStatus::Pending => #pending, + #core_types::node::BatchStatus::Error(__error) => #batch_error, + _ => #batch_failed, + }; + let __filled = __batch.len(); + if __filled < __guess { + break __batch; + } + match __hint { + #core_types::gpoll::Extent::Exactly(__total) if __total <= __filled => break __batch, + #core_types::gpoll::Extent::Exactly(__total) => __guess = __total, + #core_types::gpoll::Extent::AtLeast(__more) => __guess = (__guess * 2).max(__more), + #core_types::gpoll::Extent::Free => __guess *= 2, + } + } + } }; let #name = unsafe { #core_types::node::List::<#ty>::new(__batch) }; } diff --git a/node-graph/nodes/gcore/src/record.rs b/node-graph/nodes/gcore/src/record.rs index fb4fa9b2c6..e96b796594 100644 --- a/node-graph/nodes/gcore/src/record.rs +++ b/node-graph/nodes/gcore/src/record.rs @@ -481,6 +481,36 @@ mod tests { } } + /// A leveled source that keeps its count to itself: the extent is a lower + /// bound and lanes past the data answer the past-end signal. + struct DrainSourceNode { + layout: Layout, + count: usize, + } + + impl<'e> Node> for DrainSourceNode { + type Output = RecordValue<'e>; + + fn eval(&self, input: &ContextImpl<'e>) -> GPoll> { + let lane = input.innermost_index(); + if lane >= self.count as u64 { + return GPoll::past_end(); + } + let dst = stack::push(self.layout.frame_bytes()); + unsafe { dst.cast::().write(lane as f64) }; + 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::AtLeast(0)) + } + + fn layout(&self) -> &Layout { + &self.layout + } + } + struct IndexSourceNode { layout: Layout, } @@ -1488,6 +1518,30 @@ mod tests { assert_eq!(unsafe { out.rec(&value).element::() }, 42.); } + #[test] + fn reducer_drains_a_lower_bound_level() { + let arena = Arena::new(1 << 16).unwrap(); + let generations = []; + let scope = scope_fixture(&generations, &arena); + let ctx = ContextImpl::root(&scope); + + let leveled = leveled_f64_layout(&[]); + let out = f64_layout(&[]); + reserve_for(&[&leveled, &out]); + + // 5 lanes end inside the first guess; 20 force a full first fill, a + // hint-seeded regrow, and a short second fill. + for count in [5usize, 20] { + let source = DrainSourceNode { layout: leveled.clone(), count }; + let node = install_flip(SumNode::new(source, &leveled), &out); + let GPoll::Final(value) = node.eval(&ctx) else { + panic!("expected a final record at count {count}"); + }; + let expected = (count * (count - 1) / 2) as f64; + assert_eq!(unsafe { out.rec(&value).element::() }, expected, "count {count}"); + } + } + #[test] fn reducer_folds_a_repeated_level() { let arena = Arena::new(1024).unwrap();