using Editor.Prism.Compiler.Ir;
using Editor.Prism.Core;

namespace Editor.Prism.Compiler;

/// <summary>Which optimisations run, and how hard.</summary>
public sealed record IrOptimizerOptions
{
	/// <summary>Evaluate operations whose operands are all literals.</summary>
	public bool ConstantFolding { get; init; } = true;

	/// <summary>Apply the identities that cost nothing to check: x*1, x+0, lerp(a,b,0), select on a literal.</summary>
	public bool AlgebraicIdentities { get; init; } = true;

	/// <summary>
	/// Allow the identities that are only true for finite values: <c>x * 0 → 0</c> and
	/// <c>pow( x, 0 ) → 1</c>.
	/// <para>
	/// Every other rule under <see cref="AlgebraicIdentities"/> is exact — <c>x + 0</c> and <c>x * 1</c>
	/// give back the same bits for a NaN or an infinity. These two do not: a graph that masks a division
	/// with <c>rcp( d ) * step( eps, d )</c> folds to exactly zero where the GPU would have produced
	/// NaN. That is usually what the author wanted and it is what DXC does anyway, so this defaults on —
	/// but it is the one place the optimiser is not semantics-preserving, so it is a switch rather than
	/// a hard-coded rule, and turning it off is the second thing to try when the optimised and
	/// unoptimised shaders disagree numerically.
	/// </para>
	/// </summary>
	public bool FastMath { get; init; } = true;

	/// <summary>Inline temps whose initialiser is a literal or a plain reference.</summary>
	public bool CopyPropagation { get; init; } = true;

	/// <summary>Merge declarations that end up holding structurally identical values.</summary>
	public bool RedundantDeclarations { get; init; } = true;

	/// <summary>Delete pure declarations nothing reads, and blocks that end up empty.</summary>
	public bool DeadCodeElimination { get; init; } = true;

	/// <summary>Delete helper functions nothing calls.</summary>
	public bool DropUnusedHelpers { get; init; } = true;

	/// <summary>Keep comments. Off in a release build where nobody reads the generated text.</summary>
	public bool KeepComments { get; init; } = true;

	/// <summary>How many times to sweep before accepting the result.</summary>
	public int MaxPasses { get; init; } = 8;

	/// <summary>Everything on. What a normal compile uses.</summary>
	public static IrOptimizerOptions Default { get; } = new();

	/// <summary>Everything off. Useful when a bug report needs the emitter's raw output.</summary>
	public static IrOptimizerOptions None { get; } = new()
	{
		ConstantFolding = false,
		AlgebraicIdentities = false,
		CopyPropagation = false,
		RedundantDeclarations = false,
		DeadCodeElimination = false,
		DropUnusedHelpers = false,
		FastMath = false,
		MaxPasses = 0
	};

	/// <summary>
	/// Everything on except the rules that assume finite values. Use this to check whether an
	/// optimised shader and its unoptimised twin disagree because of <see cref="FastMath"/>.
	/// </summary>
	public static IrOptimizerOptions Exact { get; } = new()
	{
		FastMath = false
	};
}

/// <summary>What the optimiser managed to remove.</summary>
public readonly record struct OptimizerStats( int Folded, int Propagated, int Merged, int Removed, int Passes )
{
	/// <summary>Total expressions and statements the optimiser eliminated.</summary>
	public int Total => Folded + Propagated + Merged + Removed;

	/// <inheritdoc/>
	public override string ToString() =>
		$"{Folded} folded, {Propagated} propagated, {Merged} merged, {Removed} removed in {Passes} passes";
}

/// <summary>
/// Cleans up the IR after emission.
/// <para>
/// A demand-driven emitter deliberately binds almost everything to a temp: that keeps the source map
/// dense and lets hash-consing merge work across nodes. The cost is a lot of one-line temps and a lot
/// of arithmetic on literals, because the graph is full of default values that never changed. This
/// pass folds the literals, applies the identities that follow, inlines the temps that turned out to
/// be plain references, merges declarations that ended up holding the same value, and deletes whatever
/// nothing reads — typically 20-40% of what the emitter produced.
/// </para>
/// <para>
/// Statement order is never permuted. The emitter already produces a topological, demand-ordered
/// sequence, and preserving it byte for byte is what makes "regenerate, compare text, skip the
/// compile" trustworthy.
/// </para>
/// </summary>
public static class IrOptimizer
{
	/// <summary>Optimise every function in a module.</summary>
	public static OptimizerStats Optimize( IrModule module, IrOptimizerOptions options = null,
		DiagnosticSink diagnostics = null )
	{
		options ??= IrOptimizerOptions.Default;

		if ( module is null ) return default;

		var folded = 0;
		var propagated = 0;
		var merged = 0;
		var removed = 0;
		var passes = 0;

		foreach ( var function in module.Functions )
		{
			if ( function?.Body is null ) continue;

			var stats = Optimize( function.Body, options );

			folded += stats.Folded;
			propagated += stats.Propagated;
			merged += stats.Merged;
			removed += stats.Removed;
			passes = Math.Max( passes, stats.Passes );
		}

		if ( options.DropUnusedHelpers ) removed += DropUnusedHelpers( module );

		_ = diagnostics;

		return new OptimizerStats( folded, propagated, merged, removed, passes );
	}

	/// <summary>Optimise one block in place.</summary>
	public static OptimizerStats Optimize( IrBlock block, IrOptimizerOptions options = null )
	{
		options ??= IrOptimizerOptions.Default;

		if ( block is null ) return default;

		var folded = 0;
		var propagated = 0;
		var merged = 0;
		var removed = 0;
		var passes = 0;

		var assigned = new HashSet<string>( StringComparer.Ordinal );
		CollectAssignedNames( block, assigned );

		for ( ; passes < Math.Max( 0, options.MaxPasses ); passes++ )
		{
			var pass = new Pass( options, assigned );

			pass.Rewrite( block );

			if ( options.DeadCodeElimination ) pass.Removed += Prune( block, options );

			folded += pass.Folded;
			propagated += pass.Propagated;
			merged += pass.Merged;
			removed += pass.Removed;

			if ( pass.Changes == 0 ) break;
		}

		return new OptimizerStats( folded, propagated, merged, removed, passes );
	}

	/// <summary>Fold one expression tree in isolation. Exposed for tests and for the IR panel.</summary>
	public static IrExpr Fold( IrExpr expr )
	{
		if ( expr is null ) return null;

		var pass = new Pass( IrOptimizerOptions.Default, new HashSet<string>( StringComparer.Ordinal ) );

		return pass.Visit( expr );
	}

	static int DropUnusedHelpers( IrModule module )
	{
		if ( module.Helpers.Count == 0 ) return 0;

		var used = new HashSet<string>( StringComparer.Ordinal );

		foreach ( var function in module.Functions )
		{
			CollectHelperCalls( function?.Body, used );
		}

		// A helper another live helper needs stays, transitively.
		var changed = true;

		while ( changed )
		{
			changed = false;

			foreach ( var helper in module.Helpers )
			{
				if ( !used.Contains( helper.Name ) ) continue;

				foreach ( var required in helper.Requires ?? Array.Empty<HelperFunction>() )
				{
					if ( required is null ) continue;
					if ( used.Add( required.Name ) ) changed = true;
				}
			}
		}

		var before = module.Helpers.Count;
		module.Helpers.RemoveAll( x => x is not null && !used.Contains( x.Name ) );

		return before - module.Helpers.Count;
	}

	static void CollectHelperCalls( IrBlock block, HashSet<string> used )
	{
		if ( block is null ) return;

		foreach ( var statement in block.Statements )
		{
			foreach ( var expr in Expressions( statement ) )
			{
				foreach ( var node in IrExprUtil.Walk( expr ) )
				{
					if ( node is IrHelperCall call && call.Fn is not null ) used.Add( call.Fn.Name );
				}
			}

			foreach ( var child in Blocks( statement ) ) CollectHelperCalls( child, used );
		}
	}

	static void CollectAssignedNames( IrBlock block, HashSet<string> names )
	{
		if ( block is null ) return;

		foreach ( var statement in block.Statements )
		{
			if ( statement is IrAssign assign )
			{
				foreach ( var node in IrExprUtil.Walk( assign.Target ) )
				{
					if ( node is IrVar variable ) names.Add( variable.Name );
				}
			}

			if ( statement is IrFor loop ) names.Add( loop.Var );

			foreach ( var child in Blocks( statement ) ) CollectAssignedNames( child, names );
		}
	}

	// ---- dead code --------------------------------------------------------

	static int Prune( IrBlock block, IrOptimizerOptions options )
	{
		// Run to a fixed point rather than once. Use counts are taken up front and never decremented as
		// declarations are deleted, so one sweep only ever removes the *last* link of a dead chain:
		// `t0 = f(); t1 = t0; t2 = t1;` sheds t2 this sweep, t1 the next. Leaving that to the optimiser's
		// own MaxPasses budget means a chain longer than the budget ships its dead temps in the
		// generated text. Recounting is a cheap walk and the loop converges because removal is monotone.
		var total = 0;

		for ( int sweep = 0; sweep < MaxPruneSweeps; sweep++ )
		{
			var uses = new Dictionary<string, int>( StringComparer.Ordinal );
			CountUses( block, uses );

			var removed = Prune( block, uses, options );

			total += removed;

			if ( removed == 0 ) break;
		}

		return total;
	}

	/// <summary>
	/// Ceiling on the dead-code fixed point. Removal is monotone so the loop always converges long
	/// before this; the cap only exists so a future non-monotone rule cannot hang a compile.
	/// </summary>
	const int MaxPruneSweeps = 64;

	static int Prune( IrBlock block, Dictionary<string, int> uses, IrOptimizerOptions options )
	{
		if ( block is null ) return 0;

		var removed = 0;

		for ( int i = block.Statements.Count - 1; i >= 0; i-- )
		{
			var statement = block.Statements[i];

			foreach ( var child in Blocks( statement ) ) removed += Prune( child, uses, options );

			switch ( statement )
			{
				case IrDecl decl when Count( uses, decl.Name ) == 0 && ( decl.Init is null || decl.Init.Pure ):
					block.Statements.RemoveAt( i );
					removed++;
					break;

				case IrComment when !options.KeepComments:
					block.Statements.RemoveAt( i );
					removed++;
					break;

				case IrIf branch when IsEmpty( branch.Then ) && IsEmpty( branch.Else ) && branch.Cond is { Pure: true }:
					block.Statements.RemoveAt( i );
					removed++;
					break;

				// The purity guard matters as much here as on the IrIf above: an empty body does not make
				// the loop bound dead, and a count or condition built from an atomic or a texture write
				// has to run whatever the body does.
				case IrFor loop when IsEmpty( loop.Body ) && loop.Count is null or { Pure: true }:
					block.Statements.RemoveAt( i );
					removed++;
					break;

				case IrWhile loop when IsEmpty( loop.Body ) && loop.Cond is null or { Pure: true }:
					block.Statements.RemoveAt( i );
					removed++;
					break;

				case IrScope scope when IsEmpty( scope.Body ):
					block.Statements.RemoveAt( i );
					removed++;
					break;
			}
		}

		return removed;
	}

	static bool IsEmpty( IrBlock block ) => block is null || block.Statements.Count == 0;

	static int Count( Dictionary<string, int> uses, string name ) =>
		name is not null && uses.TryGetValue( name, out var count ) ? count : 0;

	static void CountUses( IrBlock block, Dictionary<string, int> uses )
	{
		if ( block is null ) return;

		foreach ( var statement in block.Statements )
		{
			foreach ( var expr in Expressions( statement ) )
			{
				foreach ( var node in IrExprUtil.Walk( expr ) )
				{
					if ( node is not IrVar variable ) continue;

					uses.TryGetValue( variable.Name, out var count );
					uses[variable.Name] = count + 1;
				}
			}

			foreach ( var child in Blocks( statement ) ) CountUses( child, uses );
		}
	}

	// ---- statement shape --------------------------------------------------

	static IEnumerable<IrExpr> Expressions( IrStmt statement )
	{
		switch ( statement )
		{
			case IrDecl decl when decl.Init is not null:
				yield return decl.Init;
				break;

			case IrAssign assign:
				// The target's own sub-expressions are reads; the outermost variable is a write.
				if ( assign.Target is not null ) yield return assign.Target;
				if ( assign.Value is not null ) yield return assign.Value;
				break;

			case IrIf branch when branch.Cond is not null:
				yield return branch.Cond;
				break;

			case IrFor loop when loop.Count is not null:
				yield return loop.Count;
				break;

			case IrWhile loop when loop.Cond is not null:
				yield return loop.Cond;
				break;

			case IrReturn ret when ret.Value is not null:
				yield return ret.Value;
				break;

			case IrExprStmt expr when expr.Value is not null:
				yield return expr.Value;
				break;
		}
	}

	static IEnumerable<IrBlock> Blocks( IrStmt statement )
	{
		switch ( statement )
		{
			case IrIf branch:
				if ( branch.Then is not null ) yield return branch.Then;
				if ( branch.Else is not null ) yield return branch.Else;
				break;

			case IrFor loop when loop.Body is not null:
				yield return loop.Body;
				break;

			case IrWhile loop when loop.Body is not null:
				yield return loop.Body;
				break;

			case IrScope scope when scope.Body is not null:
				yield return scope.Body;
				break;

			// A preprocessor guard is a real scope as far as liveness is concerned. Without this a temp
			// declared outside the guard and read only inside it counts as unused, gets pruned, and the
			// generated code then names an undeclared local.
			case IrPreprocessorIf guard:
				if ( guard.Then is not null ) yield return guard.Then;
				if ( guard.Else is not null ) yield return guard.Else;
				break;
		}
	}

	/// <summary>One rewrite sweep over a block, carrying a scoped substitution environment.</summary>
	sealed class Pass
	{
		readonly IrOptimizerOptions _options;
		readonly HashSet<string> _assigned;
		readonly List<Dictionary<string, IrExpr>> _substitutions = new();
		readonly List<Dictionary<IrExpr, string>> _available = new();

		public Pass( IrOptimizerOptions options, HashSet<string> assigned )
		{
			_options = options;
			_assigned = assigned;
		}

		public int Folded;
		public int Propagated;
		public int Merged;
		public int Removed;

		public int Changes => Folded + Propagated + Merged + Removed;

		public void Rewrite( IrBlock block )
		{
			if ( block is null ) return;

			_substitutions.Add( new Dictionary<string, IrExpr>( StringComparer.Ordinal ) );
			_available.Add( new Dictionary<IrExpr, string>() );

			for ( int i = 0; i < block.Statements.Count; i++ )
			{
				block.Statements[i] = Rewrite( block.Statements[i] );
			}

			_substitutions.RemoveAt( _substitutions.Count - 1 );
			_available.RemoveAt( _available.Count - 1 );
		}

		IrStmt Rewrite( IrStmt statement )
		{
			switch ( statement )
			{
				case IrDecl decl:
				{
					var init = Visit( decl.Init );

					// !_assigned is the same guard copy propagation carries below, and for the same
					// reason: substituting every read of a name that is later assigned to makes the
					// assignment dead and the reads answer with the pre-assignment value. The type match
					// pairs with the one on the insert, so a merged name always has the type its
					// initialiser had.
					if ( init is not null && _options.RedundantDeclarations && init.Pure &&
						!_assigned.Contains( decl.Name ) && init.Type == decl.Type &&
						TryFindAvailable( init, out var existing ) && existing != decl.Name )
					{
						Merged++;
						Substitute( decl.Name, new IrVar( decl.Type, existing ) );

						return decl with { Init = init };
					}

					if ( init is not null && _options.CopyPropagation && IsTrivial( init ) &&
						!_assigned.Contains( decl.Name ) && init.Type == decl.Type )
					{
						Propagated++;
						Substitute( decl.Name, init );
					}
					else if ( init is not null && init.Pure && !_assigned.Contains( decl.Name ) &&
						init.Type == decl.Type )
					{
						// Only offer a name whose declared type is its initialiser's. IrBuilder.Declare
						// lets a caller declare a wider or narrower type than the expression it stores,
						// and merging onto one of those would change what the reader sees.
						_available[^1][init] = decl.Name;
					}

					return ReferenceEquals( init, decl.Init ) ? decl : decl with { Init = init };
				}

				case IrAssign assign:
					return assign with { Target = VisitTarget( assign.Target ), Value = Visit( assign.Value ) };

				case IrIf branch:
				{
					var condition = Visit( branch.Cond );

					if ( _options.AlgebraicIdentities && condition is IrConst literal )
					{
						// A branch on a literal is not a branch. Keep the taken side's statements only.
						var taken = literal.Value.IsZero( 1 ) ? branch.Else : branch.Then;

						Folded++;
						Rewrite( taken );

						return new IrScope( branch.Origin, taken ?? new IrBlock() );
					}

					Rewrite( branch.Then );
					Rewrite( branch.Else );

					return branch with { Cond = condition };
				}

				case IrFor loop:
				{
					var count = Visit( loop.Count );
					Rewrite( loop.Body );
					return loop with { Count = count };
				}

				case IrWhile loop:
				{
					var condition = Visit( loop.Cond );
					Rewrite( loop.Body );
					return loop with { Cond = condition };
				}

				case IrReturn ret:
					return ret with { Value = Visit( ret.Value ) };

				case IrExprStmt expr:
					return expr with { Value = Visit( expr.Value ) };

				case IrScope scope:
					Rewrite( scope.Body );
					return scope;

				// Both sides are rewritten, but the condition is deliberately not an IrExpr and never
				// becomes one: a preprocessor condition may only name combo symbols, so there is
				// nothing here to visit or fold.
				case IrPreprocessorIf guard:
					Rewrite( guard.Then );
					Rewrite( guard.Else );
					return guard;

				default:
					return statement;
			}
		}

		void Substitute( string name, IrExpr replacement )
		{
			if ( string.IsNullOrEmpty( name ) || replacement is null ) return;

			_substitutions[^1][name] = replacement;
		}

		bool TryFindAvailable( IrExpr expr, out string name )
		{
			for ( int i = _available.Count - 1; i >= 0; i-- )
			{
				if ( _available[i].TryGetValue( expr, out name ) ) return true;
			}

			name = null;
			return false;
		}

		bool TryLookup( string name, out IrExpr replacement )
		{
			for ( int i = _substitutions.Count - 1; i >= 0; i-- )
			{
				if ( _substitutions[i].TryGetValue( name, out replacement ) ) return true;
			}

			replacement = null;
			return false;
		}

		/// <summary>
		/// Rewrite an assignment target. The outermost variable is a write and must survive; anything
		/// inside it — an index expression, say — is a read and folds normally.
		/// </summary>
		IrExpr VisitTarget( IrExpr expr ) => expr switch
		{
			null => null,
			IrVar => expr,
			IrMember member => member with { V = VisitTarget( member.V ) },
			IrSwizzle swizzle => swizzle with { V = VisitTarget( swizzle.V ) },
			IrIndex index => index with { V = VisitTarget( index.V ), I = Visit( index.I ) },
			_ => expr
		};

		/// <summary>Rewrite one expression tree bottom-up.</summary>
		public IrExpr Visit( IrExpr expr )
		{
			if ( expr is null ) return null;

			var rewritten = expr switch
			{
				IrVar variable when _options.CopyPropagation && TryLookup( variable.Name, out var replacement ) =>
					replacement,

				IrCall call => call with { Args = VisitAll( call.Args ) },
				IrHelperCall helper => helper with { Args = VisitAll( helper.Args ) },
				IrBinary bin => bin with { L = Visit( bin.L ), R = Visit( bin.R ) },
				IrUnary un => un with { V = Visit( un.V ) },
				IrSwizzle sw => sw with { V = Visit( sw.V ) },
				IrConstruct ctor => ctor with { Parts = VisitAll( ctor.Parts ) },
				IrCast cast => cast with { V = Visit( cast.V ) },
				IrSelect sel => sel with { C = Visit( sel.C ), A = Visit( sel.A ), B = Visit( sel.B ) },
				IrIndex index => index with { V = Visit( index.V ), I = Visit( index.I ) },
				IrMember member => member with { V = Visit( member.V ) },
				_ => expr
			};

			var folded = IrFolding.Fold( rewritten, _options );

			if ( !ReferenceEquals( folded, rewritten ) ) Folded++;

			return folded;
		}

		IrExpr[] VisitAll( IrExpr[] items )
		{
			if ( items is null || items.Length == 0 ) return items;

			var result = new IrExpr[items.Length];

			for ( int i = 0; i < items.Length; i++ ) result[i] = Visit( items[i] );

			return result;
		}

		static bool IsTrivial( IrExpr expr ) => expr switch
		{
			IrConst => true,
			IrVar => true,
			IrGlobalRef => true,
			IrBuiltinRef => true,
			IrSwizzle s => s.V is IrVar or IrGlobalRef or IrBuiltinRef,
			IrMember m => m.V is IrVar or IrGlobalRef,
			_ => false
		};
	}
}

/// <summary>
/// Rewrites a single expression node once its children are already folded.
/// <para>
/// Two families of rule live here. Constant evaluation replaces an operation whose operands are all
/// literals with the literal it produces — but only when the result is finite, because emitting a
/// baked-in NaN is worse than leaving the call alone. Algebraic identities remove the work a graph
/// editor generates by construction: multiplying by an untouched default of 1, adding an untouched
/// default of 0, a lerp whose factor is pinned at either end, a select on a fixed condition, and the
/// swizzle-of-construct pattern every split-and-recombine chain produces.
/// </para>
/// </summary>
public static class IrFolding
{
	/// <summary>Fold one node. Children are assumed to be folded already.</summary>
	public static IrExpr Fold( IrExpr expr, IrOptimizerOptions options )
	{
		if ( expr is null ) return null;

		options ??= IrOptimizerOptions.Default;

		return expr switch
		{
			IrCast cast => FoldCast( cast, options ),
			IrSwizzle swizzle => FoldSwizzle( swizzle, options ),
			IrConstruct ctor => FoldConstruct( ctor, options ),
			IrUnary un => FoldUnary( un, options ),
			IrBinary bin => FoldBinary( bin, options ),
			IrSelect sel => FoldSelect( sel, options ),
			IrCall call => FoldCall( call, options ),
			_ => expr
		};
	}

	// ---- casts ------------------------------------------------------------

	static IrExpr FoldCast( IrCast cast, IrOptimizerOptions options )
	{
		if ( cast.V is null ) return cast;
		if ( cast.V.Type == cast.Type ) return cast.V;

		if ( !options.ConstantFolding ) return cast;
		if ( cast.V is not IrConst literal ) return cast;

		var value = Convert( literal.Value, literal.Type, cast.Type, cast.Kind, cast.Fill );

		return Literal( cast.Type, value );
	}

	static ConstValue Convert( ConstValue value, ShaderType from, ShaderType to, CastKind kind, float fill )
	{
		var source = Math.Clamp( from.Components, 1, 4 );
		var target = Math.Clamp( to.Components, 1, 4 );

		var components = new double[4];

		for ( int i = 0; i < 4; i++ )
		{
			if ( kind == CastKind.Splat || source == 1 )
			{
				components[i] = value.X;
			}
			else if ( i < source && i < target )
			{
				components[i] = value[i];
			}
			else if ( i < target )
			{
				components[i] = fill;
			}
			else
			{
				components[i] = 0;
			}
		}

		var result = new ConstValue( components[0], components[1], components[2], components[3] );

		return Quantize( result, to );
	}

	// ---- swizzles ---------------------------------------------------------

	static IrExpr FoldSwizzle( IrSwizzle swizzle, IrOptimizerOptions options )
	{
		if ( swizzle.V is null || string.IsNullOrEmpty( swizzle.Mask ) ) return swizzle;

		// A swizzle that names every component in order is not a swizzle.
		if ( swizzle.V.Type == swizzle.Type && IsIdentity( swizzle.Mask, swizzle.V.Type ) ) return swizzle.V;

		if ( options.ConstantFolding && swizzle.V is IrConst literal )
		{
			var picked = new double[4];

			for ( int i = 0; i < 4; i++ )
			{
				var index = i < swizzle.Mask.Length ? Channel( swizzle.Mask[i] ) : Channel( swizzle.Mask[^1] );
				picked[i] = index < 0 ? 0 : literal.Value[index];
			}

			return Literal( swizzle.Type, new ConstValue( picked[0], picked[1], picked[2], picked[3] ) );
		}

		if ( !options.AlgebraicIdentities ) return swizzle;

		// .xy of .zwxy composes into .zw rather than nesting.
		if ( swizzle.V is IrSwizzle inner && !string.IsNullOrEmpty( inner.Mask ) )
		{
			var composed = new char[swizzle.Mask.Length];

			for ( int i = 0; i < swizzle.Mask.Length; i++ )
			{
				var index = Channel( swizzle.Mask[i] );

				if ( index < 0 || index >= inner.Mask.Length ) return swizzle;

				composed[i] = inner.Mask[index];
			}

			return new IrSwizzle( swizzle.Type, inner.V, new string( composed ) ) { Pure = inner.Pure };
		}

		if ( swizzle.V is IrConstruct ctor ) return FoldSwizzleOfConstruct( swizzle, ctor );

		return swizzle;
	}

	/// <summary>
	/// <c>float3( a, b, c ).xy</c> becomes <c>float2( a, b )</c>. Only applied when nothing is
	/// duplicated and every part is pure, so no side effect is repeated or dropped.
	/// </summary>
	static IrExpr FoldSwizzleOfConstruct( IrSwizzle swizzle, IrConstruct ctor )
	{
		if ( ctor.Parts is null || ctor.Parts.Length == 0 ) return swizzle;
		if ( !ctor.Pure ) return swizzle;
		if ( HasRepeats( swizzle.Mask ) ) return swizzle;

		var sources = new List<IrExpr>( 4 );

		foreach ( var part in ctor.Parts )
		{
			if ( part is null || !part.Type.IsScalarOrVector ) return swizzle;

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

			for ( int i = 0; i < width; i++ )
			{
				sources.Add( width == 1
					? part
					: new IrSwizzle( part.Type.WithComponents( 1 ), part, "xyzw"[i].ToString() ) { Pure = part.Pure } );
			}
		}

		var picked = new IrExpr[swizzle.Mask.Length];

		for ( int i = 0; i < swizzle.Mask.Length; i++ )
		{
			var index = Channel( swizzle.Mask[i] );

			if ( index < 0 || index >= sources.Count ) return swizzle;

			picked[i] = sources[index];
		}

		if ( picked.Length == 1 ) return picked[0].Type == swizzle.Type ? picked[0] : swizzle;

		return new IrConstruct( swizzle.Type, picked ) { Pure = true };
	}

	static IrExpr FoldConstruct( IrConstruct ctor, IrOptimizerOptions options )
	{
		if ( ctor.Parts is null || ctor.Parts.Length == 0 ) return ctor;

		if ( ctor.Parts.Length == 1 && ctor.Parts[0] is not null && ctor.Parts[0].Type == ctor.Type )
		{
			return ctor.Parts[0];
		}

		if ( !options.ConstantFolding ) return ctor;

		var components = new double[4];
		var filled = 0;

		foreach ( var part in ctor.Parts )
		{
			if ( part is not IrConst literal ) return ctor;

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

			for ( int i = 0; i < width && filled < 4; i++ ) components[filled++] = literal.Value[i];
		}

		if ( filled == 0 ) return ctor;

		return Literal( ctor.Type, new ConstValue( components[0], components[1], components[2], components[3] ) );
	}

	// ---- operators --------------------------------------------------------

	static IrExpr FoldUnary( IrUnary un, IrOptimizerOptions options )
	{
		if ( un.V is null ) return un;

		if ( options.AlgebraicIdentities && un.V is IrUnary inner && inner.Op == un.Op &&
			un.Op is UnaryOp.Negate or UnaryOp.LogicalNot or UnaryOp.BitNot &&
			inner.V is not null && inner.V.Type == un.Type )
		{
			return inner.V;
		}

		if ( !options.ConstantFolding || un.V is not IrConst literal ) return un;

		var width = Width( un.Type );

		var folded = un.Op switch
		{
			UnaryOp.Negate => Map( literal.Value, width, x => -x ),
			UnaryOp.LogicalNot => Map( literal.Value, width, x => x != 0 ? 0 : 1 ),
			_ => Map( literal.Value, width, x => ~(long)x )
		};

		return Literal( un.Type, folded );
	}

	static IrExpr FoldBinary( IrBinary bin, IrOptimizerOptions options )
	{
		if ( bin.L is null || bin.R is null ) return bin;

		var width = Width( bin.Type );

		if ( options.ConstantFolding && bin.L is IrConst a && bin.R is IrConst b )
		{
			if ( TryFoldBinary( bin.Op, a.Value, b.Value, width, out var value ) )
			{
				return Literal( bin.Type, value );
			}
		}

		if ( !options.AlgebraicIdentities ) return bin;

		var left = bin.L as IrConst;
		var right = bin.R as IrConst;

		switch ( bin.Op )
		{
			case BinaryOp.Add:
				if ( left is not null && left.Value.IsZero( width ) && bin.R.Type == bin.Type ) return bin.R;
				if ( right is not null && right.Value.IsZero( width ) && bin.L.Type == bin.Type ) return bin.L;
				break;

			case BinaryOp.Sub:
				if ( right is not null && right.Value.IsZero( width ) && bin.L.Type == bin.Type ) return bin.L;
				break;

			case BinaryOp.Mul:
				if ( left is not null && left.Value.IsOne( width ) && bin.R.Type == bin.Type ) return bin.R;
				if ( right is not null && right.Value.IsOne( width ) && bin.L.Type == bin.Type ) return bin.L;
				// Zero absorbs only under FastMath: NaN * 0 is NaN, not 0. See IrOptimizerOptions.FastMath.
				if ( options.FastMath && left is not null && left.Value.IsZero( width ) && bin.R.Pure )
					return Literal( bin.Type, ConstValue.Zero );
				if ( options.FastMath && right is not null && right.Value.IsZero( width ) && bin.L.Pure )
					return Literal( bin.Type, ConstValue.Zero );
				break;

			case BinaryOp.Div:
				if ( right is not null && right.Value.IsOne( width ) && bin.L.Type == bin.Type ) return bin.L;
				break;
		}

		return bin;
	}

	static bool TryFoldBinary( BinaryOp op, ConstValue a, ConstValue b, int width, out ConstValue value )
	{
		value = default;

		if ( op is BinaryOp.Div or BinaryOp.Mod )
		{
			for ( int i = 0; i < width; i++ )
			{
				if ( b[i] == 0 ) return false;
			}
		}

		value = op switch
		{
			BinaryOp.Add => Map2( a, b, width, ( x, y ) => x + y ),
			BinaryOp.Sub => Map2( a, b, width, ( x, y ) => x - y ),
			BinaryOp.Mul => Map2( a, b, width, ( x, y ) => x * y ),
			BinaryOp.Div => Map2( a, b, width, ( x, y ) => x / y ),
			BinaryOp.Mod => Map2( a, b, width, ( x, y ) => x - y * Math.Truncate( x / y ) ),
			BinaryOp.Less => Map2( a, b, width, ( x, y ) => x < y ? 1 : 0 ),
			BinaryOp.LessEqual => Map2( a, b, width, ( x, y ) => x <= y ? 1 : 0 ),
			BinaryOp.Greater => Map2( a, b, width, ( x, y ) => x > y ? 1 : 0 ),
			BinaryOp.GreaterEqual => Map2( a, b, width, ( x, y ) => x >= y ? 1 : 0 ),
			BinaryOp.Equal => Map2( a, b, width, ( x, y ) => x == y ? 1 : 0 ),
			BinaryOp.NotEqual => Map2( a, b, width, ( x, y ) => x != y ? 1 : 0 ),
			BinaryOp.LogicalAnd => Map2( a, b, width, ( x, y ) => x != 0 && y != 0 ? 1 : 0 ),
			BinaryOp.LogicalOr => Map2( a, b, width, ( x, y ) => x != 0 || y != 0 ? 1 : 0 ),
			BinaryOp.BitAnd => Map2( a, b, width, ( x, y ) => (long)x & (long)y ),
			BinaryOp.BitOr => Map2( a, b, width, ( x, y ) => (long)x | (long)y ),
			BinaryOp.BitXor => Map2( a, b, width, ( x, y ) => (long)x ^ (long)y ),
			BinaryOp.Shl => Map2( a, b, width, ( x, y ) => (long)x << (int)y ),
			_ => Map2( a, b, width, ( x, y ) => (long)x >> (int)y )
		};

		return IsFinite( value, width );
	}

	static IrExpr FoldSelect( IrSelect sel, IrOptimizerOptions options )
	{
		if ( !options.AlgebraicIdentities ) return sel;
		if ( sel.C is not IrConst condition ) return sel;
		if ( sel.A is null || sel.B is null ) return sel;

		var width = Math.Max( Width( sel.C.Type ), 1 );
		var allTrue = true;
		var allFalse = true;

		for ( int i = 0; i < width; i++ )
		{
			if ( condition.Value[i] != 0 ) allFalse = false;
			else allTrue = false;
		}

		if ( allTrue && sel.B.Pure && sel.A.Type == sel.Type ) return sel.A;
		if ( allFalse && sel.A.Pure && sel.B.Type == sel.Type ) return sel.B;

		return sel;
	}

	// ---- intrinsics -------------------------------------------------------

	static IrExpr FoldCall( IrCall call, IrOptimizerOptions options )
	{
		var args = call.Args ?? Array.Empty<IrExpr>();

		if ( options.AlgebraicIdentities )
		{
			var identity = FoldCallIdentity( call, args );
			if ( identity is not null ) return identity;
		}

		if ( !options.ConstantFolding ) return call;
		if ( !IntrinsicCatalog.Get( call.Id ).Pure ) return call;
		if ( args.Length == 0 ) return call;

		foreach ( var arg in args )
		{
			if ( arg is not IrConst ) return call;
		}

		var values = new ConstValue[args.Length];
		var widths = new int[args.Length];

		for ( int i = 0; i < args.Length; i++ )
		{
			values[i] = ( (IrConst)args[i] ).Value;
			widths[i] = Width( args[i].Type );
		}

		var width = Width( call.Type );

		if ( !TryEvaluate( call.Id, values, widths, width, out var folded ) ) return call;
		if ( !IsFinite( folded, width ) ) return call;

		return Literal( call.Type, folded );
	}

	static IrExpr FoldCallIdentity( IrCall call, IrExpr[] args )
	{
		switch ( call.Id )
		{
			case Intrinsic.Lerp when args.Length == 3 && args[2] is IrConst t:
			{
				var width = Width( args[2].Type );

				if ( t.Value.IsZero( width ) && args[0].Type == call.Type && args[1].Pure ) return args[0];
				if ( t.Value.IsOne( width ) && args[1].Type == call.Type && args[0].Pure ) return args[1];

				break;
			}

			case Intrinsic.Pow when args.Length == 2 && args[1] is IrConst exponent:
			{
				var width = Width( args[1].Type );

				if ( exponent.Value.IsOne( width ) && args[0].Type == call.Type ) return args[0];
				if ( exponent.Value.IsZero( width ) && args[0].Pure ) return Literal( call.Type, ConstValue.One );

				break;
			}

			case Intrinsic.Saturate when args.Length == 1 && args[0] is IrCall inner && inner.Id == Intrinsic.Saturate:
				return inner;

			case Intrinsic.Abs when args.Length == 1 && args[0] is IrCall abs && abs.Id == Intrinsic.Abs:
				return abs;

			case Intrinsic.Normalize when args.Length == 1 && args[0] is IrCall norm && norm.Id == Intrinsic.Normalize:
				return norm;

			case Intrinsic.Min when args.Length == 2 && args[0].Pure && Equals( args[0], args[1] ) &&
				args[0].Type == call.Type:
				return args[0];

			case Intrinsic.Max when args.Length == 2 && args[0].Pure && Equals( args[0], args[1] ) &&
				args[0].Type == call.Type:
				return args[0];
		}

		return null;
	}

	static bool TryEvaluate( Intrinsic id, ConstValue[] a, int[] widths, int width, out ConstValue value )
	{
		value = default;

		switch ( id )
		{
			case Intrinsic.Abs: value = Map( a[0], width, Math.Abs ); return true;
			case Intrinsic.Saturate: value = Map( a[0], width, x => Math.Clamp( x, 0, 1 ) ); return true;
			case Intrinsic.Frac: value = Map( a[0], width, x => x - Math.Floor( x ) ); return true;
			case Intrinsic.Ceil: value = Map( a[0], width, Math.Ceiling ); return true;
			case Intrinsic.Floor: value = Map( a[0], width, Math.Floor ); return true;
			case Intrinsic.Round: value = Map( a[0], width, x => Math.Round( x, MidpointRounding.ToEven ) ); return true;
			case Intrinsic.Trunc: value = Map( a[0], width, Math.Truncate ); return true;
			case Intrinsic.Sign: value = Map( a[0], width, x => Math.Sign( x ) ); return true;
			case Intrinsic.Rcp: value = Map( a[0], width, x => x == 0 ? double.NaN : 1.0 / x ); return true;
			case Intrinsic.Sqrt: value = Map( a[0], width, x => x < 0 ? double.NaN : Math.Sqrt( x ) ); return true;
			case Intrinsic.Rsqrt: value = Map( a[0], width, x => x <= 0 ? double.NaN : 1.0 / Math.Sqrt( x ) ); return true;
			case Intrinsic.Exp: value = Map( a[0], width, Math.Exp ); return true;
			case Intrinsic.Exp2: value = Map( a[0], width, x => Math.Pow( 2, x ) ); return true;
			case Intrinsic.Log: value = Map( a[0], width, x => x <= 0 ? double.NaN : Math.Log( x ) ); return true;
			case Intrinsic.Log2: value = Map( a[0], width, x => x <= 0 ? double.NaN : Math.Log2( x ) ); return true;
			case Intrinsic.Log10: value = Map( a[0], width, x => x <= 0 ? double.NaN : Math.Log10( x ) ); return true;
			case Intrinsic.Sin: value = Map( a[0], width, Math.Sin ); return true;
			case Intrinsic.Cos: value = Map( a[0], width, Math.Cos ); return true;
			case Intrinsic.Tan: value = Map( a[0], width, Math.Tan ); return true;
			case Intrinsic.Asin: value = Map( a[0], width, Math.Asin ); return true;
			case Intrinsic.Acos: value = Map( a[0], width, Math.Acos ); return true;
			case Intrinsic.Atan: value = Map( a[0], width, Math.Atan ); return true;
			case Intrinsic.Sinh: value = Map( a[0], width, Math.Sinh ); return true;
			case Intrinsic.Cosh: value = Map( a[0], width, Math.Cosh ); return true;
			case Intrinsic.Tanh: value = Map( a[0], width, Math.Tanh ); return true;
			case Intrinsic.Degrees: value = Map( a[0], width, x => x * ( 180.0 / Math.PI ) ); return true;
			case Intrinsic.Radians: value = Map( a[0], width, x => x * ( Math.PI / 180.0 ) ); return true;

			case Intrinsic.Min when a.Length >= 2: value = Map2( a[0], a[1], width, Math.Min ); return true;
			case Intrinsic.Max when a.Length >= 2: value = Map2( a[0], a[1], width, Math.Max ); return true;
			case Intrinsic.Pow when a.Length >= 2: value = Map2( a[0], a[1], width, Math.Pow ); return true;
			case Intrinsic.Atan2 when a.Length >= 2: value = Map2( a[0], a[1], width, Math.Atan2 ); return true;
			case Intrinsic.Step when a.Length >= 2: value = Map2( a[0], a[1], width, ( e, x ) => x >= e ? 1 : 0 ); return true;

			case Intrinsic.Fmod when a.Length >= 2:
				for ( int i = 0; i < width; i++ )
				{
					if ( a[1][i] == 0 ) return false;
				}

				value = Map2( a[0], a[1], width, ( x, y ) => x - y * Math.Truncate( x / y ) );
				return true;

			case Intrinsic.Clamp when a.Length >= 3:
				value = Map3( a[0], a[1], a[2], width, ( x, lo, hi ) => Math.Clamp( x, Math.Min( lo, hi ), Math.Max( lo, hi ) ) );
				return true;

			case Intrinsic.Lerp when a.Length >= 3:
				value = Map3( a[0], a[1], a[2], width, ( x, y, t ) => x + t * ( y - x ) );
				return true;

			case Intrinsic.Fma when a.Length >= 3:
			case Intrinsic.Mad when a.Length >= 3:
				value = Map3( a[0], a[1], a[2], width, ( x, y, z ) => x * y + z );
				return true;

			case Intrinsic.SmoothStep when a.Length >= 3:
				value = Map3( a[0], a[1], a[2], width, ( lo, hi, x ) =>
				{
					if ( hi == lo ) return x < lo ? 0 : 1;

					var t = Math.Clamp( ( x - lo ) / ( hi - lo ), 0, 1 );
					return t * t * ( 3 - 2 * t );
				} );

				return true;

			case Intrinsic.Dot when a.Length >= 2:
			{
				var sum = 0.0;
				var n = Math.Max( widths[0], widths[1] );

				for ( int i = 0; i < n; i++ ) sum += a[0][i] * a[1][i];

				value = ConstValue.From( (float)sum );
				return true;
			}

			case Intrinsic.Length when a.Length >= 1:
			{
				var sum = 0.0;

				for ( int i = 0; i < widths[0]; i++ ) sum += a[0][i] * a[0][i];

				value = ConstValue.From( (float)Math.Sqrt( sum ) );
				return true;
			}

			case Intrinsic.Distance when a.Length >= 2:
			{
				var sum = 0.0;
				var n = Math.Max( widths[0], widths[1] );

				for ( int i = 0; i < n; i++ )
				{
					var d = a[0][i] - a[1][i];
					sum += d * d;
				}

				value = ConstValue.From( (float)Math.Sqrt( sum ) );
				return true;
			}

			case Intrinsic.Normalize when a.Length >= 1:
			{
				var sum = 0.0;

				for ( int i = 0; i < widths[0]; i++ ) sum += a[0][i] * a[0][i];

				if ( sum <= 0 ) return false;

				var scale = 1.0 / Math.Sqrt( sum );
				value = Map( a[0], width, x => x * scale );

				return true;
			}

			case Intrinsic.All when a.Length >= 1:
			{
				var all = true;

				for ( int i = 0; i < widths[0]; i++ )
				{
					if ( a[0][i] == 0 ) all = false;
				}

				value = ConstValue.From( all );
				return true;
			}

			case Intrinsic.Any when a.Length >= 1:
			{
				var any = false;

				for ( int i = 0; i < widths[0]; i++ )
				{
					if ( a[0][i] != 0 ) any = true;
				}

				value = ConstValue.From( any );
				return true;
			}

			case Intrinsic.Select when a.Length >= 3:
				value = Map3( a[0], a[1], a[2], width, ( c, x, y ) => c != 0 ? x : y );
				return true;

			default:
				return false;
		}
	}

	// ---- numeric plumbing -------------------------------------------------

	static IrConst Literal( ShaderType type, ConstValue value ) => new( type, Canonical( Quantize( value, type ), type ) );

	/// <summary>
	/// Zero the lanes a type does not use.
	/// <para>
	/// <see cref="Map"/> deliberately evaluates all four components whatever the width, which keeps two
	/// structurally identical folds bit-identical — but it also means folding <c>rcp( float2( 2, 4 ) )</c>
	/// evaluates <c>1/0</c> for the two lanes the literal does not have and stores
	/// <c>(0.5, 0.25, NaN, NaN)</c>. <see cref="Literal"/> prints only the used lanes so the generated
	/// text is right, but <c>IrConst</c>'s record equality and <c>IrBuilder.StructuralHash</c> read all
	/// four — so that literal would never hash-cons or CSE against the same <c>float2(0.5, 0.25)</c>
	/// reached another way. Canonicalising here costs nothing and keeps equality meaning what it says.
	/// </para>
	/// </summary>
	static ConstValue Canonical( ConstValue value, ShaderType type )
	{
		if ( !type.IsScalarOrVector ) return value;

		return Width( type ) switch
		{
			1 => value with { Y = 0, Z = 0, W = 0 },
			2 => value with { Z = 0, W = 0 },
			3 => value with { W = 0 },
			_ => value
		};
	}

	static int Width( ShaderType type ) => Math.Clamp( type.Components, 1, 4 );

	static int Channel( char c ) => c switch
	{
		'x' or 'r' or 's' => 0,
		'y' or 'g' or 't' => 1,
		'z' or 'b' or 'p' => 2,
		'w' or 'a' or 'q' => 3,
		_ => -1
	};

	static bool HasRepeats( string mask )
	{
		for ( int i = 0; i < mask.Length; i++ )
		{
			for ( int j = i + 1; j < mask.Length; j++ )
			{
				if ( Channel( mask[i] ) == Channel( mask[j] ) ) return true;
			}
		}

		return false;
	}

	static bool IsIdentity( string mask, ShaderType type )
	{
		if ( mask.Length != Width( type ) ) return false;

		for ( int i = 0; i < mask.Length; i++ )
		{
			if ( Channel( mask[i] ) != i ) return false;
		}

		return true;
	}

	static bool IsFinite( ConstValue value, int width )
	{
		for ( int i = 0; i < width; i++ )
		{
			if ( double.IsNaN( value[i] ) || double.IsInfinity( value[i] ) ) return false;
		}

		return true;
	}

	/// <summary>Snap a folded value onto what its type can actually hold.</summary>
	static ConstValue Quantize( ConstValue value, ShaderType type )
	{
		if ( type.IsBoolean ) return Map( value, 4, x => x != 0 ? 1 : 0 );
		if ( type.IsIntegral ) return Map( value, 4, Math.Truncate );

		// Everything downstream is 32-bit; rounding here keeps the printed literal honest.
		return Map( value, 4, x => (float)x );
	}

	/// <summary>
	/// Apply a function component-wise. All four components are always evaluated — the unused ones are
	/// never printed, and doing so keeps two structurally identical folds bit-identical, which is what
	/// deterministic emission depends on.
	/// </summary>
	static ConstValue Map( ConstValue a, int width, Func<double, double> f )
	{
		_ = width;

		return new ConstValue( f( a.X ), f( a.Y ), f( a.Z ), f( a.W ) );
	}

	static ConstValue Map2( ConstValue a, ConstValue b, int width, Func<double, double, double> f )
	{
		_ = width;

		return new ConstValue( f( a.X, b.X ), f( a.Y, b.Y ), f( a.Z, b.Z ), f( a.W, b.W ) );
	}

	static ConstValue Map3( ConstValue a, ConstValue b, ConstValue c, int width,
		Func<double, double, double, double> f )
	{
		_ = width;

		return new ConstValue( f( a.X, b.X, c.X ), f( a.Y, b.Y, c.Y ), f( a.Z, b.Z, c.Z ), f( a.W, b.W, c.W ) );
	}
}