Editor/Prism/Compiler/Ir/IrPrinter.cs

IR (intermediate representation) pretty-printer for the Prism shader compiler editor. It renders IrModule, functions, blocks, statements and expressions into a readable pseudo-HLSL dump with options for indentation, metadata, declarations, origin node annotations, expression hashes and purity markers.

Reflection
using Editor.Prism.Core;
using System.Globalization;
using System.Text;

namespace Editor.Prism.Compiler.Ir;

/// <summary>How the IR dump is formatted.</summary>
public sealed record IrPrinterOptions
{
	/// <summary>One indent level.</summary>
	public string Indent { get; init; } = "\t";

	/// <summary>Line ending.</summary>
	public string NewLine { get; init; } = "\r\n";

	/// <summary>Print the module header: domain, blend mode, combos, capabilities.</summary>
	public bool ShowMetadata { get; init; } = true;

	/// <summary>Print declarations for structs, globals, helpers and varyings.</summary>
	public bool ShowDeclarations { get; init; } = true;

	/// <summary>Annotate every statement with the node that produced it. This is the source map, visible.</summary>
	public bool ShowOrigins { get; init; } = true;

	/// <summary>Annotate expressions with their structural hash. Only useful when debugging CSE.</summary>
	public bool ShowHashes { get; init; }

	/// <summary>Mark impure expressions with a bang, so a stray side effect is obvious.</summary>
	public bool ShowPurity { get; init; }

	/// <summary>The default: everything a user would want to read, nothing they would not.</summary>
	public static IrPrinterOptions Default { get; } = new();

	/// <summary>Everything, including hashes and purity. What a bug report should contain.</summary>
	public static IrPrinterOptions Verbose { get; } = new() { ShowHashes = true, ShowPurity = true };
}

/// <summary>
/// Renders an <see cref="IrModule"/> as readable pseudo-code. This is what the Code panel's "IR" tab
/// shows: the exact thing both backends are about to lower, with the originating node on every line.
/// It is a diagnostic view, not a backend — nothing here ever becomes a compiled artifact.
/// </summary>
public static class IrPrinter
{
	/// <summary>Print a whole module.</summary>
	public static string Print( IrModule module, IrPrinterOptions options = null )
	{
		var sb = new StringBuilder();
		Print( module, sb, options );
		return sb.ToString();
	}

	/// <summary>Print a whole module into an existing buffer.</summary>
	public static void Print( IrModule module, StringBuilder sb, IrPrinterOptions options = null )
	{
		if ( sb is null ) return;

		options ??= IrPrinterOptions.Default;

		if ( module is null )
		{
			sb.Append( "// <no module>" ).Append( options.NewLine );
			return;
		}

		var w = new Writer( sb, options );

		if ( options.ShowMetadata ) WriteMetadata( module, w );
		if ( options.ShowDeclarations ) WriteDeclarations( module, w );

		foreach ( var function in module.Functions )
		{
			WriteFunction( function, w );
		}
	}

	/// <summary>Print a single expression.</summary>
	public static string Print( IrExpr expr ) => Expr( expr, IrPrinterOptions.Default );

	/// <summary>Print a single expression with explicit options.</summary>
	public static string Print( IrExpr expr, IrPrinterOptions options ) => Expr( expr, options ?? IrPrinterOptions.Default );

	/// <summary>Print a single statement.</summary>
	public static string Print( IrStmt statement, IrPrinterOptions options = null )
	{
		var sb = new StringBuilder();
		var w = new Writer( sb, options ?? IrPrinterOptions.Default );

		WriteStatement( statement, w );

		return sb.ToString();
	}

	/// <summary>Print a block.</summary>
	public static string Print( IrBlock block, IrPrinterOptions options = null )
	{
		var sb = new StringBuilder();
		var w = new Writer( sb, options ?? IrPrinterOptions.Default );

		WriteBlock( block, w );

		return sb.ToString();
	}

	// ---- module sections --------------------------------------------------

	static void WriteMetadata( IrModule module, Writer w )
	{
		var meta = module.Meta;

		w.Line( $"// ---- module {meta.Name} ----" );
		w.Line( $"// domain      {meta.Domain}" );
		w.Line( $"// shading     {meta.ShadingModel}" );
		w.Line( $"// blend       {meta.BlendMode}, cull {meta.CullMode}{( meta.RenderBackfaces ? ", backfaces" : string.Empty )}" );
		w.Line( $"// stages      {meta.Stages}" );

		if ( meta.Modes.Count > 0 ) w.Line( $"// modes       {string.Join( ", ", meta.Modes )}" );
		if ( meta.UsesUv2 ) w.Line( "// uv2         yes" );

		if ( meta.Capabilities.Count > 0 )
		{
			var capabilities = meta.Capabilities.Select( x => x.ToString() ).OrderBy( x => x, StringComparer.Ordinal );
			w.Line( $"// capability  {string.Join( ", ", capabilities )}" );
		}

		foreach ( var combo in meta.Combos )
		{
			var values = combo.Values is null ? string.Empty : string.Join( ", ", combo.Values );
			w.Line( $"// combo       {combo.Kind} {combo.Name} [{values}] default {combo.Default}" );
		}

		w.Blank();
	}

	static void WriteDeclarations( IrModule module, Writer w )
	{
		if ( module.Includes.Count > 0 )
		{
			foreach ( var include in module.Includes ) w.Line( $"include \"{include}\"" );
			w.Blank();
		}

		foreach ( var global in module.Globals )
		{
			w.Line( DescribeGlobal( global ) );
		}

		if ( module.Globals.Count > 0 ) w.Blank();

		foreach ( var varying in module.Varyings )
		{
			var interpolation = varying.Interpolation == IrInterpolation.Linear
				? string.Empty
				: $" {varying.Interpolation.ToString().ToLowerInvariant()}";

			w.Line( $"varying{interpolation} {varying.Type.Hlsl} {varying.Name} : {varying.Semantic}   // slot {varying.Slot}" );
		}

		if ( module.Varyings.Count > 0 ) w.Blank();

		foreach ( var structure in module.Structs )
		{
			w.Line( $"struct {structure.Name}" );
			w.Line( "{" );
			w.Push();

			foreach ( var include in structure.Includes ) w.Line( $"include \"{include}\"" );

			foreach ( var field in structure.Fields )
			{
				var semantic = string.IsNullOrEmpty( field.Semantic ) ? string.Empty : $" : {field.Semantic}";
				w.Line( $"{field.Type.Hlsl} {field.Name}{semantic};" );
			}

			w.Pop();
			w.Line( "}" );
			w.Blank();
		}

		foreach ( var helper in module.Helpers )
		{
			w.Line( $"helper {helper.SignatureHlsl};" );
		}

		if ( module.Helpers.Count > 0 ) w.Blank();
	}

	static string DescribeGlobal( GlobalDecl global )
	{
		var sb = new StringBuilder();

		sb.Append( global.Kind.ToString().ToLowerInvariant() ).Append( ' ' );
		sb.Append( global.Type.Hlsl ).Append( ' ' ).Append( global.Name );

		if ( global.ArraySize > 0 ) sb.Append( '[' ).Append( global.ArraySize ).Append( ']' );

		if ( global.Default.HasValue ) sb.Append( " = " ).Append( global.Default.Value );
		if ( !string.IsNullOrEmpty( global.DefaultAsset ) ) sb.Append( " = \"" ).Append( global.DefaultAsset ).Append( '"' );

		var notes = new List<string>();

		if ( !string.IsNullOrEmpty( global.AttributeName ) ) notes.Add( $"attribute {global.AttributeName}" );
		if ( global.Parameter.IsValid ) notes.Add( $"param {global.Parameter}" );
		if ( global.Srgb ) notes.Add( "srgb" );
		if ( global.PreviewOnly ) notes.Add( "preview-only" );
		if ( !string.IsNullOrEmpty( global.SamplerName ) ) notes.Add( $"sampler {global.SamplerName}" );

		sb.Append( ';' );

		if ( notes.Count > 0 ) sb.Append( "   // " ).Append( string.Join( ", ", notes ) );

		return sb.ToString();
	}

	static void WriteFunction( IrFunction function, Writer w )
	{
		if ( function is null ) return;

		foreach ( var attribute in function.Attributes ) w.Line( attribute );

		var stage = function.Stage == ShaderStage.None ? string.Empty : $"   // {function.Stage.DisplayName()} stage";

		w.Line( function.SignatureHlsl + stage );
		WriteBlock( function.Body, w );
		w.Blank();
	}

	static void WriteBlock( IrBlock block, Writer w )
	{
		w.Line( "{" );
		w.Push();

		if ( block is not null )
		{
			foreach ( var statement in block.Statements ) WriteStatement( statement, w );
		}

		w.Pop();
		w.Line( "}" );
	}

	/// <summary>Write a block's statements at the current depth, without braces or a scope push.</summary>
	static void WriteStatements( IrBlock block, Writer w )
	{
		if ( block is null ) return;

		foreach ( var statement in block.Statements ) WriteStatement( statement, w );
	}

	static void WriteStatement( IrStmt statement, Writer w )
	{
		if ( statement is null ) return;

		switch ( statement )
		{
			case IrDecl decl:
				w.Line( $"{decl.Type.Hlsl} {decl.Name} = {Expr( decl.Init, w.Options )};", decl.Origin );
				break;

			case IrAssign assign:
				w.Line( $"{Expr( assign.Target, w.Options )} = {Expr( assign.Value, w.Options )};", assign.Origin );
				break;

			case IrIf branch:
				w.Line( $"if ( {Expr( branch.Cond, w.Options )} )", branch.Origin );
				WriteBlock( branch.Then, w );

				if ( branch.Else is not null && !branch.Else.IsEmpty )
				{
					w.Line( "else" );
					WriteBlock( branch.Else, w );
				}

				break;

			case IrFor loop:
				w.Line( $"for ( int {loop.Var} = 0; {loop.Var} < {Expr( loop.Count, w.Options )}; {loop.Var}++ )", loop.Origin );
				WriteBlock( loop.Body, w );
				break;

			case IrWhile loop:
				w.Line( $"while ( {Expr( loop.Cond, w.Options )} )", loop.Origin );
				WriteBlock( loop.Body, w );
				break;

			case IrBreak:
				w.Line( "break;", statement.Origin );
				break;

			case IrContinue:
				w.Line( "continue;", statement.Origin );
				break;

			case IrReturn ret:
				w.Line( ret.Value is null ? "return;" : $"return {Expr( ret.Value, w.Options )};", ret.Origin );
				break;

			case IrDiscard:
				w.Line( "discard;", statement.Origin );
				break;

			case IrExprStmt expr:
				w.Line( $"{Expr( expr.Value, w.Options )};", expr.Origin );
				break;

			case IrComment comment:
				foreach ( var line in SplitLines( comment.Text ) ) w.Line( $"// {line}" );
				break;

			case IrScope scope:
				WriteBlock( scope.Body, w );
				break;

			case IrPreprocessorIf guard:
				// Printed as the directives it actually becomes, with the guarded statements inline
				// rather than braced — a brace here would misrepresent the scope the backend emits.
				w.Line( IrPreprocessor.OpenDirective( guard.Condition ), guard.Origin );
				WriteStatements( guard.Then, w );

				if ( guard.HasElse )
				{
					w.Line( IrPreprocessor.ElseDirective );
					WriteStatements( guard.Else, w );
				}

				w.Line( IrPreprocessor.EndDirective );
				break;

			default:
				w.Line( $"// <unprintable {statement.GetType().Name}>", statement.Origin );
				break;
		}
	}

	static IEnumerable<string> SplitLines( string text )
	{
		if ( string.IsNullOrEmpty( text ) ) yield break;

		foreach ( var line in text.Replace( "\r\n", "\n" ).Split( '\n' ) ) yield return line;
	}

	// ---- expressions ------------------------------------------------------

	/// <summary>Render an expression, parenthesising only where precedence demands it.</summary>
	static string Expr( IrExpr expr, IrPrinterOptions options )
	{
		if ( expr is null ) return "<null>";

		var text = Render( expr, options );

		if ( options.ShowPurity && !expr.Pure ) text = "!" + text;
		if ( options.ShowHashes ) text += $" /*#{expr.Hash:x8}*/";

		return text;
	}

	static string Render( IrExpr expr, IrPrinterOptions options ) => expr switch
	{
		IrConst c => Literal( c ),
		IrVar v => v.Name,
		IrGlobalRef g => g.Decl?.Name ?? "<global>",
		IrBuiltinRef b => $"@{b.Id}",
		IrCall call => $"{IntrinsicCatalog.Name( call.Id )}( {Args( call.Args, options )} )",
		IrHelperCall helper => $"{helper.Fn?.Name ?? "<helper>"}( {Args( helper.Args, options )} )",
		IrBinary bin => Binary( bin, options ),
		IrUnary un => $"{UnaryOps.Symbol( un.Op )}{Operand( un.V, 100, options )}",
		IrSwizzle sw => $"{Operand( sw.V, 100, options )}.{sw.Mask}",
		IrConstruct ctor => $"{ctor.Type.Hlsl}( {Args( ctor.Parts, options )} )",
		IrCast cast => Cast( cast, options ),
		IrSelect sel => $"select( {Expr( sel.C, options )}, {Expr( sel.A, options )}, {Expr( sel.B, options )} )",
		IrIndex index => $"{Operand( index.V, 100, options )}[{Expr( index.I, options )}]",
		IrMember member => $"{Operand( member.V, 100, options )}.{member.Field}",
		_ => $"<{expr.GetType().Name}>"
	};

	static string Cast( IrCast cast, IrPrinterOptions options )
	{
		var inner = Expr( cast.V, options );

		return cast.Kind switch
		{
			CastKind.Splat => $"{cast.Type.Hlsl}( {inner} )",
			CastKind.Truncate => $"{Operand( cast.V, 100, options )}{TypeRules.SwizzleFor( cast.Type.Components )}",
			CastKind.Pad => $"{cast.Type.Hlsl}( {inner}, {Number( cast.Fill )} )",
			CastKind.Bitcast => $"asbits<{cast.Type.Hlsl}>( {inner} )",
			_ => $"({cast.Type.Hlsl}){Operand( cast.V, 100, options )}"
		};
	}

	static string Binary( IrBinary bin, IrPrinterOptions options )
	{
		var precedence = BinaryOps.Precedence( bin.Op );
		var left = Operand( bin.L, precedence, options );
		var right = Operand( bin.R, precedence + 1, options );

		return $"{left} {BinaryOps.Symbol( bin.Op )} {right}";
	}

	static string Operand( IrExpr expr, int precedence, IrPrinterOptions options )
	{
		var text = Expr( expr, options );

		if ( expr is IrBinary bin && BinaryOps.Precedence( bin.Op ) < precedence ) return $"( {text} )";
		if ( expr is IrUnary && precedence >= 100 ) return $"( {text} )";
		if ( expr is IrSelect && precedence >= 100 ) return text;

		return text;
	}

	static string Args( IrExpr[] args, IrPrinterOptions options )
	{
		if ( args is null || args.Length == 0 ) return string.Empty;

		return string.Join( ", ", args.Select( x => Expr( x, options ) ) );
	}

	static string Literal( IrConst c )
	{
		var type = c.Type;
		var components = Math.Clamp( type.Components, 1, 4 );

		if ( type.IsBoolean )
		{
			var parts = Enumerable.Range( 0, components ).Select( i => c.Value[i] != 0 ? "true" : "false" );
			return components == 1 ? parts.First() : $"{type.Hlsl}( {string.Join( ", ", parts )} )";
		}

		if ( type.IsIntegral )
		{
			var parts = Enumerable.Range( 0, components )
				.Select( i => ( (long)c.Value[i] ).ToString( CultureInfo.InvariantCulture ) );

			return components == 1 ? parts.First() : $"{type.Hlsl}( {string.Join( ", ", parts )} )";
		}

		var numbers = Enumerable.Range( 0, components ).Select( i => Number( (float)c.Value[i] ) );

		return components == 1 ? numbers.First() : $"{type.Hlsl}( {string.Join( ", ", numbers )} )";
	}

	/// <summary>
	/// Round-trippable float formatting. <c>"R"</c> would give the shortest round-trip form, but a
	/// trailing decimal point keeps the value unambiguously floating point in the dump.
	/// </summary>
	static string Number( float value )
	{
		if ( float.IsNaN( value ) ) return "nan";
		if ( float.IsPositiveInfinity( value ) ) return "inf";
		if ( float.IsNegativeInfinity( value ) ) return "-inf";

		var text = value.ToString( "R", CultureInfo.InvariantCulture );

		if ( text.Contains( '.' ) || text.Contains( 'E' ) || text.Contains( 'e' ) ) return text;

		return text + ".0";
	}

	sealed class Writer
	{
		readonly StringBuilder _sb;

		public Writer( StringBuilder sb, IrPrinterOptions options )
		{
			_sb = sb;
			Options = options;
		}

		public IrPrinterOptions Options { get; }

		int _depth;

		public void Push() => _depth++;

		public void Pop() => _depth = Math.Max( 0, _depth - 1 );

		public void Blank() => _sb.Append( Options.NewLine );

		public void Line( string text ) => Line( text, NodeId.None );

		public void Line( string text, NodeId origin )
		{
			for ( int i = 0; i < _depth; i++ ) _sb.Append( Options.Indent );

			_sb.Append( text );

			if ( Options.ShowOrigins && origin.IsValid ) _sb.Append( "   // @" ).Append( origin.Value );

			_sb.Append( Options.NewLine );
		}
	}
}