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

9 files changed, 165 insertions, 47 deletions

diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
index 50b9bc9f..85347dfd 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs
@@ -10,9 +10,13 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
     {
         public const string Tab = "    ";
 
+        private readonly StructuredProgramInfo _info;
+
+        public StructuredFunction CurrentFunction { get; set; }
+
         public ShaderConfig Config { get; }
 
-        public bool CbIndexable { get; }
+        public bool CbIndexable => _info.UsesCbIndexing;
 
         public List<BufferDescriptor>  CBufferDescriptors { get; }
         public List<BufferDescriptor>  SBufferDescriptors { get; }
@@ -27,10 +31,10 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         private string _indentation;
 
-        public CodeGenContext(ShaderConfig config, bool cbIndexable)
+        public CodeGenContext(StructuredProgramInfo info, ShaderConfig config)
         {
+            _info = info;
             Config = config;
-            CbIndexable = cbIndexable;
 
             CBufferDescriptors = new List<BufferDescriptor>();
             SBufferDescriptors = new List<BufferDescriptor>();
@@ -95,6 +99,11 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 descriptor.CbufOffset == cBufOffset);
         }
 
+        public StructuredFunction GetFunction(int id)
+        {
+            return _info.Functions[id];
+        }
+
         private void UpdateIndentation()
         {
             _indentation = GetIndentation(_level);
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
index 08279839..6f5e75aa 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
@@ -187,9 +187,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             }
         }
 
-        public static void DeclareLocals(CodeGenContext context, StructuredProgramInfo info)
+        public static void DeclareLocals(CodeGenContext context, StructuredFunction function)
         {
-            foreach (AstOperand decl in info.Locals)
+            foreach (AstOperand decl in function.Locals)
             {
                 string name = context.OperandManager.DeclareLocal(decl);
 
@@ -197,13 +197,14 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             }
         }
 
-        private static string GetVarTypeName(VariableType type)
+        public static string GetVarTypeName(VariableType type)
         {
             switch (type)
             {
                 case VariableType.Bool: return "bool";
                 case VariableType.F32:  return "precise float";
                 case VariableType.F64:  return "double";
+                case VariableType.None: return "void";
                 case VariableType.S32:  return "int";
                 case VariableType.U32:  return "uint";
             }
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
index d1cf4636..cd9ca96e 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
@@ -22,6 +22,8 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
         public const string LocalMemoryName  = "local_mem";
         public const string SharedMemoryName = "shared_mem";
 
+        public const string ArgumentNamePrefix = "a";
+
         public const string UndefinedName = "undef";
 
         public const string IsBgraName = "is_bgra";
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
index 00a32262..276544fc 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs
@@ -10,13 +10,32 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 {
     static class GlslGenerator
     {
+        private const string MainFunctionName = "main";
+
         public static GlslProgram Generate(StructuredProgramInfo info, ShaderConfig config)
         {
-            CodeGenContext context = new CodeGenContext(config, info.UsesCbIndexing);
+            CodeGenContext context = new CodeGenContext(info, config);
 
             Declarations.Declare(context, info);
 
-            PrintMainBlock(context, info);
+            if (info.Functions.Count != 0)
+            {
+                for (int i = 1; i < info.Functions.Count; i++)
+                {
+                    context.AppendLine($"{GetFunctionSignature(info.Functions[i])};");
+                }
+
+                context.AppendLine();
+
+                for (int i = 1; i < info.Functions.Count; i++)
+                {
+                    PrintFunction(context, info, info.Functions[i]);
+
+                    context.AppendLine();
+                }
+            }
+
+            PrintFunction(context, info, info.Functions[0], MainFunctionName);
 
             return new GlslProgram(
                 context.CBufferDescriptors.ToArray(),
@@ -26,55 +45,78 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 context.GetCode());
         }
 
-        private static void PrintMainBlock(CodeGenContext context, StructuredProgramInfo info)
+        private static void PrintFunction(CodeGenContext context, StructuredProgramInfo info, StructuredFunction function, string funcName = null)
         {
-            context.AppendLine("void main()");
+            context.CurrentFunction = function;
 
+            context.AppendLine(GetFunctionSignature(function, funcName));
             context.EnterScope();
 
-            Declarations.DeclareLocals(context, info);
-
-            // Some games will leave some elements of gl_Position uninitialized,
-            // in those cases, the elements will contain undefined values according
-            // to the spec, but on NVIDIA they seems to be always initialized to (0, 0, 0, 1),
-            // so we do explicit initialization to avoid UB on non-NVIDIA gpus.
-            if (context.Config.Stage == ShaderStage.Vertex)
-            {
-                context.AppendLine("gl_Position = vec4(0.0, 0.0, 0.0, 1.0);");
-            }
+            Declarations.DeclareLocals(context, function);
 
-            // Ensure that unused attributes are set, otherwise the downstream
-            // compiler may eliminate them.
-            // (Not needed for fragment shader as it is the last stage).
-            if (context.Config.Stage != ShaderStage.Compute &&
-                context.Config.Stage != ShaderStage.Fragment)
+            if (funcName == MainFunctionName)
             {
-                for (int attr = 0; attr < Declarations.MaxAttributes; attr++)
+                // Some games will leave some elements of gl_Position uninitialized,
+                // in those cases, the elements will contain undefined values according
+                // to the spec, but on NVIDIA they seems to be always initialized to (0, 0, 0, 1),
+                // so we do explicit initialization to avoid UB on non-NVIDIA gpus.
+                if (context.Config.Stage == ShaderStage.Vertex)
                 {
-                    if (info.OAttributes.Contains(attr))
-                    {
-                        continue;
-                    }
+                    context.AppendLine("gl_Position = vec4(0.0, 0.0, 0.0, 1.0);");
+                }
 
-                    if ((context.Config.Flags & TranslationFlags.Feedback) != 0)
-                    {
-                        context.AppendLine($"{DefaultNames.OAttributePrefix}{attr}_x = 0;");
-                        context.AppendLine($"{DefaultNames.OAttributePrefix}{attr}_y = 0;");
-                        context.AppendLine($"{DefaultNames.OAttributePrefix}{attr}_z = 0;");
-                        context.AppendLine($"{DefaultNames.OAttributePrefix}{attr}_w = 0;");
-                    }
-                    else
+                // Ensure that unused attributes are set, otherwise the downstream
+                // compiler may eliminate them.
+                // (Not needed for fragment shader as it is the last stage).
+                if (context.Config.Stage != ShaderStage.Compute &&
+                    context.Config.Stage != ShaderStage.Fragment)
+                {
+                    for (int attr = 0; attr < Declarations.MaxAttributes; attr++)
                     {
-                        context.AppendLine($"{DefaultNames.OAttributePrefix}{attr} = vec4(0);");
+                        if (info.OAttributes.Contains(attr))
+                        {
+                            continue;
+                        }
+
+                        if ((context.Config.Flags & TranslationFlags.Feedback) != 0)
+                        {
+                            context.AppendLine($"{DefaultNames.OAttributePrefix}{attr}_x = 0;");
+                            context.AppendLine($"{DefaultNames.OAttributePrefix}{attr}_y = 0;");
+                            context.AppendLine($"{DefaultNames.OAttributePrefix}{attr}_z = 0;");
+                            context.AppendLine($"{DefaultNames.OAttributePrefix}{attr}_w = 0;");
+                        }
+                        else
+                        {
+                            context.AppendLine($"{DefaultNames.OAttributePrefix}{attr} = vec4(0);");
+                        }
                     }
                 }
             }
 
-            PrintBlock(context, info.MainBlock);
+            PrintBlock(context, function.MainBlock);
 
             context.LeaveScope();
         }
 
+        private static string GetFunctionSignature(StructuredFunction function, string funcName = null)
+        {
+            string[] args = new string[function.InArguments.Length + function.OutArguments.Length];
+
+            for (int i = 0; i < function.InArguments.Length; i++)
+            {
+                args[i] = $"{Declarations.GetVarTypeName(function.InArguments[i])} {OperandManager.GetArgumentName(i)}";
+            }
+
+            for (int i = 0; i < function.OutArguments.Length; i++)
+            {
+                int j = i + function.InArguments.Length;
+
+                args[j] = $"out {Declarations.GetVarTypeName(function.OutArguments[i])} {OperandManager.GetArgumentName(j)}";
+            }
+
+            return $"{Declarations.GetVarTypeName(function.ReturnType)} {funcName ?? function.Name}({string.Join(", ", args)})";
+        }
+
         private static void PrintBlock(CodeGenContext context, AstBlock block)
         {
             AstBlockVisitor visitor = new AstBlockVisitor(block);
@@ -123,8 +165,8 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 }
                 else if (node is AstAssignment assignment)
                 {
-                    VariableType srcType = OperandManager.GetNodeDestType(assignment.Source);
-                    VariableType dstType = OperandManager.GetNodeDestType(assignment.Destination);
+                    VariableType srcType = OperandManager.GetNodeDestType(context, assignment.Source);
+                    VariableType dstType = OperandManager.GetNodeDestType(context, assignment.Destination);
 
                     string dest;
 
@@ -154,7 +196,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         private static string GetCondExpr(CodeGenContext context, IAstNode cond)
         {
-            VariableType srcType = OperandManager.GetNodeDestType(cond);
+            VariableType srcType = OperandManager.GetNodeDestType(context, cond);
 
             return ReinterpretCast(context, cond, srcType, VariableType.Bool);
         }
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
index f1c741e6..388f0c25 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
@@ -2,6 +2,7 @@ using Ryujinx.Graphics.Shader.IntermediateRepresentation;
 using Ryujinx.Graphics.Shader.StructuredIr;
 using System;
 
+using static Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions.InstGenCall;
 using static Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions.InstGenHelper;
 using static Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions.InstGenMemory;
 using static Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions.InstGenPacking;
@@ -82,6 +83,12 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             {
                 string op = info.OpName;
 
+                // Return may optionally have a return value (and in this case it is unary).
+                if (inst == Instruction.Return && operation.SourcesCount != 0)
+                {
+                    return $"{op} {GetSoureExpr(context, operation.GetSource(0), context.CurrentFunction.ReturnType)}";
+                }
+
                 int arity = (int)(info.Type & InstType.ArityMask);
 
                 string[] expr = new string[arity];
@@ -116,6 +123,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             {
                 switch (inst)
                 {
+                    case Instruction.Call:
+                        return Call(context, operation);
+
                     case Instruction.ImageLoad:
                         return ImageLoadOrStore(context, operation);
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenCall.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenCall.cs
new file mode 100644
index 00000000..2df6960d
--- /dev/null
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenCall.cs
@@ -0,0 +1,29 @@
+using Ryujinx.Graphics.Shader.IntermediateRepresentation;
+using Ryujinx.Graphics.Shader.StructuredIr;
+using System.Diagnostics;
+
+using static Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions.InstGenHelper;
+
+namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
+{
+    static class InstGenCall
+    {
+        public static string Call(CodeGenContext context, AstOperation operation)
+        {
+            AstOperand funcId = (AstOperand)operation.GetSource(0);
+
+            Debug.Assert(funcId.Type == OperandType.Constant);
+
+            var function = context.GetFunction(funcId.Value);
+
+            string[] args = new string[operation.SourcesCount - 1];
+
+            for (int i = 0; i < args.Length; i++)
+            {
+                args[i] = GetSoureExpr(context, operation.GetSource(i + 1), function.GetArgumentType(i));
+            }
+
+            return $"{function.Name}({string.Join(", ", args)})";
+        }
+    }
+}
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
index 15f9b666..1b1efe9d 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
@@ -36,6 +36,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             Add(Instruction.BitwiseExclusiveOr,       InstType.OpBinaryCom,    "^",               7);
             Add(Instruction.BitwiseNot,               InstType.OpUnary,        "~",               0);
             Add(Instruction.BitwiseOr,                InstType.OpBinaryCom,    "|",               8);
+            Add(Instruction.Call,                     InstType.Special);
             Add(Instruction.Ceiling,                  InstType.CallUnary,      "ceil");
             Add(Instruction.Clamp,                    InstType.CallTernary,    "clamp");
             Add(Instruction.ClampU32,                 InstType.CallTernary,    "clamp");
@@ -135,7 +136,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 
         public static string GetSoureExpr(CodeGenContext context, IAstNode node, VariableType dstType)
         {
-            return ReinterpretCast(context, node, OperandManager.GetNodeDestType(node), dstType);
+            return ReinterpretCast(context, node, OperandManager.GetNodeDestType(context, node), dstType);
         }
 
         public static string Enclose(string expr, IAstNode node, Instruction pInst, bool isLhs)
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
index cb339f05..7fdca138 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
@@ -226,7 +226,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 
             string offsetExpr = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
 
-            VariableType srcType = OperandManager.GetNodeDestType(src2);
+            VariableType srcType = OperandManager.GetNodeDestType(context, src2);
 
             string src = TypeConversion.ReinterpretCast(context, src2, srcType, VariableType.U32);
 
@@ -242,7 +242,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             string indexExpr  = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
             string offsetExpr = GetSoureExpr(context, src2, GetSrcVarType(operation.Inst, 1));
 
-            VariableType srcType = OperandManager.GetNodeDestType(src3);
+            VariableType srcType = OperandManager.GetNodeDestType(context, src3);
 
             string src = TypeConversion.ReinterpretCast(context, src3, srcType, VariableType.U32);
 
diff --git a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
index 459b60c4..14ea7032 100644
--- a/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
+++ b/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
@@ -3,6 +3,7 @@ using Ryujinx.Graphics.Shader.StructuredIr;
 using Ryujinx.Graphics.Shader.Translation;
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 
 using static Ryujinx.Graphics.Shader.StructuredIr.InstructionInfo;
 
@@ -96,6 +97,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
         {
             switch (operand.Type)
             {
+                case OperandType.Argument:
+                    return GetArgumentName(operand.Value);
+
                 case OperandType.Attribute:
                     return GetAttributeName(operand, config);
 
@@ -287,7 +291,12 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return "xyzw"[value];
         }
 
-        public static VariableType GetNodeDestType(IAstNode node)
+        public static string GetArgumentName(int argIndex)
+        {
+            return $"{DefaultNames.ArgumentNamePrefix}{argIndex}";
+        }
+
+        public static VariableType GetNodeDestType(CodeGenContext context, IAstNode node)
         {
             if (node is AstOperation operation)
             {
@@ -298,6 +307,14 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 {
                     return GetOperandVarType((AstOperand)operation.GetSource(0));
                 }
+                else if (operation.Inst == Instruction.Call)
+                {
+                    AstOperand funcId = (AstOperand)operation.GetSource(0);
+
+                    Debug.Assert(funcId.Type == OperandType.Constant);
+
+                    return context.GetFunction(funcId.Value).ReturnType;
+                }
                 else if (operation is AstTextureOperation texOp &&
                          (texOp.Inst == Instruction.ImageLoad ||
                           texOp.Inst == Instruction.ImageStore))
@@ -309,6 +326,13 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             }
             else if (node is AstOperand operand)
             {
+                if (operand.Type == OperandType.Argument)
+                {
+                    int argIndex = operand.Value;
+
+                    return context.CurrentFunction.GetArgumentType(argIndex);
+                }
+
                 return GetOperandVarType(operand);
             }
             else