1 files changed, 156 insertions, 14 deletions

diff --git a/Ryujinx.Graphics/Gal/Shader/ShaderDecoder.cs b/Ryujinx.Graphics/Gal/Shader/ShaderDecoder.cs
index 024fa364..85522ff9 100644
--- a/Ryujinx.Graphics/Gal/Shader/ShaderDecoder.cs
+++ b/Ryujinx.Graphics/Gal/Shader/ShaderDecoder.cs
@@ -1,19 +1,133 @@
+using System.Collections.Generic;
+
 namespace Ryujinx.Graphics.Gal.Shader
 {
     static class ShaderDecoder
     {
         private const bool AddDbgComments = true;
 
-        public static ShaderIrBlock DecodeBasicBlock(IGalMemory Memory, long Position)
+        public static ShaderIrBlock[] Decode(IGalMemory Memory, long Start)
         {
-            ShaderIrBlock Block = new ShaderIrBlock();
+            Dictionary<long, ShaderIrBlock> Visited    = new Dictionary<long, ShaderIrBlock>();
+            Dictionary<long, ShaderIrBlock> VisitedEnd = new Dictionary<long, ShaderIrBlock>();
+
+            Queue<ShaderIrBlock> Blocks = new Queue<ShaderIrBlock>();
+
+            ShaderIrBlock Enqueue(long Position, ShaderIrBlock Source = null)
+            {
+                if (!Visited.TryGetValue(Position, out ShaderIrBlock Output))
+                {
+                    Output = new ShaderIrBlock(Position);
+
+                    Blocks.Enqueue(Output);
+
+                    Visited.Add(Position, Output);
+                }
+
+                if (Source != null)
+                {
+                    Output.Sources.Add(Source);
+                }
+
+                return Output;
+            }
+
+            ShaderIrBlock Entry = Enqueue(Start);
+
+            while (Blocks.Count > 0)
+            {
+                ShaderIrBlock Current = Blocks.Dequeue();
+
+                FillBlock(Memory, Current);
+
+                //Set child blocks. "Branch" is the block the branch instruction
+                //points to (when taken), "Next" is the block at the next address,
+                //executed when the branch is not taken. For Unconditional Branches
+                //or end of shader, Next is null.
+                if (Current.Nodes.Count > 0)
+                {
+                    ShaderIrNode LastNode = Current.GetLastNode();
+
+                    ShaderIrOp Op = GetInnermostOp(LastNode);
+
+                    if (Op?.Inst == ShaderIrInst.Bra)
+                    {
+                        int Offset = ((ShaderIrOperImm)Op.OperandA).Value;
+
+                        long Target = Current.EndPosition + Offset;
+
+                        Current.Branch = Enqueue(Target, Current);
+                    }
+
+                    if (NodeHasNext(LastNode))
+                    {
+                        Current.Next = Enqueue(Current.EndPosition);
+                    }
+                }
+
+                //If we have on the graph two blocks with the same end position,
+                //then we need to split the bigger block and have two small blocks,
+                //the end position of the bigger "Current" block should then be == to
+                //the position of the "Smaller" block.
+                while (VisitedEnd.TryGetValue(Current.EndPosition, out ShaderIrBlock Smaller))
+                {
+                    if (Current.Position > Smaller.Position)
+                    {
+                        ShaderIrBlock Temp = Smaller;
+
+                        Smaller = Current;
+                        Current = Temp;
+                    }
+
+                    Current.EndPosition = Smaller.Position;
+                    Current.Next        = Smaller;
+                    Current.Branch      = null;
+
+                    Current.Nodes.RemoveRange(
+                        Current.Nodes.Count - Smaller.Nodes.Count,
+                        Smaller.Nodes.Count);
+
+                    VisitedEnd[Smaller.EndPosition] = Smaller;
+                }
+
+                VisitedEnd.Add(Current.EndPosition, Current);
+            }
+
+            //Make and sort Graph blocks array by position.
+            ShaderIrBlock[] Graph = new ShaderIrBlock[Visited.Count];
 
-            while (true)
+            while (Visited.Count > 0)
             {
-                Block.Position = Position;
+                ulong FirstPos = ulong.MaxValue;
+
+                foreach (ShaderIrBlock Block in Visited.Values)
+                {
+                    if (FirstPos > (ulong)Block.Position)
+                        FirstPos = (ulong)Block.Position;
+                }
 
-                Block.MarkLabel(Position);
+                ShaderIrBlock Current = Visited[(long)FirstPos];
 
+                do
+                {
+                    Graph[Graph.Length - Visited.Count] = Current;
+
+                    Visited.Remove(Current.Position);
+
+                    Current = Current.Next;
+                }
+                while (Current != null);
+            }
+
+            return Graph;
+        }
+
+        private static void FillBlock(IGalMemory Memory, ShaderIrBlock Block)
+        {
+            long Position = Block.Position;
+
+            do
+            {
                 //Ignore scheduling instructions, which are written every 32 bytes.
                 if ((Position & 0x1f) == 0)
                 {
@@ -33,9 +147,20 @@ namespace Ryujinx.Graphics.Gal.Shader
 
                 if (AddDbgComments)
                 {
-                    string DbgOpCode = $"0x{Position:x16}: 0x{OpCode:x16} ";
+                    string DbgOpCode = $"0x{(Position - 8):x16}: 0x{OpCode:x16} ";
+
+                    DbgOpCode += (Decode?.Method.Name ?? "???");
 
-                    Block.AddNode(new ShaderIrCmnt(DbgOpCode + (Decode?.Method.Name ?? "???")));
+                    if (Decode == ShaderDecode.Bra)
+                    {
+                        int Offset = ((int)(OpCode >> 20) << 8) >> 8;
+
+                        long Target = Position + Offset;
+
+                        DbgOpCode += " (0x" + Target.ToString("x16") + ")";
+                    }
+
+                    Block.AddNode(new ShaderIrCmnt(DbgOpCode));
                 }
 
                 if (Decode == null)
@@ -44,19 +169,36 @@ namespace Ryujinx.Graphics.Gal.Shader
                 }
 
                 Decode(Block, OpCode);
+            }
+            while (!IsFlowChange(Block.GetLastNode()));
 
-                if (Block.GetLastNode() is ShaderIrOp Op && Op.Inst == ShaderIrInst.Exit)
-                {
-                    break;
-                }
+            Block.EndPosition = Position;
+        }
+
+        private static bool IsFlowChange(ShaderIrNode Node)
+        {
+            return !NodeHasNext(GetInnermostOp(Node));
+        }
+
+        private static ShaderIrOp GetInnermostOp(ShaderIrNode Node)
+        {
+            if (Node is ShaderIrCond Cond)
+            {
+                Node = Cond.Child;
             }
 
-            return Block;
+            return Node is ShaderIrOp Op ? Op : null;
         }
 
-        private static bool IsFlowChange(ShaderIrInst Inst)
+        private static bool NodeHasNext(ShaderIrNode Node)
         {
-            return Inst == ShaderIrInst.Exit;
+            if (!(Node is ShaderIrOp Op))
+            {
+                return true;
+            }
+
+            return Op.Inst != ShaderIrInst.Exit &&
+                   Op.Inst != ShaderIrInst.Bra;
         }
     }
 }
\ No newline at end of file