Clean up duplicated code used for recursively flattening graphic types (#3836)

* Reduce recusive flattening algorithm duplication

* Generalize further

* Avoid code duplication in the 'Flatten Path' node

* Avoid cloning

* Include intermediate levels of alpha blending composition

Author: Keavon

node-graph/libraries/graphic-types/src/graphic.rs

@@ -104,201 +104,94 @@ impl From<Table<GradientStops>> for Graphic {
 	}
 }
 
-// Local trait to convert types to Table<Graphic> (avoids orphan rule issues)
-pub trait IntoGraphicTable {
-	fn into_graphic_table(self) -> Table<Graphic>;
+/// Deeply flattens a graphic table, collecting only elements matching a specific variant (extracted by `extract_variant`)
+/// and discarding all other non-matching content. Recursion through `Graphic::Graphic` sub-tables composes transforms and opacity.
+fn flatten_graphic_table<T>(content: Table<Graphic>, extract_variant: fn(Graphic) -> Option<Table<T>>) -> Table<T> {
+	fn compose_alpha_blending(parent: AlphaBlending, child: AlphaBlending) -> AlphaBlending {
+		AlphaBlending {
+			blend_mode: child.blend_mode,
+			opacity: parent.opacity * child.opacity,
+			fill: child.fill,
+			clip: child.clip,
+		}
+	}
 
-	/// Deeply flattens any vector content within a graphic table, discarding non-vector content, and returning a table of only vector elements.
-	fn into_flattened_vector_table(self) -> Table<Vector>
-	where
-		Self: std::marker::Sized,
-	{
-		let content = self.into_graphic_table();
-
-		// TODO: Avoid mutable reference, instead return a new Table<Graphic>?
-		fn flatten_table(output_vector_table: &mut Table<Vector>, current_graphic_table: Table<Graphic>) {
-			for current_graphic_row in current_graphic_table.iter() {
-				let current_graphic = current_graphic_row.element.clone();
-				let source_node_id = *current_graphic_row.source_node_id;
-
-				match current_graphic {
-					// If we're allowed to recurse, flatten any tables we encounter
-					Graphic::Graphic(mut current_graphic_table) => {
-						// Apply the parent graphic's transform to all child elements
-						for graphic in current_graphic_table.iter_mut() {
-							*graphic.transform = *current_graphic_row.transform * *graphic.transform;
-						}
+	fn flatten_recursive<T>(output: &mut Table<T>, current_graphic_table: Table<Graphic>, extract_variant: fn(Graphic) -> Option<Table<T>>) {
+		for current_graphic_row in current_graphic_table.into_iter() {
+			let source_node_id = current_graphic_row.source_node_id;
 
-						flatten_table(output_vector_table, current_graphic_table);
+			match current_graphic_row.element {
+				// Recurse into nested graphic tables, composing the parent's transform onto each child
+				Graphic::Graphic(mut sub_table) => {
+					for graphic in sub_table.iter_mut() {
+						*graphic.transform = current_graphic_row.transform * *graphic.transform;
+						*graphic.alpha_blending = compose_alpha_blending(current_graphic_row.alpha_blending, *graphic.alpha_blending);
 					}
-					// Push any leaf Vector elements we encounter
-					Graphic::Vector(vector_table) => {
-						for current_vector_row in vector_table.iter() {
-							output_vector_table.push(TableRow {
-								element: current_vector_row.element.clone(),
-								transform: *current_graphic_row.transform * *current_vector_row.transform,
-								alpha_blending: AlphaBlending {
-									blend_mode: current_vector_row.alpha_blending.blend_mode,
-									opacity: current_graphic_row.alpha_blending.opacity * current_vector_row.alpha_blending.opacity,
-									fill: current_vector_row.alpha_blending.fill,
-									clip: current_vector_row.alpha_blending.clip,
-								},
+
+					flatten_recursive(output, sub_table, extract_variant);
+				}
+				// Try to extract the target variant; if it matches, push its rows with composed transform and opacity
+				other => {
+					if let Some(typed_table) = extract_variant(other) {
+						for row in typed_table.into_iter() {
+							output.push(TableRow {
+								element: row.element,
+								transform: current_graphic_row.transform * row.transform,
+								alpha_blending: compose_alpha_blending(current_graphic_row.alpha_blending, row.alpha_blending),
 								source_node_id,
 							});
 						}
 					}
-					_ => {}
 				}
 			}
 		}
-
-		let mut output = Table::new();
-		flatten_table(&mut output, content);
-		output
 	}
 
-	/// Deeply flattens any raster content within a graphic table, discarding non-raster content, and returning a table of only raster elements.
-	fn into_flattened_raster_table(self) -> Table<Raster<CPU>>
-	where
-		Self: std::marker::Sized,
-	{
-		let content = self.into_graphic_table();
-
-		fn flatten_table(output_raster_table: &mut Table<Raster<CPU>>, current_graphic_table: Table<Graphic>) {
-			for current_graphic_row in current_graphic_table.iter() {
-				let current_graphic = current_graphic_row.element.clone();
-				let source_node_id = *current_graphic_row.source_node_id;
-
-				match current_graphic {
-					// If we're allowed to recurse, flatten any tables we encounter
-					Graphic::Graphic(mut current_graphic_table) => {
-						// Apply the parent graphic's transform to all child elements
-						for graphic in current_graphic_table.iter_mut() {
-							*graphic.transform = *current_graphic_row.transform * *graphic.transform;
-						}
+	let mut output = Table::new();
+	flatten_recursive(&mut output, content, extract_variant);
+	output
+}
 
-						flatten_table(output_raster_table, current_graphic_table);
-					}
-					// Push any leaf RasterCPU elements we encounter
-					Graphic::RasterCPU(raster_table) => {
-						for current_raster_row in raster_table.iter() {
-							output_raster_table.push(TableRow {
-								element: current_raster_row.element.clone(),
-								transform: *current_graphic_row.transform * *current_raster_row.transform,
-								alpha_blending: AlphaBlending {
-									blend_mode: current_raster_row.alpha_blending.blend_mode,
-									opacity: current_graphic_row.alpha_blending.opacity * current_raster_row.alpha_blending.opacity,
-									fill: current_raster_row.alpha_blending.fill,
-									clip: current_raster_row.alpha_blending.clip,
-								},
-								source_node_id,
-							});
-						}
-					}
-					_ => {}
-				}
-			}
-		}
+/// Maps from a concrete element type to its corresponding `Graphic` enum variant,
+/// enabling type-directed casting of typed tables from a `Graphic` value.
+pub trait TryFromGraphic: Clone + Sized {
+	fn try_from_graphic(graphic: Graphic) -> Option<Table<Self>>;
+}
 
-		let mut output = Table::new();
-		flatten_table(&mut output, content);
-		output
+impl TryFromGraphic for Vector {
+	fn try_from_graphic(graphic: Graphic) -> Option<Table<Self>> {
+		if let Graphic::Vector(t) = graphic { Some(t) } else { None }
 	}
+}
 
-	/// Deeply flattens any color content within a graphic table, discarding non-color content, and returning a table of only color elements.
-	fn into_flattened_color_table(self) -> Table<Color>
-	where
-		Self: std::marker::Sized,
-	{
-		let content = self.into_graphic_table();
-
-		fn flatten_table(output_color_table: &mut Table<Color>, current_graphic_table: Table<Graphic>) {
-			for current_graphic_row in current_graphic_table.iter() {
-				let current_graphic = current_graphic_row.element.clone();
-				let source_node_id = *current_graphic_row.source_node_id;
-
-				match current_graphic {
-					// If we're allowed to recurse, flatten any tables we encounter
-					Graphic::Graphic(mut current_graphic_table) => {
-						// Apply the parent graphic's transform to all child elements
-						for graphic in current_graphic_table.iter_mut() {
-							*graphic.transform = *current_graphic_row.transform * *graphic.transform;
-						}
+impl TryFromGraphic for Raster<CPU> {
+	fn try_from_graphic(graphic: Graphic) -> Option<Table<Self>> {
+		if let Graphic::RasterCPU(t) = graphic { Some(t) } else { None }
+	}
+}
 
-						flatten_table(output_color_table, current_graphic_table);
-					}
-					// Push any leaf Color elements we encounter
-					Graphic::Color(color_table) => {
-						for current_color_row in color_table.iter() {
-							output_color_table.push(TableRow {
-								element: *current_color_row.element,
-								transform: *current_graphic_row.transform * *current_color_row.transform,
-								alpha_blending: AlphaBlending {
-									blend_mode: current_color_row.alpha_blending.blend_mode,
-									opacity: current_graphic_row.alpha_blending.opacity * current_color_row.alpha_blending.opacity,
-									fill: current_color_row.alpha_blending.fill,
-									clip: current_color_row.alpha_blending.clip,
-								},
-								source_node_id,
-							});
-						}
-					}
-					_ => {}
-				}
-			}
-		}
+impl TryFromGraphic for Color {
+	fn try_from_graphic(graphic: Graphic) -> Option<Table<Self>> {
+		if let Graphic::Color(t) = graphic { Some(t) } else { None }
+	}
+}
 
-		let mut output = Table::new();
-		flatten_table(&mut output, content);
-		output
+impl TryFromGraphic for GradientStops {
+	fn try_from_graphic(graphic: Graphic) -> Option<Table<Self>> {
+		if let Graphic::Gradient(t) = graphic { Some(t) } else { None }
 	}
+}
 
-	/// Deeply flattens any gradient content within a graphic table, discarding non-gradient content, and returning a table of only gradient elements.
-	fn into_flattened_gradient_table(self) -> Table<GradientStops>
+// Local trait to convert types to Table<Graphic> (avoids orphan rule issues)
+pub trait IntoGraphicTable {
+	fn into_graphic_table(self) -> Table<Graphic>;
+
+	/// Deeply flattens any content of type `T` within a graphic table, discarding all other content, and returning a flat table of only `T` elements.
+	fn into_flattened_table<T: TryFromGraphic>(self) -> Table<T>
 	where
 		Self: std::marker::Sized,
 	{
-		let content = self.into_graphic_table();
-
-		fn flatten_table(output_gradient_table: &mut Table<GradientStops>, current_graphic_table: Table<Graphic>) {
-			for current_graphic_row in current_graphic_table.iter() {
-				let current_graphic = current_graphic_row.element.clone();
-				let source_node_id = *current_graphic_row.source_node_id;
-
-				match current_graphic {
-					// If we're allowed to recurse, flatten any tables we encounter
-					Graphic::Graphic(mut current_graphic_table) => {
-						// Apply the parent graphic's transform to all child elements
-						for graphic in current_graphic_table.iter_mut() {
-							*graphic.transform = *current_graphic_row.transform * *graphic.transform;
-						}
-
-						flatten_table(output_gradient_table, current_graphic_table);
-					}
-					// Push any leaf GradientStops elements we encounter
-					Graphic::Gradient(gradient_table) => {
-						for current_gradient_row in gradient_table.iter() {
-							output_gradient_table.push(TableRow {
-								element: current_gradient_row.element.clone(),
-								transform: *current_graphic_row.transform * *current_gradient_row.transform,
-								alpha_blending: AlphaBlending {
-									blend_mode: current_gradient_row.alpha_blending.blend_mode,
-									opacity: current_graphic_row.alpha_blending.opacity * current_gradient_row.alpha_blending.opacity,
-									fill: current_gradient_row.alpha_blending.fill,
-									clip: current_gradient_row.alpha_blending.clip,
-								},
-								source_node_id,
-							});
-						}
-					}
-					_ => {}
-				}
-			}
-		}
-
-		let mut output = Table::new();
-		flatten_table(&mut output, content);
-		output
+		flatten_graphic_table(self.into_graphic_table(), T::try_from_graphic)
 	}
 }
 

node-graph/libraries/graphic-types/src/lib.rs

@@ -8,7 +8,7 @@ pub use vector_types;
 
 // Re-export commonly used types at the crate root
 pub use artboard::Artboard;
-pub use graphic::{Graphic, IntoGraphicTable, Vector};
+pub use graphic::{Graphic, IntoGraphicTable, TryFromGraphic, Vector};
 
 pub mod migrations {
 	use core_types::{

node-graph/nodes/graphic/src/graphic.rs

@@ -295,31 +295,31 @@ pub async fn flatten_graphic(_: impl Ctx, content: Table<Graphic>, fully_flatten
 /// Converts a graphic table into a vector table by deeply flattening any vector content it contains, and discarding any non-vector content.
 #[node_macro::node(category("Vector"))]
 pub async fn flatten_vector<T: IntoGraphicTable + 'n + Send + Clone>(_: impl Ctx, #[implementations(Table<Graphic>, Table<Vector>)] content: T) -> Table<Vector> {
-	content.into_flattened_vector_table()
+	content.into_flattened_table()
 }
 
 /// Converts a graphic table into a raster table by deeply flattening any raster content it contains, and discarding any non-raster content.
 #[node_macro::node(category("Raster"))]
 pub async fn flatten_raster<T: IntoGraphicTable + 'n + Send + Clone>(_: impl Ctx, #[implementations(Table<Graphic>, Table<Raster<CPU>>)] content: T) -> Table<Raster<CPU>> {
-	content.into_flattened_raster_table()
+	content.into_flattened_table()
 }
 
 /// Converts a graphic table into a color table by deeply flattening any color content it contains, and discarding any non-color content.
 #[node_macro::node(category("General"))]
 pub async fn flatten_color<T: IntoGraphicTable + 'n + Send + Clone>(_: impl Ctx, #[implementations(Table<Graphic>, Table<Color>)] content: T) -> Table<Color> {
-	content.into_flattened_color_table()
+	content.into_flattened_table()
 }
 
 /// Converts a graphic table into a gradient table by deeply flattening any gradient content it contains, and discarding any non-gradient content.
 #[node_macro::node(category("General"))]
 pub async fn flatten_gradient<T: IntoGraphicTable + 'n + Send + Clone>(_: impl Ctx, #[implementations(Table<Graphic>, Table<GradientStops>)] content: T) -> Table<GradientStops> {
-	content.into_flattened_gradient_table()
+	content.into_flattened_table()
 }
 
 /// Constructs a gradient from a table of colors, where the colors are evenly distributed as gradient stops across the range from 0 to 1.
 #[node_macro::node(category("Color"))]
 fn colors_to_gradient<T: IntoGraphicTable + 'n + Send + Clone>(_: impl Ctx, #[implementations(Table<Graphic>, Table<Color>)] colors: T) -> GradientStops {
-	let colors = colors.into_flattened_color_table();
+	let colors = colors.into_flattened_table::<Color>();
 	let total_colors = colors.len();
 
 	if total_colors == 0 {