1 files changed, 73 insertions, 3 deletions

diff --git a/src/Ryujinx.Graphics.Shader/Translation/Optimizations/Utils.cs b/src/Ryujinx.Graphics.Shader/Translation/Optimizations/Utils.cs
index 74a6d5a1..23180ff8 100644
--- a/src/Ryujinx.Graphics.Shader/Translation/Optimizations/Utils.cs
+++ b/src/Ryujinx.Graphics.Shader/Translation/Optimizations/Utils.cs
@@ -34,6 +34,50 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
             return elemIndexSrc.Type == OperandType.Constant && elemIndexSrc.Value == elemIndex;
         }
 
+        private static bool IsSameOperand(Operand x, Operand y)
+        {
+            if (x.Type != y.Type || x.Value != y.Value)
+            {
+                return false;
+            }
+
+            // TODO: Handle Load operations with the same storage and the same constant parameters.
+            return x == y || x.Type == OperandType.Constant || x.Type == OperandType.ConstantBuffer;
+        }
+
+        private static bool AreAllSourcesEqual(INode node, INode otherNode)
+        {
+            if (node.SourcesCount != otherNode.SourcesCount)
+            {
+                return false;
+            }
+
+            for (int index = 0; index < node.SourcesCount; index++)
+            {
+                if (!IsSameOperand(node.GetSource(index), otherNode.GetSource(index)))
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }
+
+        public static bool AreAllSourcesTheSameOperand(INode node)
+        {
+            Operand firstSrc = node.GetSource(0);
+
+            for (int index = 1; index < node.SourcesCount; index++)
+            {
+                if (!IsSameOperand(firstSrc, node.GetSource(index)))
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }
+
         private static Operation FindBranchSource(BasicBlock block)
         {
             foreach (BasicBlock sourceBlock in block.Predecessors)
@@ -55,6 +99,19 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
             return inst == Instruction.BranchIfFalse || inst == Instruction.BranchIfTrue;
         }
 
+        private static bool IsSameCondition(Operand currentCondition, Operand queryCondition)
+        {
+            if (currentCondition == queryCondition)
+            {
+                return true;
+            }
+
+            return currentCondition.AsgOp is Operation currentOperation &&
+                queryCondition.AsgOp is Operation queryOperation &&
+                currentOperation.Inst == queryOperation.Inst &&
+                AreAllSourcesEqual(currentOperation, queryOperation);
+        }
+
         private static bool BlockConditionsMatch(BasicBlock currentBlock, BasicBlock queryBlock)
         {
             // Check if all the conditions for the query block are satisfied by the current block.
@@ -70,10 +127,10 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
 
             return currentBranch != null && queryBranch != null &&
                    currentBranch.Inst == queryBranch.Inst &&
-                   currentCondition == queryCondition;
+                   IsSameCondition(currentCondition, queryCondition);
         }
 
-        public static Operand FindLastOperation(Operand source, BasicBlock block)
+        public static Operand FindLastOperation(Operand source, BasicBlock block, bool recurse = true)
         {
             if (source.AsgOp is PhiNode phiNode)
             {
@@ -84,10 +141,23 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
                 for (int i = phiNode.SourcesCount - 1; i >= 0; i--)
                 {
                     BasicBlock phiBlock = phiNode.GetBlock(i);
+                    Operand phiSource = phiNode.GetSource(i);
 
                     if (BlockConditionsMatch(block, phiBlock))
                     {
-                        return phiNode.GetSource(i);
+                        return phiSource;
+                    }
+                    else if (recurse && phiSource.AsgOp is PhiNode)
+                    {
+                        // Phi source is another phi.
+                        // Let's check if that phi has a block that matches our condition.
+
+                        Operand match = FindLastOperation(phiSource, block, false);
+
+                        if (match != phiSource)
+                        {
+                            return match;
+                        }
                     }
                 }
             }