Graphite/node-graph/nodes/raster/src/blending_nodes.rs

use crate::adjust::Adjust;
#[cfg(feature = "std")]
use core_types::list::Item;
use no_std_types::Ctx;
use no_std_types::blending::BlendMode;
use no_std_types::color::{Color, Pixel};
#[cfg(not(feature = "std"))]
use no_std_types::list::ShaderItem as Item;
use no_std_types::registry::types::PercentageF32;
#[cfg(feature = "std")]
use raster_types::{CPU, Raster};
#[cfg(feature = "std")]
use vector_types::{Gradient, GradientStop};

pub trait Blend<P: Pixel> {
	fn blend(&self, under: &Self, blend_fn: impl Fn(P, P) -> P) -> Self;
}
impl Blend<Color> for Color {
	fn blend(&self, under: &Self, blend_fn: impl Fn(Color, Color) -> Color) -> Self {
		blend_fn(*self, *under)
	}
}

#[cfg(feature = "std")]
mod blend_std {
	use super::*;
	use core::cmp::Ordering;
	use raster_types::Image;
	use raster_types::Raster;

	impl Blend<Color> for Raster<CPU> {
		fn blend(&self, under: &Self, blend_fn: impl Fn(Color, Color) -> Color) -> Self {
			let data = self.data.iter().zip(under.data.iter()).map(|(a, b)| blend_fn(*a, *b)).collect();

			Raster::new_cpu(Image {
				data,
				width: self.width,
				height: self.height,
				base64_string: None,
			})
		}
	}
	impl Blend<Color> for Gradient {
		// TODO: This joining is unfaithful in several ways: it samples only at stop positions so midpoint curves flatten away;
		// TODO: it evaluates both sources with default whole-ramp attributes rather than their own (which this element-level impl cannot read);
		// TODO: and the output keeps over's attributes despite being sampled with defaults
		fn blend(&self, under: &Self, blend_fn: impl Fn(Color, Color) -> Color) -> Self {
			let mut combined_stops = self.positions(false).into_iter().chain(under.positions(false)).collect::<Vec<_>>();
			combined_stops.sort_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal));
			combined_stops.dedup_by(|a, b| (*a - *b).abs() < 1e-6);
			let over_evaluator = self.evaluator(Default::default());
			let under_evaluator = under.evaluator(Default::default());
			let stops = combined_stops.into_iter().map(|position| {
				let color = blend_fn(over_evaluator.evaluate(position), under_evaluator.evaluate(position));
				GradientStop { position, midpoint: 0.5, color }
			});

			// Positions stay explicit because eliding them needs the cyclic flag this impl can't read, and a wrong guess would relocate the stops
			Gradient::new(stops)
		}
	}
}

pub use no_std_types::blending::{apply_blend_mode, blend_colors};

#[node_macro::node(category("Raster"), cfg(feature = "std"))]
fn mix<T: Blend<Color> + Send>(
	_: impl Ctx,
	#[implementations(
		Raster<CPU>,
		Color,
		Gradient,
	)]
	#[gpu_image]
	over: Item<T>,
	#[expose]
	#[implementations(
		Raster<CPU>,
		Color,
		Gradient,
	)]
	#[gpu_image]
	under: Item<T>,
	blend_mode: Item<BlendMode>,
	#[default(100.)] opacity: Item<PercentageF32>,
) -> Item<T> {
	let mut over = over;
	let blend_mode = blend_mode.into_element();
	let opacity = opacity.into_element();

	let blended = over.element().blend(under.element(), |a, b| blend_colors(a, b, blend_mode, opacity / 100.));
	*over.element_mut() = blended;
	over
}

#[node_macro::node(category("Raster: Adjustment"), shader_node(PerPixelAdjust))]
fn color_overlay<T: Adjust<Color>>(
	_: impl Ctx,
	#[implementations(
		Raster<CPU>,
		Color,
		Gradient,
	)]
	#[gpu_image]
	image: Item<T>,
	#[default(Color::BLACK)] color: Item<Color>,
	blend_mode: Item<BlendMode>,
	#[default(100.)] opacity: Item<PercentageF32>,
) -> Item<T> {
	let mut image = image;
	let color = color.into_element();
	let blend_mode = blend_mode.into_element();
	let opacity = opacity.into_element();

	let opacity = (opacity / 100.).clamp(0., 1.);

	image.element_mut().adjust(|pixel| {
		let image = pixel.map_rgb(|channel| channel * (1. - opacity));

		// The apply blend mode function divides rgb by the alpha channel for the background. This undoes that.
		let associated_pixel = Color::from_rgbaf32_unchecked(pixel.r() * pixel.a(), pixel.g() * pixel.a(), pixel.b() * pixel.a(), pixel.a());
		let overlay = apply_blend_mode(color, associated_pixel, blend_mode).map_rgb(|channel| channel * opacity);

		Color::from_rgbaf32_unchecked(image.r() + overlay.r(), image.g() + overlay.g(), image.b() + overlay.b(), pixel.a())
	});
	image
}

#[cfg(all(feature = "std", test))]
mod test {
	use core_types::blending::BlendMode;
	use core_types::color::Color;
	use core_types::list::Item;
	use raster_types::Image;
	use raster_types::Raster;

	#[tokio::test]
	async fn color_overlay_multiply() {
		let image_color = Color::from_rgbaf32_unchecked(0.7, 0.6, 0.5, 0.4);
		let image = Image::new(1, 1, image_color);

		// Color { red: 0., green: 1., blue: 0., alpha: 1. }
		let overlay_color = Color::GREEN;

		// 100% of the output should come from the multiplied value
		let opacity = 100.;

		let result = super::color_overlay(
			(),
			Item::new_from_element(Raster::new_cpu(image.clone())),
			overlay_color.into(),
			BlendMode::Multiply.into(),
			opacity.into(),
		);
		let result = result.into_element();

		// The output should just be the original green and alpha channels (as we multiply them by 1 and other channels by 0)
		assert_eq!(result.data[0], Color::from_rgbaf32_unchecked(0., image_color.g(), 0., image_color.a()));
	}
}