Editor/Prism/Nodes/ChannelNodes.cs

Editor node library for Prism channel operations. Defines utilities for vector component handling and several PrismNode implementations: Split, Combine, Swizzle, Append, Channel Mask, Channel Flip, and Component Select, each validating inputs and emitting IR using swizzles, constructs, and masks.

Native Interop
using Editor.Prism.Compiler;
using Editor.Prism.Compiler.Ir;
using Editor.Prism.Core;
using Editor.Prism.Model;

namespace Editor.Prism.Nodes;

// ---------------------------------------------------------------------------------------------------
// WP-6 · node library A — taking vectors apart and putting them back together.
//
// These are the nodes that most often sit between two others of different widths, so every one of them
// degrades rather than fails: asking a float2 for its .w gives zero and a note, not a broken compile,
// and an over-long swizzle is clamped to what the source actually has.
// ---------------------------------------------------------------------------------------------------

/// <summary>One component of a vector.</summary>
public enum PrismVectorChannel
{
	/// <summary>The first component — <c>x</c> or <c>r</c>.</summary>
	X,
	/// <summary>The second component — <c>y</c> or <c>g</c>.</summary>
	Y,
	/// <summary>The third component — <c>z</c> or <c>b</c>.</summary>
	Z,
	/// <summary>The fourth component — <c>w</c> or <c>a</c>.</summary>
	W
}

/// <summary>Component bookkeeping shared by the channel nodes.</summary>
internal static class PrismChannelHelpers
{
	/// <summary>The canonical swizzle alphabet, in component order.</summary>
	public const string Alphabet = "xyzw";

	/// <summary>How many components a value has, treating anything that is not a vector as one.</summary>
	public static int Width( IrValue value ) =>
		value.IsValid && value.Type.IsScalarOrVector ? Math.Max( 1, value.Type.Components ) : 0;

	/// <summary>The component index a swizzle character selects, or -1 when it is not a swizzle character.</summary>
	public static int IndexOf( char c ) => c switch
	{
		'x' or 'X' or 'r' or 'R' => 0,
		'y' or 'Y' or 'g' or 'G' => 1,
		'z' or 'Z' or 'b' or 'B' => 2,
		'w' or 'W' or 'a' or 'A' => 3,
		_ => -1
	};

	/// <summary>
	/// Rewrite a user-authored mask into the canonical alphabet, clamping every component that the
	/// source does not have and dropping anything past four. Never returns an empty mask.
	/// </summary>
	public static string Sanitize( string mask, int width, out bool clamped )
	{
		clamped = false;
		width = Math.Clamp( width, 1, 4 );

		if ( string.IsNullOrWhiteSpace( mask ) ) return "x";

		var trimmed = mask.Trim().TrimStart( '.' );
		var chars = new List<char>( 4 );

		foreach ( var c in trimmed )
		{
			if ( chars.Count == 4 )
			{
				clamped = true;
				break;
			}

			var index = IndexOf( c );

			if ( index < 0 )
			{
				clamped = true;
				continue;
			}

			if ( index >= width )
			{
				clamped = true;
				index = width - 1;
			}

			chars.Add( Alphabet[index] );
		}

		if ( chars.Count == 0 ) return "x";

		return new string( chars.ToArray() );
	}

	/// <summary>True when every character of a mask is a legal swizzle letter and the mask is 1 to 4 long.</summary>
	public static bool IsWellFormed( string mask )
	{
		if ( string.IsNullOrWhiteSpace( mask ) ) return false;

		var trimmed = mask.Trim().TrimStart( '.' );

		if ( trimmed.Length is < 1 or > 4 ) return false;

		foreach ( var c in trimmed )
		{
			if ( IndexOf( c ) < 0 ) return false;
		}

		return true;
	}

	/// <summary>One component of a value, or zero when the value is too narrow to have it.</summary>
	public static IrValue Component( EmitContext ctx, IrValue value, int index )
	{
		var width = Width( value );

		if ( width <= 0 || index < 0 || index >= width ) return ctx.Const( 0f );

		return ctx.Swizzle( value, Alphabet[index].ToString() );
	}
}

/// <summary>Take a vector apart into its four components.</summary>
[NodeInfo( Id = "prism.channel.split", Title = "Split", Category = "Channel", Icon = "call_split",
	Keywords = ["split", "components", "xyzw", "rgba", "break", "decompose"] )]
[NodeVersion( 1 )]
public sealed class SplitVectorNode : PrismNode
{
	/// <summary>The vector to take apart.</summary>
	[In( "T", Name = "In" )] public PortRef In { get; set; }

	/// <summary>The first component. Zero when the input has none.</summary>
	[Out( "T.scalar", Name = "X" )] public PortRef X { get; set; }

	/// <summary>The second component. Zero when the input has none.</summary>
	[Out( "T.scalar", Name = "Y" )] public PortRef Y { get; set; }

	/// <summary>The third component. Zero when the input has none.</summary>
	[Out( "T.scalar", Name = "Z" )] public PortRef Z { get; set; }

	/// <summary>The fourth component. Zero when the input has none.</summary>
	[Out( "T.scalar", Name = "W" )] public PortRef W { get; set; }

	/// <summary>The literal used when <c>In</c> is unconnected.</summary>
	[InlineValue( nameof( In ) ), Title( "In" )] public Vector4 DefaultIn { get; set; } = new( 0f, 0f, 0f, 1f );

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var value = ctx.In( nameof( In ) );

		if ( !value.IsValid ) return;

		if ( !value.Type.IsScalarOrVector )
		{
			ctx.Error( $"'{value.Type.Hlsl}' has no components to split", nameof( In ) );
			return;
		}

		value = ctx.Let( "split", value );

		ctx.Out( nameof( X ), PrismChannelHelpers.Component( ctx, value, 0 ) );
		ctx.Out( nameof( Y ), PrismChannelHelpers.Component( ctx, value, 1 ) );
		ctx.Out( nameof( Z ), PrismChannelHelpers.Component( ctx, value, 2 ) );
		ctx.Out( nameof( W ), PrismChannelHelpers.Component( ctx, value, 3 ) );
	}
}

/// <summary>Build a vector out of separate components.</summary>
[NodeInfo( Id = "prism.channel.combine", Title = "Combine", Category = "Channel", Icon = "call_merge",
	Keywords = ["combine", "merge", "make", "xyzw", "rgba", "construct"] )]
[NodeVersion( 1 )]
public sealed class CombineVectorNode : PrismNode
{
	/// <summary>The first component.</summary>
	[In( "float", Name = "X" )] public PortRef X { get; set; }

	/// <summary>The second component.</summary>
	[In( "float", Name = "Y" )] public PortRef Y { get; set; }

	/// <summary>The third component.</summary>
	[In( "float", Name = "Z" )] public PortRef Z { get; set; }

	/// <summary>The fourth component.</summary>
	[In( "float", Name = "W" )] public PortRef W { get; set; }

	/// <summary>All four components.</summary>
	[Out( "float4", Name = "XYZW" )] public PortRef XYZW { get; set; }

	/// <summary>The first three components.</summary>
	[Out( "float3", Name = "XYZ" )] public PortRef XYZ { get; set; }

	/// <summary>The first two components.</summary>
	[Out( "float2", Name = "XY" )] public PortRef XY { get; set; }

	/// <summary>The literal used when <c>X</c> is unconnected.</summary>
	[InlineValue( nameof( X ) ), Title( "X" )] public float DefaultX { get; set; }

	/// <summary>The literal used when <c>Y</c> is unconnected.</summary>
	[InlineValue( nameof( Y ) ), Title( "Y" )] public float DefaultY { get; set; }

	/// <summary>The literal used when <c>Z</c> is unconnected.</summary>
	[InlineValue( nameof( Z ) ), Title( "Z" )] public float DefaultZ { get; set; }

	/// <summary>The literal used when <c>W</c> is unconnected.</summary>
	[InlineValue( nameof( W ) ), Title( "W" )] public float DefaultW { get; set; } = 1f;

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var x = ctx.InAs( nameof( X ), ShaderType.Float );
		var y = ctx.InAs( nameof( Y ), ShaderType.Float );
		var z = ctx.InAs( nameof( Z ), ShaderType.Float );
		var w = ctx.InAs( nameof( W ), ShaderType.Float );

		if ( !x.IsValid || !y.IsValid || !z.IsValid || !w.IsValid ) return;

		ctx.Out( nameof( XY ), ctx.Construct( ShaderType.Float2, x, y ) );
		ctx.Out( nameof( XYZ ), ctx.Construct( ShaderType.Float3, x, y, z ) );
		ctx.Out( nameof( XYZW ), ctx.Construct( ShaderType.Float4, x, y, z, w ) );
	}
}

/// <summary>Reorder, repeat or drop a vector's components with a free-form mask.</summary>
[NodeInfo( Id = "prism.channel.swizzle", Title = "Swizzle", Category = "Channel", Icon = "swap_horiz",
	Keywords = ["swizzle", "reorder", "shuffle", "mask", "xyzw", "rgba"] )]
[NodeVersion( 1 )]
public sealed class SwizzleVectorNode : PrismNode
{
	/// <summary>The vector to reorder.</summary>
	[In( "T", Name = "In" )] public PortRef In { get; set; }

	/// <summary>The reordered vector. Its width is the length of the mask.</summary>
	[Out( "any", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>In</c> is unconnected.</summary>
	[InlineValue( nameof( In ) ), Title( "In" )] public Vector4 DefaultIn { get; set; } = new( 0f, 0f, 0f, 1f );

	/// <summary>
	/// Which components to take and in what order, one to four of <c>xyzw</c> or <c>rgba</c>. Components
	/// the source does not have are clamped to its last one rather than failing the compile.
	/// </summary>
	public string Mask { get; set; } = "xyzw";

	/// <inheritdoc/>
	public override void OnValidate( ValidationContext ctx )
	{
		if ( ctx is null ) return;

		if ( !PrismChannelHelpers.IsWellFormed( Mask ) )
		{
			ctx.Error( $"'{Mask}' is not a swizzle; use one to four of x, y, z, w or r, g, b, a" );
			return;
		}

		var type = ctx.TypeOf( nameof( In ) );

		if ( !type.IsScalarOrVector || type.IsVoid ) return;

		var width = Math.Max( 1, type.Components );

		foreach ( var c in Mask.Trim().TrimStart( '.' ) )
		{
			if ( PrismChannelHelpers.IndexOf( c ) < width ) continue;

			ctx.Warn( $"'{type.Hlsl}' has no '{c}' component; it will be clamped to the last one", nameof( In ) );
			return;
		}
	}

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var value = ctx.In( nameof( In ) );

		if ( !value.IsValid ) return;

		if ( !value.Type.IsScalarOrVector )
		{
			ctx.Error( $"'{value.Type.Hlsl}' cannot be swizzled", nameof( In ) );
			return;
		}

		var mask = PrismChannelHelpers.Sanitize( Mask, PrismChannelHelpers.Width( value ), out var clamped );

		if ( clamped )
		{
			ctx.Info( $"The mask '{Mask}' does not fit a {value.Type.Hlsl}; it was read as '.{mask}'", nameof( In ) );
		}

		ctx.Out( nameof( Out ), ctx.Swizzle( value, mask ) );
	}
}

/// <summary>Join two values end to end into one wider vector.</summary>
[NodeInfo( Id = "prism.channel.append", Title = "Append", Category = "Channel", Icon = "add_box",
	Keywords = ["append", "concat", "join", "widen", "pack"] )]
[NodeVersion( 1 )]
public sealed class AppendVectorNode : PrismNode
{
	/// <summary>The leading components.</summary>
	[In( "T", Name = "A" )] public PortRef A { get; set; }

	/// <summary>The trailing components.</summary>
	[In( "any", Name = "B" )] public PortRef B { get; set; }

	/// <summary>The two inputs joined together.</summary>
	[Out( "any", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>A</c> is unconnected.</summary>
	[InlineValue( nameof( A ) ), Title( "A" )] public float DefaultA { get; set; }

	/// <summary>The literal used when <c>B</c> is unconnected.</summary>
	[InlineValue( nameof( B ) ), Title( "B" )] public float DefaultB { get; set; }

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var a = ctx.In( nameof( A ) );
		var b = ctx.In( nameof( B ) );

		if ( !a.IsValid || !b.IsValid ) return;

		var left = PrismChannelHelpers.Width( a );
		var right = PrismChannelHelpers.Width( b );

		if ( left <= 0 || right <= 0 )
		{
			ctx.Error( $"'{a.Type.Hlsl}' and '{b.Type.Hlsl}' cannot be appended" );
			return;
		}

		var total = left + right;

		if ( total > 4 )
		{
			ctx.Warn( $"A {a.Type.Hlsl} and a {b.Type.Hlsl} make {total} components; the result was cut to four" );
			total = 4;
		}

		var scalar = TypeRules.PromoteScalar( a.Type.Scalar, b.Type.Scalar );

		ctx.Out( nameof( Out ), ctx.Construct( ShaderType.Vec( scalar, total ), a, b ) );
	}
}

/// <summary>Zero out whichever components you do not want.</summary>
[NodeInfo( Id = "prism.channel.mask", Title = "Channel Mask", Category = "Channel", Icon = "filter_alt",
	Keywords = ["mask", "channel", "zero", "isolate", "rgba"] )]
[NodeVersion( 1 )]
public sealed class ChannelMaskNode : PrismNode
{
	/// <summary>The value to mask.</summary>
	[In( "T", Name = "In" )] public PortRef In { get; set; }

	/// <summary>The masked value: kept components pass through, the rest are zero.</summary>
	[Out( "T", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>In</c> is unconnected.</summary>
	[InlineValue( nameof( In ) ), Title( "In" )] public Vector4 DefaultIn { get; set; } = new( 1f, 1f, 1f, 1f );

	/// <summary>Keep the first component.</summary>
	[Title( "R" )] public bool Red { get; set; } = true;

	/// <summary>Keep the second component.</summary>
	[Title( "G" )] public bool Green { get; set; } = true;

	/// <summary>Keep the third component.</summary>
	[Title( "B" )] public bool Blue { get; set; } = true;

	/// <summary>Keep the fourth component.</summary>
	[Title( "A" )] public bool Alpha { get; set; } = true;

	/// <summary>The mask as a numeric constant, one per component.</summary>
	internal ConstValue Weights => new( Red ? 1d : 0d, Green ? 1d : 0d, Blue ? 1d : 0d, Alpha ? 1d : 0d );

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var value = ctx.In( nameof( In ) );

		if ( !value.IsValid ) return;

		var width = PrismChannelHelpers.Width( value );

		if ( width <= 0 )
		{
			ctx.Error( $"'{value.Type.Hlsl}' has no channels to mask", nameof( In ) );
			return;
		}

		var mask = ctx.Const( ShaderType.Vec( ScalarKind.Float, width ), Weights );

		ctx.Out( nameof( Out ), ctx.Bin( BinaryOp.Mul, value, mask ) );
	}
}

/// <summary>Invert whichever components you choose, leaving the rest alone.</summary>
[NodeInfo( Id = "prism.channel.flip", Title = "Flip", Category = "Channel", Icon = "flip",
	Keywords = ["flip", "invert", "one minus", "channel", "negate"] )]
[NodeVersion( 1 )]
public sealed class ChannelFlipNode : PrismNode
{
	/// <summary>The value to flip.</summary>
	[In( "T", Name = "In" )] public PortRef In { get; set; }

	/// <summary>The flipped value.</summary>
	[Out( "T", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>In</c> is unconnected.</summary>
	[InlineValue( nameof( In ) ), Title( "In" )] public Vector4 DefaultIn { get; set; } = new( 0f, 0f, 0f, 1f );

	/// <summary>Replace the first component with one minus itself.</summary>
	[Title( "R" )] public bool Red { get; set; }

	/// <summary>Replace the second component with one minus itself.</summary>
	[Title( "G" )] public bool Green { get; set; }

	/// <summary>Replace the third component with one minus itself.</summary>
	[Title( "B" )] public bool Blue { get; set; }

	/// <summary>Replace the fourth component with one minus itself.</summary>
	[Title( "A" )] public bool Alpha { get; set; }

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var value = ctx.In( nameof( In ) );

		if ( !value.IsValid ) return;

		var width = PrismChannelHelpers.Width( value );

		if ( width <= 0 )
		{
			ctx.Error( $"'{value.Type.Hlsl}' has no channels to flip", nameof( In ) );
			return;
		}

		if ( !Red && !Green && !Blue && !Alpha )
		{
			ctx.Out( nameof( Out ), value );
			return;
		}

		value = ctx.Let( "flip", value );

		var weights = new ConstValue( Red ? 1d : 0d, Green ? 1d : 0d, Blue ? 1d : 0d, Alpha ? 1d : 0d );
		var mask = ctx.Const( ShaderType.Vec( ScalarKind.Float, width ), weights );
		var inverted = ctx.Bin( BinaryOp.Sub, ctx.Const( 1f ), value );

		// lerp with a 0/1 mask picks per component and stays exact at both ends.
		ctx.Out( nameof( Out ), ctx.Call( Intrinsic.Lerp, value, inverted, mask ) );
	}
}

// The channel mixer lives with the other colour-adjustment nodes, as prism.color.channelMixer,
// which is where the reference catalogue puts it. Nothing is declared for it here.

/// <summary>Pull a single named component out of a vector.</summary>
[NodeInfo( Id = "prism.channel.select", Title = "Component Select", Category = "Channel",
	Icon = "filter_1", Keywords = ["component", "select", "pick", "channel", "extract", "index"] )]
[NodeVersion( 1 )]
public sealed class ComponentSelectNode : PrismNode
{
	/// <summary>The vector to read from.</summary>
	[In( "T", Name = "In" )] public PortRef In { get; set; }

	/// <summary>The chosen component.</summary>
	[Out( "T.scalar", Name = "Out" )] public PortRef Out { get; set; }

	/// <summary>The literal used when <c>In</c> is unconnected.</summary>
	[InlineValue( nameof( In ) ), Title( "In" )] public Vector4 DefaultIn { get; set; } = new( 0f, 0f, 0f, 1f );

	/// <summary>Which component to take.</summary>
	public PrismVectorChannel Channel { get; set; } = PrismVectorChannel.X;

	/// <inheritdoc/>
	public override void OnValidate( ValidationContext ctx )
	{
		if ( ctx is null ) return;

		var type = ctx.TypeOf( nameof( In ) );

		if ( !type.IsScalarOrVector || type.IsVoid ) return;
		if ( (int)Channel < Math.Max( 1, type.Components ) ) return;

		ctx.Warn( $"'{type.Hlsl}' has no {Channel} component; the result will be zero", nameof( In ) );
	}

	/// <inheritdoc/>
	public override void Emit( EmitContext ctx )
	{
		if ( ctx is null ) return;

		var value = ctx.In( nameof( In ) );

		if ( !value.IsValid ) return;

		if ( !value.Type.IsScalarOrVector )
		{
			ctx.Error( $"'{value.Type.Hlsl}' has no components to select", nameof( In ) );
			return;
		}

		ctx.Out( nameof( Out ), PrismChannelHelpers.Component( ctx, value, (int)Channel ) );
	}
}