Implement CAL and RET shader instructions (#1618)

* Add support for CAL and RET shader instructions

* Remove unused stuff

* Fix a bug that could cause the wrong values to be passed to a function

* Avoid repopulating function id dictionary every time

* PR feedback

* Fix vertex shader A/B merge

8 files changed, 171 insertions, 48 deletions

diff --git a/Ryujinx.Graphics.Shader/StructuredIr/AstOperation.cs b/Ryujinx.Graphics.Shader/StructuredIr/AstOperation.cs
index 76eee71e..a8474955 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/AstOperation.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/AstOperation.cs
@@ -14,24 +14,35 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
 
         public int SourcesCount => _sources.Length;
 
-        public AstOperation(Instruction inst, params IAstNode[] sources)
+        public AstOperation(Instruction inst, IAstNode[] sources, int sourcesCount)
         {
             Inst     = inst;
             _sources = sources;
 
-            foreach (IAstNode source in sources)
+            for (int index = 0; index < sources.Length; index++)
             {
-                AddUse(source, this);
+                if (index < sourcesCount)
+                {
+                    AddUse(sources[index], this);
+                }
+                else
+                {
+                    AddDef(sources[index], this);
+                }
             }
 
             Index = 0;
         }
 
-        public AstOperation(Instruction inst, int index, params IAstNode[] sources) : this(inst, sources)
+        public AstOperation(Instruction inst, int index, IAstNode[] sources, int sourcesCount) : this(inst, sources, sourcesCount)
         {
             Index = index;
         }
 
+        public AstOperation(Instruction inst, params IAstNode[] sources) : this(inst, sources, sources.Length)
+        {
+        }
+
         public IAstNode GetSource(int index)
         {
             return _sources[index];
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/AstOptimizer.cs b/Ryujinx.Graphics.Shader/StructuredIr/AstOptimizer.cs
index a37e1a3e..4c6d17a0 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/AstOptimizer.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/AstOptimizer.cs
@@ -11,7 +11,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
     {
         public static void Optimize(StructuredProgramContext context)
         {
-            AstBlock mainBlock = context.Info.MainBlock;
+            AstBlock mainBlock = context.CurrentFunction.MainBlock;
 
             // When debug mode is enabled, we disable expression propagation
             // (this makes comparison with the disassembly easier).
@@ -34,7 +34,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
                         {
                             visitor.Block.Remove(assignment);
 
-                            context.Info.Locals.Remove(propVar);
+                            context.CurrentFunction.Locals.Remove(propVar);
                         }
                     }
                 }
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/AstTextureOperation.cs b/Ryujinx.Graphics.Shader/StructuredIr/AstTextureOperation.cs
index a3fa3e3a..188bf919 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/AstTextureOperation.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/AstTextureOperation.cs
@@ -19,7 +19,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
             int               handle,
             int               arraySize,
             int               index,
-            params IAstNode[] sources) : base(inst, index, sources)
+            params IAstNode[] sources) : base(inst, index, sources, sources.Length)
         {
             Type      = type;
             Format    = format;
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/InstructionInfo.cs b/Ryujinx.Graphics.Shader/StructuredIr/InstructionInfo.cs
index 3fcc5f11..fcf39cc0 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/InstructionInfo.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/InstructionInfo.cs
@@ -49,6 +49,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
             Add(Instruction.BitwiseOr,                VariableType.Int,    VariableType.Int,    VariableType.Int);
             Add(Instruction.BranchIfTrue,             VariableType.None,   VariableType.Bool);
             Add(Instruction.BranchIfFalse,            VariableType.None,   VariableType.Bool);
+            Add(Instruction.Call,                     VariableType.Scalar);
             Add(Instruction.Ceiling,                  VariableType.Scalar, VariableType.Scalar, VariableType.Scalar);
             Add(Instruction.Clamp,                    VariableType.Scalar, VariableType.Scalar, VariableType.Scalar, VariableType.Scalar);
             Add(Instruction.ClampU32,                 VariableType.U32,    VariableType.U32,    VariableType.U32,    VariableType.U32);
@@ -151,6 +152,10 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
             {
                 return VariableType.F32;
             }
+            else if (inst == Instruction.Call)
+            {
+                return VariableType.S32;
+            }
 
             return GetFinalVarType(_infoTbl[(int)(inst & Instruction.Mask)].SrcTypes[index], inst);
         }
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/StructuredFunction.cs b/Ryujinx.Graphics.Shader/StructuredIr/StructuredFunction.cs
new file mode 100644
index 00000000..3723f259
--- /dev/null
+++ b/Ryujinx.Graphics.Shader/StructuredIr/StructuredFunction.cs
@@ -0,0 +1,41 @@
+using System.Collections.Generic;
+
+namespace Ryujinx.Graphics.Shader.StructuredIr
+{
+    class StructuredFunction
+    {
+        public AstBlock MainBlock { get; }
+
+        public string Name { get; }
+
+        public VariableType ReturnType { get; }
+
+        public VariableType[] InArguments { get; }
+        public VariableType[] OutArguments { get; }
+
+        public HashSet<AstOperand> Locals { get; }
+
+        public StructuredFunction(
+            AstBlock mainBlock,
+            string name,
+            VariableType returnType,
+            VariableType[] inArguments,
+            VariableType[] outArguments)
+        {
+            MainBlock = mainBlock;
+            Name = name;
+            ReturnType = returnType;
+            InArguments = inArguments;
+            OutArguments = outArguments;
+
+            Locals = new HashSet<AstOperand>();
+        }
+
+        public VariableType GetArgumentType(int index)
+        {
+            return index >= InArguments.Length
+                ? OutArguments[index - InArguments.Length]
+                : InArguments[index];
+        }
+    }
+}
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs b/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs
index 65de5218..66570dc9 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs
@@ -8,51 +8,108 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
 {
     static class StructuredProgram
     {
-        public static StructuredProgramInfo MakeStructuredProgram(BasicBlock[] blocks, ShaderConfig config)
+        public static StructuredProgramInfo MakeStructuredProgram(Function[] functions, ShaderConfig config)
         {
-            PhiFunctions.Remove(blocks);
+            StructuredProgramContext context = new StructuredProgramContext(config);
 
-            StructuredProgramContext context = new StructuredProgramContext(blocks.Length, config);
-
-            for (int blkIndex = 0; blkIndex < blocks.Length; blkIndex++)
+            for (int funcIndex = 0; funcIndex < functions.Length; funcIndex++)
             {
-                BasicBlock block = blocks[blkIndex];
+                Function function = functions[funcIndex];
+
+                BasicBlock[] blocks = function.Blocks;
+
+                VariableType returnType = function.ReturnsValue ? VariableType.S32 : VariableType.None;
 
-                context.EnterBlock(block);
+                VariableType[] inArguments  = new VariableType[function.InArgumentsCount];
+                VariableType[] outArguments = new VariableType[function.OutArgumentsCount];
 
-                foreach (INode node in block.Operations)
+                for (int i = 0; i < inArguments.Length; i++)
                 {
-                    Operation operation = (Operation)node;
+                    inArguments[i] = VariableType.S32;
+                }
 
-                    if (IsBranchInst(operation.Inst))
-                    {
-                        context.LeaveBlock(block, operation);
-                    }
-                    else
+                for (int i = 0; i < outArguments.Length; i++)
+                {
+                    outArguments[i] = VariableType.S32;
+                }
+
+                context.EnterFunction(blocks.Length, function.Name, returnType, inArguments, outArguments);
+
+                PhiFunctions.Remove(blocks);
+
+                for (int blkIndex = 0; blkIndex < blocks.Length; blkIndex++)
+                {
+                    BasicBlock block = blocks[blkIndex];
+
+                    context.EnterBlock(block);
+
+                    for (LinkedListNode<INode> opNode = block.Operations.First; opNode != null; opNode = opNode.Next)
                     {
-                        AddOperation(context, operation);
+                        Operation operation = (Operation)opNode.Value;
+
+                        if (IsBranchInst(operation.Inst))
+                        {
+                            context.LeaveBlock(block, operation);
+                        }
+                        else if (operation.Inst != Instruction.CallOutArgument)
+                        {
+                            AddOperation(context, opNode);
+                        }
                     }
                 }
-            }
 
-            GotoElimination.Eliminate(context.GetGotos());
+                GotoElimination.Eliminate(context.GetGotos());
 
-            AstOptimizer.Optimize(context);
+                AstOptimizer.Optimize(context);
+
+                context.LeaveFunction();
+            }
 
             return context.Info;
         }
 
-        private static void AddOperation(StructuredProgramContext context, Operation operation)
+        private static void AddOperation(StructuredProgramContext context, LinkedListNode<INode> opNode)
         {
+            Operation operation = (Operation)opNode.Value;
+
             Instruction inst = operation.Inst;
 
-            IAstNode[] sources = new IAstNode[operation.SourcesCount];
+            bool isCall = inst == Instruction.Call;
 
-            for (int index = 0; index < sources.Length; index++)
+            int sourcesCount = operation.SourcesCount;
+
+            List<Operand> callOutOperands = new List<Operand>();
+
+            if (isCall)
+            {
+                LinkedListNode<INode> scan = opNode.Next;
+
+                while (scan != null && scan.Value is Operation nextOp && nextOp.Inst == Instruction.CallOutArgument)
+                {
+                    callOutOperands.Add(nextOp.Dest);
+                    scan = scan.Next;
+                }
+
+                sourcesCount += callOutOperands.Count;
+            }
+
+            IAstNode[] sources = new IAstNode[sourcesCount];
+
+            for (int index = 0; index < operation.SourcesCount; index++)
             {
                 sources[index] = context.GetOperandUse(operation.GetSource(index));
             }
 
+            if (isCall)
+            {
+                for (int index = 0; index < callOutOperands.Count; index++)
+                {
+                    sources[operation.SourcesCount + index] = context.GetOperandDef(callOutOperands[index]);
+                }
+
+                callOutOperands.Clear();
+            }
+
             AstTextureOperation GetAstTextureOperation(TextureOperation texOp)
             {
                 return new AstTextureOperation(
@@ -98,8 +155,6 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
                     AddSBufferUse(context.Info.SBuffers, operation);
                 }
 
-                AstAssignment assignment;
-
                 // If all the sources are bool, it's better to use short-circuiting
                 // logical operations, rather than forcing a cast to int and doing
                 // a bitwise operation with the value, as it is likely to be used as
@@ -152,16 +207,14 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
                 }
                 else if (!isCopy)
                 {
-                    source = new AstOperation(inst, operation.Index, sources);
+                    source = new AstOperation(inst, operation.Index, sources, operation.SourcesCount);
                 }
                 else
                 {
                     source = sources[0];
                 }
 
-                assignment = new AstAssignment(dest, source);
-
-                context.AddNode(assignment);
+                context.AddNode(new AstAssignment(dest, source));
             }
             else if (operation.Inst == Instruction.Comment)
             {
@@ -182,7 +235,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
                     AddSBufferUse(context.Info.SBuffers, operation);
                 }
 
-                context.AddNode(new AstOperation(inst, operation.Index, sources));
+                context.AddNode(new AstOperation(inst, operation.Index, sources, operation.SourcesCount));
             }
 
             // Those instructions needs to be emulated by using helper functions,
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramContext.cs b/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramContext.cs
index b7d5efbe..2667be1d 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramContext.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramContext.cs
@@ -24,11 +24,25 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
         private int _currEndIndex;
         private int _loopEndIndex;
 
+        public StructuredFunction CurrentFunction { get; private set; }
+
         public StructuredProgramInfo Info { get; }
 
         public ShaderConfig Config { get; }
 
-        public StructuredProgramContext(int blocksCount, ShaderConfig config)
+        public StructuredProgramContext(ShaderConfig config)
+        {
+            Info = new StructuredProgramInfo();
+
+            Config = config;
+        }
+
+        public void EnterFunction(
+            int blocksCount,
+            string name,
+            VariableType returnType,
+            VariableType[] inArguments,
+            VariableType[] outArguments)
         {
             _loopTails = new HashSet<BasicBlock>();
 
@@ -45,9 +59,12 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
             _currEndIndex = blocksCount;
             _loopEndIndex = blocksCount;
 
-            Info = new StructuredProgramInfo(_currBlock);
+            CurrentFunction = new StructuredFunction(_currBlock, name, returnType, inArguments, outArguments);
+        }
 
-            Config = config;
+        public void LeaveFunction()
+        {
+            Info.Functions.Add(CurrentFunction);
         }
 
         public void EnterBlock(BasicBlock block)
@@ -185,7 +202,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
             // so it is reset to false by the "local" assignment anyway.
             if (block.Index != 0)
             {
-                Info.MainBlock.AddFirst(Assign(gotoTempAsg.Destination, Const(IrConsts.False)));
+                CurrentFunction.MainBlock.AddFirst(Assign(gotoTempAsg.Destination, Const(IrConsts.False)));
             }
         }
 
@@ -253,7 +270,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
         {
             AstOperand newTemp = Local(type);
 
-            Info.Locals.Add(newTemp);
+            CurrentFunction.Locals.Add(newTemp);
 
             return newTemp;
         }
@@ -304,7 +321,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
 
                 _localsMap.Add(operand, astOperand);
 
-                Info.Locals.Add(astOperand);
+                CurrentFunction.Locals.Add(astOperand);
             }
 
             return astOperand;
diff --git a/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramInfo.cs b/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramInfo.cs
index ef3b3eca..16a27f51 100644
--- a/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramInfo.cs
+++ b/Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramInfo.cs
@@ -4,9 +4,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
 {
     class StructuredProgramInfo
     {
-        public AstBlock MainBlock { get; }
-
-        public HashSet<AstOperand> Locals { get; }
+        public List<StructuredFunction> Functions { get; }
 
         public HashSet<int> CBuffers { get; }
         public HashSet<int> SBuffers { get; }
@@ -22,11 +20,9 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
         public HashSet<AstTextureOperation> Samplers { get; }
         public HashSet<AstTextureOperation> Images   { get; }
 
-        public StructuredProgramInfo(AstBlock mainBlock)
+        public StructuredProgramInfo()
         {
-            MainBlock = mainBlock;
-
-            Locals = new HashSet<AstOperand>();
+            Functions = new List<StructuredFunction>();
 
             CBuffers = new HashSet<int>();
             SBuffers = new HashSet<int>();