1  using Ryujinx.Graphics.Shader.IntermediateRepresentation;
  2  using Ryujinx.Graphics.Shader.Translation;
  3  using System.Collections.Generic;
  4  using System.Linq;
  5  
  6  using static Ryujinx.Graphics.Shader.StructuredIr.AstHelper;
  7  
  8  namespace Ryujinx.Graphics.Shader.StructuredIr
  9  {
 10      static class AstOptimizer
 11      {
 12          public static void Optimize(StructuredProgramContext context)
 13          {
 14              AstBlock mainBlock = context.CurrentFunction.MainBlock;
 15  
 16              // When debug mode is enabled, we disable expression propagation
 17              // (this makes comparison with the disassembly easier).
 18              if (!context.DebugMode)
 19              {
 20                  AstBlockVisitor visitor = new(mainBlock);
 21  
 22                  foreach (IAstNode node in visitor.Visit())
 23                  {
 24                      if (node is AstAssignment assignment && assignment.Destination is AstOperand propVar)
 25                      {
 26                          bool isWorthPropagating = propVar.Uses.Count == 1 || IsWorthPropagating(assignment.Source);
 27  
 28                          if (propVar.Defs.Count == 1 && isWorthPropagating)
 29                          {
 30                              PropagateExpression(propVar, assignment.Source);
 31                          }
 32  
 33                          if (propVar.Type == OperandType.LocalVariable && propVar.Uses.Count == 0)
 34                          {
 35                              visitor.Block.Remove(assignment);
 36  
 37                              context.CurrentFunction.Locals.Remove(propVar);
 38                          }
 39                      }
 40                  }
 41              }
 42  
 43              RemoveEmptyBlocks(mainBlock);
 44          }
 45  
 46          private static bool IsWorthPropagating(IAstNode source)
 47          {
 48              if (source is not AstOperation srcOp)
 49              {
 50                  return false;
 51              }
 52  
 53              if (!InstructionInfo.IsUnary(srcOp.Inst))
 54              {
 55                  return false;
 56              }
 57  
 58              return srcOp.GetSource(0) is AstOperand || srcOp.Inst == Instruction.Copy;
 59          }
 60  
 61          private static void PropagateExpression(AstOperand propVar, IAstNode source)
 62          {
 63              IAstNode[] uses = propVar.Uses.ToArray();
 64  
 65              foreach (IAstNode useNode in uses)
 66              {
 67                  if (useNode is AstBlock useBlock)
 68                  {
 69                      useBlock.Condition = source;
 70                  }
 71                  else if (useNode is AstOperation useOperation)
 72                  {
 73                      for (int srcIndex = 0; srcIndex < useOperation.SourcesCount; srcIndex++)
 74                      {
 75                          if (useOperation.GetSource(srcIndex) == propVar)
 76                          {
 77                              useOperation.SetSource(srcIndex, source);
 78                          }
 79                      }
 80                  }
 81                  else if (useNode is AstAssignment useAssignment)
 82                  {
 83                      useAssignment.Source = source;
 84                  }
 85              }
 86          }
 87  
 88          private static void RemoveEmptyBlocks(AstBlock mainBlock)
 89          {
 90              Queue<AstBlock> pending = new();
 91  
 92              pending.Enqueue(mainBlock);
 93  
 94              while (pending.TryDequeue(out AstBlock block))
 95              {
 96                  foreach (IAstNode node in block)
 97                  {
 98                      if (node is AstBlock childBlock)
 99                      {
100                          pending.Enqueue(childBlock);
101                      }
102                  }
103  
104                  AstBlock parent = block.Parent;
105  
106                  if (parent == null)
107                  {
108                      continue;
109                  }
110  
111                  AstBlock nextBlock = Next(block) as AstBlock;
112  
113                  bool hasElse = nextBlock != null && nextBlock.Type == AstBlockType.Else;
114  
115                  bool isIf = block.Type == AstBlockType.If;
116  
117                  if (block.Count == 0)
118                  {
119                      if (isIf)
120                      {
121                          if (hasElse)
122                          {
123                              nextBlock.TurnIntoIf(InverseCond(block.Condition));
124                          }
125  
126                          parent.Remove(block);
127                      }
128                      else if (block.Type == AstBlockType.Else)
129                      {
130                          parent.Remove(block);
131                      }
132                  }
133                  else if (isIf && parent.Type == AstBlockType.Else && parent.Count == (hasElse ? 2 : 1))
134                  {
135                      AstBlock parentOfParent = parent.Parent;
136  
137                      parent.Remove(block);
138  
139                      parentOfParent.AddAfter(parent, block);
140  
141                      if (hasElse)
142                      {
143                          parent.Remove(nextBlock);
144  
145                          parentOfParent.AddAfter(block, nextBlock);
146                      }
147  
148                      parentOfParent.Remove(parent);
149  
150                      block.TurnIntoElseIf();
151                  }
152              }
153          }
154      }
155  }