Split HLSL program into statements

AndresTraks · Dec 15, 2024 · 02bff80 · 02bff80
1 parent 86e9b74
commit 02bff80
Show file tree

Hide file tree

Showing 17 changed files with 174 additions and 127 deletions.
diff --git a/Hlsl/Compiler/NodeGrouper.cs b/Hlsl/Compiler/NodeGrouper.cs
@@ -169,11 +169,6 @@ public bool CanGroupComponents(HlslTreeNode node1, HlslTreeNode node2, bool allo
                         CanGroupComponents(compare1.LessValue, compare2.LessValue) &&
                         CanGroupComponents(compare1.GreaterEqualValue, compare2.GreaterEqualValue);
                 }
-                else if (operation1 is ClipOperation clip1 &&
-                         operation2 is ClipOperation clip2)
-                {
-                    return CanGroupComponents(clip1.Value, clip2.Value);
-                }
                 else if (operation1 is LengthOperation length1 &&
                          operation2 is LengthOperation length2)
                 {

diff --git a/Hlsl/FlowControl/BreakStatement.cs b/Hlsl/FlowControl/BreakStatement.cs
@@ -0,0 +1,6 @@
+namespace HlslDecompiler.Hlsl.FlowControl
+{
+    public class BreakStatement : IStatement
+    {
+    }
+}
diff --git a/Hlsl/FlowControl/ClipStatement.cs b/Hlsl/FlowControl/ClipStatement.cs
@@ -0,0 +1,12 @@
+namespace HlslDecompiler.Hlsl.FlowControl
+{
+    public class ClipStatement : IStatement
+    {
+        public HlslTreeNode Value { get; }
+
+        public ClipStatement(HlslTreeNode value)
+        {
+            Value = value;
+        }
+    }
+}
diff --git a/Hlsl/FlowControl/IStatement.cs b/Hlsl/FlowControl/IStatement.cs
@@ -0,0 +1,6 @@
+namespace HlslDecompiler.Hlsl.FlowControl
+{
+    public interface IStatement
+    {
+    }
+}
diff --git a/Hlsl/FlowControl/StatementSequence.cs b/Hlsl/FlowControl/StatementSequence.cs
@@ -0,0 +1,65 @@
+using HlslDecompiler.DirectXShaderModel;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace HlslDecompiler.Hlsl.FlowControl
+{
+    public class StatementSequence : IStatement
+    {
+        public Dictionary<RegisterComponentKey, HlslTreeNode> Outputs { get; set; }
+
+        public Dictionary<RegisterKey, HlslTreeNode> GroupAssignments()
+        {
+            return GroupComponents(Outputs.Where(IsAssignment));
+        }
+
+        public Dictionary<RegisterKey, HlslTreeNode> GroupOutputs()
+        {
+            return GroupComponents(Outputs.Where(IsOutput));
+        }
+
+        private static Dictionary<RegisterKey, HlslTreeNode> GroupComponents(IEnumerable<KeyValuePair<RegisterComponentKey, HlslTreeNode>> outputsByComponent)
+        {
+            return outputsByComponent
+                .OrderBy(o => o.Key.ComponentIndex)
+                .GroupBy(o => o.Key.RegisterKey)
+                .ToDictionary(
+                    o => o.Key,
+                    o => (HlslTreeNode)new GroupNode(o.Select(o => o.Value).ToArray()));
+        }
+
+        private static bool IsAssignment(KeyValuePair<RegisterComponentKey, HlslTreeNode> operation)
+        {
+            if (operation.Key.RegisterKey is D3D9RegisterKey key9)
+            {
+                return key9.Type == RegisterType.Temp;
+            }
+            else if (operation.Key.RegisterKey is D3D10RegisterKey key10)
+            {
+                return key10.OperandType == OperandType.Temp;
+            }
+            else
+            {
+                throw new NotImplementedException();
+            }
+        }
+
+        private static bool IsOutput(KeyValuePair<RegisterComponentKey, HlslTreeNode> operation)
+        {
+            if (operation.Key.RegisterKey is D3D9RegisterKey key9)
+            {
+                RegisterType type = key9.Type;
+                return type == RegisterType.Output || type == RegisterType.ColorOut || type == RegisterType.DepthOut;
+            }
+            else if (operation.Key.RegisterKey is D3D10RegisterKey key10)
+            {
+                return key10.OperandType == OperandType.Output;
+            }
+            else
+            {
+                throw new NotImplementedException();
+            }
+        }
+    }
+}
diff --git a/Hlsl/HlslAst.cs b/Hlsl/HlslAst.cs
@@ -1,62 +1,17 @@
-using HlslDecompiler.DirectXShaderModel;
-using HlslDecompiler.Hlsl.TemplateMatch;
-using System;
+using HlslDecompiler.Hlsl.FlowControl;
 using System.Collections.Generic;
-using System.Linq;
 
 namespace HlslDecompiler.Hlsl
 {
     public class HlslAst
     {
-        public List<StatementSequence> Statements { get; private set; }
+        public List<IStatement> Statements { get; private set; }
         public RegisterState RegisterState { get; private set; }
 
-        public HlslAst(List<StatementSequence> statements, RegisterState registerState)
+        public HlslAst(List<IStatement> statements, RegisterState registerState)
         {
             Statements = statements;
             RegisterState = registerState;
         }
-
-        public List<Dictionary<RegisterKey, HlslTreeNode>> ReduceTree(NodeGrouper nodeGrouper)
-        {
-            var templateMatcher = new TemplateMatcher(nodeGrouper);
-            return Statements
-                .Select(GroupOutputs)
-                .Select(outputs => outputs.ToDictionary(r => r.Key, r => templateMatcher.Reduce(r.Value)))
-                .ToList();
-        }
-
-        public static Dictionary<RegisterKey, HlslTreeNode> GroupOutputs(StatementSequence statements)
-        {
-            IEnumerable<KeyValuePair<RegisterComponentKey, HlslTreeNode>> outputsByComponent =
-                statements.Outputs.Where(o =>
-            {
-                if (o.Value is ClipOperation)
-                {
-                    return true; 
-                }
-
-                if (o.Key.RegisterKey is D3D9RegisterKey key9)
-                {
-                    RegisterType type = key9.Type;
-                    return type == RegisterType.Output || type == RegisterType.ColorOut || type == RegisterType.DepthOut;
-                }
-                else if (o.Key.RegisterKey is D3D10RegisterKey key10)
-                {
-                    return key10.OperandType == OperandType.Output;
-                }
-                else
-                {
-                    throw new NotImplementedException();
-                }
-            });
-            var outputsByRegister = outputsByComponent
-                .OrderBy(o => o.Key.ComponentIndex)
-                .GroupBy(o => o.Key.RegisterKey)
-                .ToDictionary(
-                    o => o.Key,
-                    o => (HlslTreeNode)new GroupNode(o.Select(o => o.Value).ToArray()));
-            return outputsByRegister;
-        }
     }
 }
diff --git a/Hlsl/HlslAstWriter.cs b/Hlsl/HlslAstWriter.cs
@@ -1,4 +1,6 @@
 using HlslDecompiler.DirectXShaderModel;
+using HlslDecompiler.Hlsl.FlowControl;
+using HlslDecompiler.Hlsl.TemplateMatch;
 using System.Collections.Generic;
 using System.Linq;
 
@@ -26,34 +28,54 @@ protected override void WriteMethodBody()
         private void WriteAst(HlslAst ast)
         {
             var compiler = new NodeCompiler(_registers);
+            var nodeGrouper = new NodeGrouper(_registers);
+            var templateMatcher = new TemplateMatcher(nodeGrouper);
 
-            List<Dictionary<RegisterKey, HlslTreeNode>> sequenceRoots = ast.ReduceTree(new NodeGrouper(_registers));
-            foreach (var roots in sequenceRoots)
+            for (int i = 0; i < ast.Statements.Count; i++)
             {
-                if (roots.Count == 1)
+                IStatement statement = ast.Statements[i];
+                if (statement is StatementSequence sequence)
                 {
-                    if (roots.First().Value.Inputs.First() is ClipOperation)
+                    bool isLastStatement = i == ast.Statements.Count - 1;
+
+                    Dictionary<RegisterKey, HlslTreeNode> roots = isLastStatement
+                        ? sequence.GroupOutputs()
+                        : sequence.GroupAssignments();
+                    roots = roots.ToDictionary(r => r.Key, r => templateMatcher.Reduce(r.Value));
+
+                    if (isLastStatement)
                     {
-                        string statement = compiler.Compile(roots.First().Value);
-                        WriteLine($"{statement};");
+                        if (roots.Count == 1)
+                        {
+                            string compiled = compiler.Compile(roots.Single().Value);
+                            WriteLine($"return {compiled};");
+                        }
+                        else
+                        {
+                            foreach (var rootGroup in roots)
+                            {
+                                RegisterDeclaration outputRegister = _registers.MethodOutputRegisters[rootGroup.Key];
+                                string compiled = compiler.Compile(rootGroup.Value);
+                                WriteLine($"o.{outputRegister.Name} = {compiled};");
+                            }
+                            WriteLine();
+                            WriteLine($"return o;");
+                        }
                     }
                     else
                     {
-                        string statement = compiler.Compile(roots.Single().Value);
-                        WriteLine($"return {statement};");
+                        foreach (var rootGroup in roots)
+                        {
+                            string registerName = _registers.GetRegisterName(rootGroup.Key);
+                            string compiled = compiler.Compile(rootGroup.Value);
+                            WriteLine($"float4 {registerName} = {compiled};");
+                        }
                     }
                 }
-                else
+                else if (statement is ClipStatement clip)
                 {
-                    foreach (var rootGroup in roots)
-                    {
-                        RegisterDeclaration outputRegister = _registers.MethodOutputRegisters[rootGroup.Key];
-                        string statement = compiler.Compile(rootGroup.Value);
-                        WriteLine($"o.{outputRegister.Name} = {statement};");
-                    }
-
-                    WriteLine();
-                    WriteLine($"return o;");
+                    string compiled = compiler.Compile(clip.Value);
+                    WriteLine($"clip({compiled});");
                 }
             }
         }

diff --git a/Hlsl/InstructionParser.cs b/Hlsl/InstructionParser.cs
@@ -1,4 +1,5 @@
 using HlslDecompiler.DirectXShaderModel;
+using HlslDecompiler.Hlsl.FlowControl;
 using HlslDecompiler.Util;
 using System;
 using System.Collections.Generic;
@@ -12,7 +13,7 @@ class InstructionParser
     {
         private Dictionary<RegisterComponentKey, HlslTreeNode> _activeOutputs;
         private RegisterState _registerState;
-        private List<StatementSequence> _statementSequences;
+        private List<IStatement> _sequences;
         private StatementSequence _currentStatementSequence;
 
         public static HlslAst Parse(ShaderModel shader)
@@ -25,10 +26,10 @@ private HlslAst ParseToAst(ShaderModel shader)
         {
             _activeOutputs = new Dictionary<RegisterComponentKey, HlslTreeNode>();
             _registerState = new RegisterState(shader);
-            _statementSequences = new List<StatementSequence>();
+            _sequences = new List<IStatement>();
 
             _currentStatementSequence = new StatementSequence();
-            _statementSequences.Add(_currentStatementSequence);
+            _sequences.Add(_currentStatementSequence);
 
             LoadConstantOutputs(shader);
 
@@ -53,14 +54,22 @@ private HlslAst ParseToAst(ShaderModel shader)
 
             _currentStatementSequence.Outputs = new Dictionary<RegisterComponentKey, HlslTreeNode>(_activeOutputs);
 
-            return new HlslAst(_statementSequences, _registerState);
+            return new HlslAst(_sequences, _registerState);
         }
 
         private void ParseInstruction(D3D9Instruction instruction)
         {
             if (instruction.HasDestination)
             {
-                ParseAssignmentInstruction(instruction);
+                if (instruction.Opcode == Opcode.TexKill)
+                {
+                    InsertClip(instruction);
+                    EndStatementSequence(_activeOutputs);
+                }
+                else
+                {
+                    ParseAssignmentInstruction(instruction);
+                }
             }
             else
             {
@@ -101,14 +110,8 @@ private void ParseInstruction(D3D10Instruction instruction)
                 {
                     case D3D10Opcode.Discard:
                         {
-                            RegisterKey registerKey = instruction.GetParamRegisterKey(0);
-                            _currentStatementSequence.Outputs = new Dictionary<RegisterComponentKey, HlslTreeNode>();
-                            RegisterComponentKey registerComponentKey = new RegisterComponentKey(registerKey, 0);
-                            ClipOperation clip = new ClipOperation(new RegisterInputNode(registerComponentKey));
-                            _currentStatementSequence.Outputs.Add(registerComponentKey, clip);
-
-                            _currentStatementSequence = new StatementSequence();
-                            _statementSequences.Add(_currentStatementSequence);
+                            EndStatementSequence(_activeOutputs);
+                            InsertClip(instruction);
                             break;
                         }
                     case D3D10Opcode.DclTemps:
@@ -290,16 +293,9 @@ private void ParseAssignmentInstruction(D3D9Instruction instruction)
                 newOutputs[destinationKey] = instructionTree;
             }
 
-            if (instruction.Opcode == Opcode.TexKill)
-            {
-                EndStatementSequence(newOutputs);
-            }
-            else
+            foreach (var output in newOutputs)
             {
-                foreach (var output in newOutputs)
-                {
-                    _activeOutputs[output.Key] = output.Value;
-                }
+                _activeOutputs[output.Key] = output.Value;
             }
         }
 
@@ -308,7 +304,15 @@ private void EndStatementSequence(Dictionary<RegisterComponentKey, HlslTreeNode>
             _currentStatementSequence.Outputs = new Dictionary<RegisterComponentKey, HlslTreeNode>(newOutputs);
 
             _currentStatementSequence = new StatementSequence();
-            _statementSequences.Add(_currentStatementSequence);
+            _sequences.Add(_currentStatementSequence);
+        }
+
+        private void InsertClip(Instruction instruction)
+        {
+            RegisterKey registerKey = instruction.GetParamRegisterKey(0);
+            var registerComponentKey = new RegisterComponentKey(registerKey, 0);
+            var clip = new ClipStatement(new RegisterInputNode(registerComponentKey));
+            _sequences.Add(clip);
         }
 
         private void ParseAssignmentInstruction(D3D10Instruction instruction)
@@ -445,8 +449,6 @@ private HlslTreeNode CreateInstructionTree(D3D9Instruction instruction, Register
                                 return new SignGreaterOrEqualOperation(inputs[0], inputs[1]);
                             case Opcode.Slt:
                                 return new SignLessOperation(inputs[0], inputs[1]);
-                            case Opcode.TexKill:
-                                return new ClipOperation(inputs[0]);
                             default:
                                 throw new NotImplementedException();
                         }

diff --git a/Hlsl/Operations/ClipOperation.cs b/Hlsl/Operations/ClipOperation.cs
diff --git a/Hlsl/StatementSequence.cs b/Hlsl/StatementSequence.cs
diff --git a/HlslDecompiler.Tests/DecompileDxbcTests.cs b/HlslDecompiler.Tests/DecompileDxbcTests.cs
@@ -2,6 +2,7 @@
 using HlslDecompiler.DirectXShaderModel;
 using NUnit.Framework;
 using System.IO;
+using NUnit.Framework.Legacy;
 
 namespace HlslDecompiler.Tests
 {