Add the rank-model repeat on points flattening every points row

node-graph/nodes/repeat/src/leveled.rs

@@ -5,11 +5,12 @@
 use core::f64::consts::TAU;
 use core_types::attribute::{Attr, Transform};
 use core_types::context::{ExtractIndex, IndexLink};
-use core_types::extent::{ExtentIn, LevelIn, ValueIn};
-use core_types::gpoll::{Extent, GPoll, Interrupt};
+use core_types::extent::{ExtentIn, LevelIn, ListIn, ValueIn};
+use core_types::gpoll::{Extent, GPoll, GraphError, Interrupt};
 use core_types::registry::types::{Angle, PixelSize};
-use core_types::{Ctx, DeriveCtx};
+use core_types::{Ctx, DeriveCtx, InjectIndex};
 use glam::{DAffine2, DVec2};
+use graphic_types::Vector;
 
 /// The rank-model Repeat Array: each copy evaluates the lazy content at its
 /// own index and composes the linear step onto the row's transform.
@@ -83,6 +84,57 @@ fn repeat_radial_extent(content: ExtentIn<'_>, _start_angle: ValueIn<'_, Angle>,
 	}
 }
 
+/// The rank-model Repeat on Points: the pushed level flattens every point of
+/// every points row, mirroring the eager iteration order (rows in order, a
+/// row's points reversed when `reverse` is set); each copy evaluates the
+/// content with its point's transformed position pushed, then lands the
+/// content row's transform on that position.
+#[node_macro::node(category("Test"), name("Repeat on Points"), extent(repeat_on_points_extent))]
+fn repeat_on_points<T>(
+	ctx: impl Ctx + DeriveCtx + ExtractIndex + InjectIndex + Copy,
+	content: impl Node<Context<'_>, Output = (T, Attr<Transform>)>,
+	points: IList<Vector>,
+	reverse: bool,
+) -> Result<IList<(T, Attr<Transform>)>, Interrupt> {
+	let inner = content.inner_extent(ctx)?;
+	let (copy, rest) = ctx.split_innermost(inner);
+
+	let mut remaining = copy as usize;
+	for row_index in 0..points.len() {
+		let vector = points.element_ref(row_index);
+		let positions = vector.point_domain.positions();
+		if remaining >= positions.len() {
+			remaining -= positions.len();
+			continue;
+		}
+		let index = match reverse {
+			true => positions.len() - 1 - remaining,
+			false => remaining,
+		};
+		let transform: DAffine2 = points.lane(row_index).attr::<Transform>();
+		let transformed_point = transform.transform_point2(positions[index]);
+
+		let scoped = ctx.push_position(transformed_point);
+		let mut frame = IndexLink { index: 0, outer: None };
+		let (element, local) = content.eval(&scoped.ctx().push_level(&mut frame, index as u64, rest))?;
+		let mut composed = *local;
+		composed.translation = transformed_point;
+		return Ok(emit(element, Attr(composed)));
+	}
+	Err(GraphError::new("repeat on points addressed past its point count").into())
+}
+
+/// The pushed level's extent is the flattened point count across the points
+/// rows; inner levels forward to the content, uniform across copies.
+fn repeat_on_points_extent(content: ExtentIn<'_>, points: ListIn<'_, Vector>, _reverse: ValueIn<'_, bool>, level: LevelIn) -> GPoll<Extent> {
+	match level.pushed() {
+		true => points
+			.get()
+			.map(|points| Extent::Exactly((0..points.len()).map(|row| points.element_ref(row).point_domain.positions().len()).sum())),
+		false => content.at(level),
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -92,6 +144,7 @@ mod tests {
 	use core_types::context::{ContextImpl, EvalScope};
 	use core_types::node::Node;
 	use core_types::record::{FieldWrite, Layout, Rec, RecordSource, RecordValue, element_write, stack};
+	use vector_types::subpath::Subpath;
 
 	struct ValueNode<T>(T);
 
@@ -128,6 +181,67 @@ mod tests {
 		EvalScope::new(Some(0.5), None, None, generations, arena)
 	}
 
+	struct VectorRows {
+		layout: Layout,
+		rows: Vec<(Vector, DAffine2)>,
+	}
+
+	impl<'e> Node<ContextImpl<'e>> for VectorRows {
+		type Output = RecordValue<'e>;
+
+		fn eval(&self, input: &ContextImpl<'e>) -> GPoll<RecordValue<'e>> {
+			use core_types::context::{ExtractArena, ExtractIndex};
+			let (vector, transform) = &self.rows[input.innermost_index() as usize % self.rows.len()];
+			let dst = stack::push(self.layout.frame_bytes());
+			// SAFETY: dst is the claimed frame of this layout; offsets are the layout's own.
+			unsafe {
+				if core_types::record::write_element(dst, vector.clone(), input.arena()).is_none() {
+					return GPoll::Error(Box::new(core_types::gpoll::GraphError::new("arena exhausted")));
+				}
+				core_types::record::write_field(dst, self.layout.offset_of(<Transform as AttributeMarker>::NAME, 0).unwrap(), *transform);
+			}
+			stack::pop(dst);
+			GPoll::Final(RecordValue::spilled(unsafe { Rec::new(dst.cast_const()) }))
+		}
+
+		fn extent_at(&self, _input: &ContextImpl<'e>, _level: u8) -> GPoll<Extent> {
+			GPoll::Final(Extent::Exactly(self.rows.len()))
+		}
+
+		fn layout(&self) -> &Layout {
+			&self.layout
+		}
+	}
+
+	fn vector_rows_layout() -> Layout {
+		Layout::default().with_writes(1, element_write::<Vector>(), &[FieldWrite::of::<Transform>(0)])
+	}
+
+	struct PositionProbe {
+		layout: Layout,
+	}
+
+	impl<'e> Node<ContextImpl<'e>> for PositionProbe {
+		type Output = RecordValue<'e>;
+
+		fn eval(&self, input: &ContextImpl<'e>) -> GPoll<RecordValue<'e>> {
+			use core_types::context::ExtractPosition;
+			let position = input.try_position().and_then(|mut positions| positions.next()).unwrap_or(DVec2::ZERO);
+			let dst = stack::push(self.layout.frame_bytes());
+			// SAFETY: dst is the claimed frame of this layout; offsets are the layout's own.
+			unsafe {
+				dst.cast::<f64>().write(position.x);
+				core_types::record::write_field(dst, self.layout.offset_of(<Transform as AttributeMarker>::NAME, 0).unwrap(), DAffine2::IDENTITY);
+			}
+			stack::pop(dst);
+			GPoll::Final(RecordValue::spilled(unsafe { Rec::new(dst.cast_const()) }))
+		}
+
+		fn layout(&self) -> &Layout {
+			&self.layout
+		}
+	}
+
 	fn transform_layout() -> Layout {
 		Layout::default().with_writes(0, element_write::<f64>(), &[FieldWrite::of::<Transform>(0)])
 	}
@@ -214,4 +328,87 @@ mod tests {
 			unsafe { stack::rewind(mark) };
 		}
 	}
+	#[test]
+	fn repeat_on_points_lands_each_copy_on_its_transformed_point() {
+		let arena = Arena::new(1 << 16).unwrap();
+		let generations = [];
+		let scope = scope_fixture(&generations, &arena);
+		let ctx = ContextImpl::root(&scope);
+
+		let row0: Vec<DVec2> = vec![DVec2::new(40., 20.), DVec2::ONE];
+		let row1: Vec<DVec2> = vec![DVec2::new(-42., 9.), DVec2::new(10., 345.), DVec2::new(3., 4.)];
+		let row0_transform = DAffine2::from_translation(DVec2::new(100., 0.));
+		let points = VectorRows {
+			layout: vector_rows_layout(),
+			rows: vec![
+				(Vector::from_subpath(Subpath::from_anchors(row0.clone(), false)), row0_transform),
+				(Vector::from_subpath(Subpath::from_anchors(row1.clone(), false)), DAffine2::IDENTITY),
+			],
+		};
+		let content_layout = transform_layout();
+		let content = PositionProbe { layout: content_layout.clone() };
+
+		let mut node = RepeatOnPointsNode::new(RecordSource::new(content, &content_layout, &content_layout), points, ValueNode(false), &content_layout);
+		Node::<ContextImpl>::set_layout(&mut node, repeat_on_points_layout_meta().resolve(&[Some(&content_layout)]));
+		let leveled = Node::<ContextImpl>::layout(&node).clone();
+		assert_eq!(leveled.depth, 1);
+		assert_eq!(node.extent_at(&ctx, 0), GPoll::Final(Extent::Exactly(5)), "the pushed level flattens both rows' points");
+
+		let expected: Vec<DVec2> = row0
+			.iter()
+			.map(|&point| row0_transform.transform_point2(point))
+			.chain(row1.iter().copied())
+			.collect();
+
+		let head = ctx.index_head();
+		for (flat, &point) in expected.iter().enumerate() {
+			let mark = stack::sp();
+			let lane = ctx.promoted(&head, flat as u64);
+			let GPoll::Final(value) = node.eval(&lane) else {
+				panic!("expected a final record");
+			};
+			let rec = leveled.rec(&value);
+			// The content saw the pushed position, and the output transform lands on it.
+			assert_eq!(unsafe { rec.element::<f64>() }, point.x);
+			let composed: DAffine2 = unsafe { rec.read(leveled.offset_of(<Transform as AttributeMarker>::NAME, 0).unwrap()) };
+			assert_eq!(composed.translation, point);
+			// SAFETY: the element and transform were read out above, so no borrow into this lane's frames remains.
+			unsafe { stack::rewind(mark) };
+		}
+	}
+
+	#[test]
+	fn repeat_on_points_reverse_flips_each_rows_points() {
+		let arena = Arena::new(1 << 16).unwrap();
+		let generations = [];
+		let scope = scope_fixture(&generations, &arena);
+		let ctx = ContextImpl::root(&scope);
+
+		let positions: Vec<DVec2> = vec![DVec2::new(40., 20.), DVec2::ONE, DVec2::new(-42., 9.), DVec2::new(10., 345.)];
+		let points = VectorRows {
+			layout: vector_rows_layout(),
+			rows: vec![(Vector::from_subpath(Subpath::from_anchors(positions.clone(), false)), DAffine2::IDENTITY)],
+		};
+		let content_layout = transform_layout();
+		let content = PositionProbe { layout: content_layout.clone() };
+
+		let mut node = RepeatOnPointsNode::new(RecordSource::new(content, &content_layout, &content_layout), points, ValueNode(true), &content_layout);
+		Node::<ContextImpl>::set_layout(&mut node, repeat_on_points_layout_meta().resolve(&[Some(&content_layout)]));
+		let leveled = Node::<ContextImpl>::layout(&node).clone();
+
+		let mut expected = positions.clone();
+		expected.reverse();
+		let head = ctx.index_head();
+		for (flat, &point) in expected.iter().enumerate() {
+			let mark = stack::sp();
+			let lane = ctx.promoted(&head, flat as u64);
+			let GPoll::Final(value) = node.eval(&lane) else {
+				panic!("expected a final record");
+			};
+			let composed: DAffine2 = unsafe { leveled.rec(&value).read(leveled.offset_of(<Transform as AttributeMarker>::NAME, 0).unwrap()) };
+			assert_eq!(composed.translation, point);
+			// SAFETY: the transform was read out above, so no borrow into this lane's frames remains.
+			unsafe { stack::rewind(mark) };
+		}
+	}
 }