Materialize ranked non-subject inputs whole and expose them to extent overrides

This commit is contained in:
Dennis Kobert
2026-08-17 09:26:59 +00:00
parent a03db0ac6c
commit 0f3e3f7337
5 changed files with 81 additions and 6 deletions

View File

@@ -42,6 +42,23 @@ impl<'a> ExtentIn<'a> {
}
}
/// A ranked (`IList`) input materialized whole: `get` drives the batch and
/// yields the level as a [`List`](crate::node::List), for extents that depend
/// on the input's data rather than its counts alone.
pub struct ListIn<'a, T> {
get: &'a dyn Fn() -> GPoll<crate::node::List<'a, T>>,
}
impl<'a, T> ListIn<'a, T> {
pub fn new(get: &'a dyn Fn() -> GPoll<crate::node::List<'a, T>>) -> Self {
Self { get }
}
pub fn get(&self) -> GPoll<crate::node::List<'a, T>> {
(self.get)()
}
}
/// The queried absolute level (innermost `0`), paired with the node's depth.
#[derive(Clone, Copy, Debug)]
pub struct LevelIn {

View File

@@ -177,7 +177,7 @@ pub struct List<'a, T> {
_element: PhantomData<T>,
}
impl<'a, T: Copy> List<'a, T> {
impl<'a, T> List<'a, T> {
/// # Safety
/// `T` must be the batch's record element type, proven at the consumer's wiring.
pub unsafe fn new(batch: RecordBatch<'a>) -> Self {
@@ -192,12 +192,30 @@ impl<'a, T: Copy> List<'a, T> {
self.batch.is_empty()
}
pub fn get(&self, index: usize) -> T {
pub fn get(&self, index: usize) -> T
where
T: Copy,
{
// SAFETY: `List::new` established that `T` is the batch's element type.
unsafe { self.batch.get(index).element::<T>() }
}
pub fn iter(&self) -> impl Iterator<Item = T> + '_ {
/// Borrows lane `index`'s element, through the park for droppable types.
pub fn element_ref(&self, index: usize) -> &T {
// SAFETY: `List::new` established that `T` is the batch's element type,
// and the borrow lives within the batch's own lifetime.
unsafe { crate::record::borrow_element::<T>(self.batch.get(index).rec()) }
}
/// Lane `index`'s record, for attribute reads beside the element.
pub fn lane(&self, index: usize) -> RecordLane<'a> {
self.batch.get(index)
}
pub fn iter(&self) -> impl Iterator<Item = T> + '_
where
T: Copy,
{
(0..self.len()).map(move |index| self.get(index))
}
}

View File

@@ -978,6 +978,9 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
let record_value_ty: Type = syn::parse_quote!(#core_types::record::RecordValue<'__record>);
let node_bounds = regular_fields.iter().enumerate().zip(&node_generics).map(|((index, field), node_generic)| match &field.ty {
// A ranked input rides a record edge whatever the node kind; the
// materialized batch reads its lanes.
ParsedFieldType::Regular(_) if ir::materialized_levels(&node, index) > 0 => quote!(#node_generic: #core_types::node::Node<#ctx_ident, Output = #record_value_ty>),
ParsedFieldType::Regular(_) if flip => quote!(#node_generic: #core_types::node::Node<#ctx_ident, Output = #record_value_ty>),
ParsedFieldType::Node(_) if flip => match derives {
true => quote! {
@@ -1283,6 +1286,29 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
_ => extent_edge(&query, &arg),
},
ParsedFieldType::Regular(RegularParsedField { ty, .. }) => match ir::value_binding(&node, index) {
// A ranked input materializes whole, so a data-dependent
// extent can walk its lanes.
ValueBinding::Materialized => {
let levels = ir::materialized_levels(&node, index);
quote! {
let #query = || {
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,
#core_types::gpoll::GPoll::Pending => return #core_types::gpoll::GPoll::Pending,
_ => return #core_types::gpoll::GPoll::Error(::std::boxed::Box::new(#core_types::gpoll::GraphError::new("extent over a non-exact ranked input"))),
};
match #core_types::record::materialize_batch(&self.#name, __input, 0..__count as u64, __arena) {
#core_types::node::BatchStatus::Lent(__batch, _) => #core_types::gpoll::GPoll::Final(unsafe { #core_types::node::List::<#ty>::new(__batch) }),
#core_types::node::BatchStatus::Filled(__batch, _) => #core_types::gpoll::GPoll::Final(unsafe { #core_types::node::List::<#ty>::new(__batch.into_shared()) }),
#core_types::node::BatchStatus::Pending => #core_types::gpoll::GPoll::Pending,
#core_types::node::BatchStatus::Error(__error) => #core_types::gpoll::GPoll::Error(::std::boxed::Box::new(__error)),
_ => #core_types::gpoll::GPoll::Error(::std::boxed::Box::new(#core_types::gpoll::GraphError::new("extent could not materialize a ranked input"))),
}
};
let #arg = #core_types::extent::ListIn::new(&#query);
}
}
ValueBinding::RecordElement | ValueBinding::ReadingSecondary => {
let slot = format_ident!("__in_{index}");
quote! {

View File

@@ -188,6 +188,8 @@ enum SlotKind {
Value(Type),
/// A record edge whose element extracts to the node's plain value input.
Extracted(Type),
/// A ranked record edge consumed whole; no layout rides to the constructor.
Ranked(Type),
/// A plain value edge.
Plain(Type),
/// A lazy node edge.
@@ -227,6 +229,7 @@ fn single_row_entries(parsed: &ParsedNodeFn, struct_name: &Ident, regular_fields
match &field.ty {
ParsedFieldType::Node(NodeParsedField { output_type, .. }) => SlotKind::Lazy(output_type.clone()),
ParsedFieldType::Regular(RegularParsedField { ty, .. }) => match ir::value_binding(&node, index) {
ir::ValueBinding::Materialized => SlotKind::Ranked(ty.clone()),
ir::ValueBinding::ReadingSecondary | ir::ValueBinding::RecordElement => SlotKind::Value(ty.clone()),
// One wire kind: a record node's plain value still rides a
// record edge, extracted to its element at construction.
@@ -240,7 +243,9 @@ fn single_row_entries(parsed: &ParsedNodeFn, struct_name: &Ident, regular_fields
// Every non-base value/plain/lazy input must be concrete.
let values_concrete = regular_fields.iter().zip(&slots).all(|(field, slot)| match slot {
SlotKind::BaseGeneric(_) | SlotKind::BaseConcrete(_) => true,
SlotKind::Value(ty) | SlotKind::Extracted(ty) | SlotKind::Plain(ty) | SlotKind::Lazy(ty) => !contains_open_generic(parsed, ty) && (lend(field) || !type_disqualifies(ty)),
SlotKind::Value(ty) | SlotKind::Extracted(ty) | SlotKind::Ranked(ty) | SlotKind::Plain(ty) | SlotKind::Lazy(ty) => {
!contains_open_generic(parsed, ty) && (lend(field) || !type_disqualifies(ty))
}
});
if !values_concrete {
return quote!();
@@ -252,7 +257,7 @@ fn single_row_entries(parsed: &ParsedNodeFn, struct_name: &Ident, regular_fields
let input_types = slots.iter().map(|slot| match slot {
SlotKind::BaseGeneric(name) => quote!(gcore::registry::generic_record_edge_type(#name)),
SlotKind::BaseConcrete(ty) | SlotKind::Value(ty) | SlotKind::Extracted(ty) => quote!(gcore::registry::record_edge_type::<#ty>()),
SlotKind::BaseConcrete(ty) | SlotKind::Value(ty) | SlotKind::Extracted(ty) | SlotKind::Ranked(ty) => quote!(gcore::registry::record_edge_type::<#ty>()),
SlotKind::Plain(ty) | SlotKind::Lazy(ty) => quote!(gcore::registry::edge_type::<#ty>()),
});
@@ -277,6 +282,9 @@ fn single_row_entries(parsed: &ParsedNodeFn, struct_name: &Ident, regular_fields
let #layout = #handle.layout().clone();
let #name = gcore::record::RecordExtract::<#value_ty, _>::new(#handle.downcast_record::<#value_ty>()?, &#layout);
},
SlotKind::Ranked(value_ty) => quote! {
let #name = inputs.next().unwrap().downcast_record::<#value_ty>()?;
},
SlotKind::Plain(value_ty) | SlotKind::Lazy(value_ty) => quote!(let #name = inputs.next().unwrap().downcast::<#value_ty>()?;),
}
});

View File

@@ -322,7 +322,13 @@ fn has_attr_io(node: &Node) -> bool {
/// Levels of `input[index]` the output does not carry; `> 0` folds the input
/// into a `List` before the kernel.
pub(crate) fn materialized_levels(node: &Node, index: usize) -> u8 {
node.inputs[index].shape.depth.saturating_sub(node.output.shape.depth)
let input = &node.inputs[index];
// A ranked subject folds the levels the output collapses; a ranked
// non-subject input is consumed whole.
match input.subject {
true => input.shape.depth.saturating_sub(node.output.shape.depth),
false => input.shape.depth,
}
}
pub(crate) fn value_binding(node: &Node, index: usize) -> ValueBinding {