1  using Ryujinx.Graphics.Shader.Translation;
  2  using System;
  3  using System.Collections.Generic;
  4  using System.Linq;
  5  using System.Runtime.CompilerServices;
  6  using static Ryujinx.Graphics.Shader.IntermediateRepresentation.OperandHelper;
  7  
  8  namespace Ryujinx.Graphics.Shader.Decoders
  9  {
 10      static class Decoder
 11      {
 12          private class Context
 13          {
 14              public AttributeUsage AttributeUsage { get; }
 15              public FeatureFlags UsedFeatures { get; private set; }
 16              public byte ClipDistancesWritten { get; private set; }
 17              public int Cb1DataSize { get; private set; }
 18  
 19              private readonly IGpuAccessor _gpuAccessor;
 20  
 21              public Context(IGpuAccessor gpuAccessor)
 22              {
 23                  _gpuAccessor = gpuAccessor;
 24                  AttributeUsage = new(gpuAccessor);
 25              }
 26  
 27              public uint ConstantBuffer1Read(int offset)
 28              {
 29                  if (Cb1DataSize < offset + 4)
 30                  {
 31                      Cb1DataSize = offset + 4;
 32                  }
 33  
 34                  return _gpuAccessor.ConstantBuffer1Read(offset);
 35              }
 36  
 37              public void SetUsedFeature(FeatureFlags flags)
 38              {
 39                  UsedFeatures |= flags;
 40              }
 41  
 42              public void SetClipDistanceWritten(int index)
 43              {
 44                  ClipDistancesWritten |= (byte)(1 << index);
 45              }
 46          }
 47  
 48          public static DecodedProgram Decode(ShaderDefinitions definitions, IGpuAccessor gpuAccessor, ulong startAddress)
 49          {
 50              Context context = new(gpuAccessor);
 51              Queue<DecodedFunction> functionsQueue = new();
 52              Dictionary<ulong, DecodedFunction> functionsVisited = new();
 53  
 54              DecodedFunction EnqueueFunction(ulong address)
 55              {
 56                  if (!functionsVisited.TryGetValue(address, out DecodedFunction function))
 57                  {
 58                      functionsVisited.Add(address, function = new DecodedFunction(address));
 59                      functionsQueue.Enqueue(function);
 60                  }
 61  
 62                  return function;
 63              }
 64  
 65              DecodedFunction mainFunction = EnqueueFunction(0);
 66  
 67              while (functionsQueue.TryDequeue(out DecodedFunction currentFunction))
 68              {
 69                  List<Block> blocks = new();
 70                  Queue<Block> workQueue = new();
 71                  Dictionary<ulong, Block> visited = new();
 72  
 73                  Block GetBlock(ulong blkAddress)
 74                  {
 75                      if (!visited.TryGetValue(blkAddress, out Block block))
 76                      {
 77                          block = new Block(blkAddress);
 78  
 79                          workQueue.Enqueue(block);
 80                          visited.Add(blkAddress, block);
 81                      }
 82  
 83                      return block;
 84                  }
 85  
 86                  GetBlock(currentFunction.Address);
 87  
 88                  bool hasNewTarget;
 89  
 90                  do
 91                  {
 92                      while (workQueue.TryDequeue(out Block currBlock))
 93                      {
 94                          // Check if the current block is inside another block.
 95                          if (BinarySearch(blocks, currBlock.Address, out int nBlkIndex))
 96                          {
 97                              Block nBlock = blocks[nBlkIndex];
 98  
 99                              if (nBlock.Address == currBlock.Address)
100                              {
101                                  throw new InvalidOperationException("Found duplicate block address on the list.");
102                              }
103  
104                              nBlock.Split(currBlock);
105                              blocks.Insert(nBlkIndex + 1, currBlock);
106  
107                              continue;
108                          }
109  
110                          // If we have a block after the current one, set the limit address.
111                          ulong limitAddress = ulong.MaxValue;
112  
113                          if (nBlkIndex != blocks.Count)
114                          {
115                              Block nBlock = blocks[nBlkIndex];
116  
117                              int nextIndex = nBlkIndex + 1;
118  
119                              if (nBlock.Address < currBlock.Address && nextIndex < blocks.Count)
120                              {
121                                  limitAddress = blocks[nextIndex].Address;
122                              }
123                              else if (nBlock.Address > currBlock.Address)
124                              {
125                                  limitAddress = blocks[nBlkIndex].Address;
126                              }
127                          }
128  
129                          FillBlock(definitions, gpuAccessor, context, currBlock, limitAddress, startAddress);
130  
131                          if (currBlock.OpCodes.Count != 0)
132                          {
133                              // We should have blocks for all possible branch targets,
134                              // including those from PBK/PCNT/SSY instructions.
135                              foreach (PushOpInfo pushOp in currBlock.PushOpCodes)
136                              {
137                                  GetBlock(pushOp.Op.GetAbsoluteAddress());
138                              }
139  
140                              // Set child blocks. "Branch" is the block the branch instruction
141                              // points to (when taken), "Next" is the block at the next address,
142                              // executed when the branch is not taken. For Unconditional Branches
143                              // or end of program, Next is null.
144                              InstOp lastOp = currBlock.GetLastOp();
145  
146                              if (lastOp.Name == InstName.Cal)
147                              {
148                                  EnqueueFunction(lastOp.GetAbsoluteAddress()).AddCaller(currentFunction);
149                              }
150                              else if (lastOp.Name == InstName.Bra)
151                              {
152                                  Block succBlock = GetBlock(lastOp.GetAbsoluteAddress());
153                                  currBlock.Successors.Add(succBlock);
154                                  succBlock.Predecessors.Add(currBlock);
155                              }
156  
157                              if (!IsUnconditionalBranch(ref lastOp))
158                              {
159                                  Block succBlock = GetBlock(currBlock.EndAddress);
160                                  currBlock.Successors.Insert(0, succBlock);
161                                  succBlock.Predecessors.Add(currBlock);
162                              }
163                          }
164  
165                          // Insert the new block on the list (sorted by address).
166                          if (blocks.Count != 0)
167                          {
168                              Block nBlock = blocks[nBlkIndex];
169  
170                              blocks.Insert(nBlkIndex + (nBlock.Address < currBlock.Address ? 1 : 0), currBlock);
171                          }
172                          else
173                          {
174                              blocks.Add(currBlock);
175                          }
176                      }
177  
178                      // Propagate SSY/PBK addresses into their uses (SYNC/BRK).
179                      foreach (Block block in blocks.Where(x => x.PushOpCodes.Count != 0))
180                      {
181                          for (int pushOpIndex = 0; pushOpIndex < block.PushOpCodes.Count; pushOpIndex++)
182                          {
183                              PropagatePushOp(visited, block, pushOpIndex);
184                          }
185                      }
186  
187                      // Try to find targets for BRX (indirect branch) instructions.
188                      hasNewTarget = FindBrxTargets(context, blocks, GetBlock);
189  
190                      // If we discovered new branch targets from the BRX instruction,
191                      // we need another round of decoding to decode the new blocks.
192                      // Additionally, we may have more SSY/PBK targets to propagate,
193                      // and new BRX instructions.
194                  }
195                  while (hasNewTarget);
196  
197                  currentFunction.SetBlocks(blocks.ToArray());
198              }
199  
200              return new DecodedProgram(
201                  mainFunction,
202                  functionsVisited,
203                  context.AttributeUsage,
204                  context.UsedFeatures,
205                  context.ClipDistancesWritten,
206                  context.Cb1DataSize);
207          }
208  
209          private static bool BinarySearch(List<Block> blocks, ulong address, out int index)
210          {
211              index = 0;
212  
213              int left = 0;
214              int right = blocks.Count - 1;
215  
216              while (left <= right)
217              {
218                  int size = right - left;
219  
220                  int middle = left + (size >> 1);
221  
222                  Block block = blocks[middle];
223  
224                  index = middle;
225  
226                  if (address >= block.Address && address < block.EndAddress)
227                  {
228                      return true;
229                  }
230  
231                  if (address < block.Address)
232                  {
233                      right = middle - 1;
234                  }
235                  else
236                  {
237                      left = middle + 1;
238                  }
239              }
240  
241              return false;
242          }
243  
244          private static void FillBlock(
245              ShaderDefinitions definitions,
246              IGpuAccessor gpuAccessor,
247              Context context,
248              Block block,
249              ulong limitAddress,
250              ulong startAddress)
251          {
252              ulong address = block.Address;
253              int bufferOffset = 0;
254              ReadOnlySpan<ulong> buffer = ReadOnlySpan<ulong>.Empty;
255  
256              InstOp op = default;
257  
258              do
259              {
260                  if (address + 7 >= limitAddress)
261                  {
262                      break;
263                  }
264  
265                  // Ignore scheduling instructions, which are written every 32 bytes.
266                  if ((address & 0x1f) == 0)
267                  {
268                      address += 8;
269                      bufferOffset++;
270                      continue;
271                  }
272  
273                  if (bufferOffset >= buffer.Length)
274                  {
275                      buffer = gpuAccessor.GetCode(startAddress + address, 8);
276                      bufferOffset = 0;
277                  }
278  
279                  ulong opCode = buffer[bufferOffset++];
280  
281                  op = InstTable.GetOp(address, opCode);
282  
283                  if (op.Props.HasFlag(InstProps.TexB))
284                  {
285                      context.SetUsedFeature(FeatureFlags.Bindless);
286                  }
287  
288                  switch (op.Name)
289                  {
290                      case InstName.Ald:
291                      case InstName.Ast:
292                      case InstName.Ipa:
293                          SetUserAttributeUses(definitions, context, op.Name, opCode);
294                          break;
295                      case InstName.Pbk:
296                      case InstName.Pcnt:
297                      case InstName.Ssy:
298                          block.AddPushOp(op);
299                          break;
300                      case InstName.Shfl:
301                          context.SetUsedFeature(FeatureFlags.Shuffle);
302                          break;
303                      case InstName.Ldl:
304                      case InstName.Stl:
305                          context.SetUsedFeature(FeatureFlags.LocalMemory);
306                          break;
307                      case InstName.Atoms:
308                      case InstName.AtomsCas:
309                      case InstName.Lds:
310                      case InstName.Sts:
311                          context.SetUsedFeature(FeatureFlags.SharedMemory);
312                          break;
313                      case InstName.Atom:
314                      case InstName.AtomCas:
315                      case InstName.Red:
316                      case InstName.Stg:
317                      case InstName.Suatom:
318                      case InstName.SuatomB:
319                      case InstName.SuatomB2:
320                      case InstName.SuatomCas:
321                      case InstName.SuatomCasB:
322                      case InstName.Sured:
323                      case InstName.SuredB:
324                      case InstName.Sust:
325                      case InstName.SustB:
326                      case InstName.SustD:
327                      case InstName.SustDB:
328                          context.SetUsedFeature(FeatureFlags.Store);
329                          break;
330                  }
331  
332                  block.OpCodes.Add(op);
333  
334                  address += 8;
335              }
336              while (!op.Props.HasFlag(InstProps.Bra));
337  
338              block.EndAddress = address;
339          }
340  
341          private static void SetUserAttributeUses(ShaderDefinitions definitions, Context context, InstName name, ulong opCode)
342          {
343              int offset;
344              int count = 1;
345              bool isStore = false;
346              bool indexed;
347              bool perPatch = false;
348  
349              if (name == InstName.Ast)
350              {
351                  InstAst opAst = new(opCode);
352                  count = (int)opAst.AlSize + 1;
353                  offset = opAst.Imm11;
354                  indexed = opAst.Phys;
355                  perPatch = opAst.P;
356                  isStore = true;
357              }
358              else if (name == InstName.Ald)
359              {
360                  InstAld opAld = new(opCode);
361                  count = (int)opAld.AlSize + 1;
362                  offset = opAld.Imm11;
363                  indexed = opAld.Phys;
364                  perPatch = opAld.P;
365                  isStore = opAld.O;
366              }
367              else /* if (name == InstName.Ipa) */
368              {
369                  InstIpa opIpa = new(opCode);
370                  offset = opIpa.Imm10;
371                  indexed = opIpa.Idx;
372              }
373  
374              if (indexed)
375              {
376                  if (isStore)
377                  {
378                      context.AttributeUsage.SetAllOutputUserAttributes();
379                      definitions.EnableOutputIndexing();
380                  }
381                  else
382                  {
383                      context.AttributeUsage.SetAllInputUserAttributes();
384                      definitions.EnableInputIndexing();
385                  }
386              }
387              else
388              {
389                  for (int elemIndex = 0; elemIndex < count; elemIndex++)
390                  {
391                      int attr = offset + elemIndex * 4;
392  
393                      if (perPatch)
394                      {
395                          if (attr >= AttributeConsts.UserAttributePerPatchBase && attr < AttributeConsts.UserAttributePerPatchEnd)
396                          {
397                              int userAttr = attr - AttributeConsts.UserAttributePerPatchBase;
398                              int index = userAttr / 16;
399  
400                              if (isStore)
401                              {
402                                  context.AttributeUsage.SetOutputUserAttributePerPatch(index);
403                              }
404                              else
405                              {
406                                  context.AttributeUsage.SetInputUserAttributePerPatch(index);
407                              }
408                          }
409                      }
410                      else if (attr >= AttributeConsts.UserAttributeBase && attr < AttributeConsts.UserAttributeEnd)
411                      {
412                          int userAttr = attr - AttributeConsts.UserAttributeBase;
413                          int index = userAttr / 16;
414  
415                          if (isStore)
416                          {
417                              context.AttributeUsage.SetOutputUserAttribute(index);
418                          }
419                          else
420                          {
421                              context.AttributeUsage.SetInputUserAttribute(index, (userAttr >> 2) & 3);
422                          }
423                      }
424  
425                      if (!isStore &&
426                          (attr == AttributeConsts.FogCoord ||
427                          (attr >= AttributeConsts.FrontColorDiffuseR && attr < AttributeConsts.ClipDistance0) ||
428                          (attr >= AttributeConsts.TexCoordBase && attr < AttributeConsts.TexCoordEnd)))
429                      {
430                          context.SetUsedFeature(FeatureFlags.FixedFuncAttr);
431                      }
432                      else
433                      {
434                          if (isStore)
435                          {
436                              switch (attr)
437                              {
438                                  case AttributeConsts.Layer:
439                                      if (definitions.Stage != ShaderStage.Compute && definitions.Stage != ShaderStage.Fragment)
440                                      {
441                                          context.SetUsedFeature(FeatureFlags.RtLayer);
442                                      }
443                                      break;
444                                  case AttributeConsts.ViewportIndex:
445                                      if (definitions.Stage != ShaderStage.Fragment)
446                                      {
447                                          context.SetUsedFeature(FeatureFlags.ViewportIndex);
448                                      }
449                                      break;
450                                  case AttributeConsts.ClipDistance0:
451                                  case AttributeConsts.ClipDistance1:
452                                  case AttributeConsts.ClipDistance2:
453                                  case AttributeConsts.ClipDistance3:
454                                  case AttributeConsts.ClipDistance4:
455                                  case AttributeConsts.ClipDistance5:
456                                  case AttributeConsts.ClipDistance6:
457                                  case AttributeConsts.ClipDistance7:
458                                      if (definitions.Stage.IsVtg())
459                                      {
460                                          context.SetClipDistanceWritten((attr - AttributeConsts.ClipDistance0) / 4);
461                                      }
462                                      break;
463                                  case AttributeConsts.ViewportMask:
464                                      if (definitions.Stage != ShaderStage.Fragment)
465                                      {
466                                          context.SetUsedFeature(FeatureFlags.ViewportMask);
467                                      }
468                                      break;
469                              }
470                          }
471                          else
472                          {
473                              switch (attr)
474                              {
475                                  case AttributeConsts.PositionX:
476                                  case AttributeConsts.PositionY:
477                                      if (definitions.Stage == ShaderStage.Fragment)
478                                      {
479                                          context.SetUsedFeature(FeatureFlags.FragCoordXY);
480                                      }
481                                      break;
482                                  case AttributeConsts.InstanceId:
483                                      if (definitions.Stage == ShaderStage.Vertex)
484                                      {
485                                          context.SetUsedFeature(FeatureFlags.InstanceId);
486                                      }
487                                      break;
488                              }
489                          }
490                      }
491                  }
492              }
493          }
494  
495          public static bool IsUnconditionalBranch(ref InstOp op)
496          {
497              return IsUnconditional(ref op) && op.Props.HasFlag(InstProps.Bra);
498          }
499  
500          private static bool IsUnconditional(ref InstOp op)
501          {
502              InstConditional condOp = new(op.RawOpCode);
503  
504              if ((op.Name == InstName.Bra || op.Name == InstName.Exit) && condOp.Ccc != Ccc.T)
505              {
506                  return false;
507              }
508  
509              return condOp.Pred == RegisterConsts.PredicateTrueIndex && !condOp.PredInv;
510          }
511  
512          private static bool FindBrxTargets(Context context, IEnumerable<Block> blocks, Func<ulong, Block> getBlock)
513          {
514              bool hasNewTarget = false;
515  
516              foreach (Block block in blocks)
517              {
518                  InstOp lastOp = block.GetLastOp();
519                  bool hasNext = block.HasNext();
520  
521                  if (lastOp.Name == InstName.Brx && block.Successors.Count == (hasNext ? 1 : 0))
522                  {
523                      HashSet<ulong> visited = new();
524  
525                      InstBrx opBrx = new(lastOp.RawOpCode);
526                      ulong baseOffset = lastOp.GetAbsoluteAddress();
527  
528                      // An indirect branch could go anywhere,
529                      // try to get the possible target offsets from the constant buffer.
530                      (int cbBaseOffset, int cbOffsetsCount) = FindBrxTargetRange(block, opBrx.SrcA);
531  
532                      if (cbOffsetsCount != 0)
533                      {
534                          hasNewTarget = true;
535                      }
536  
537                      for (int i = 0; i < cbOffsetsCount; i++)
538                      {
539                          uint targetOffset = context.ConstantBuffer1Read(cbBaseOffset + i * 4);
540                          ulong targetAddress = baseOffset + targetOffset;
541  
542                          if (visited.Add(targetAddress))
543                          {
544                              Block target = getBlock(targetAddress);
545                              target.Predecessors.Add(block);
546                              block.Successors.Add(target);
547                          }
548                      }
549                  }
550              }
551  
552              return hasNewTarget;
553          }
554  
555          private static (int, int) FindBrxTargetRange(Block block, int brxReg)
556          {
557              // Try to match the following pattern:
558              //
559              // IMNMX.U32 Rx, Rx, UpperBound, PT
560              // SHL Rx, Rx, 0x2
561              // LDC Rx, c[0x1][Rx+BaseOffset]
562              //
563              // Here, Rx is an arbitrary register, "UpperBound" and "BaseOffset" are constants.
564              // The above pattern is assumed to be generated by the compiler before BRX,
565              // as the instruction is usually used to implement jump tables for switch statement optimizations.
566              // On a successful match, "BaseOffset" is the offset in bytes where the jump offsets are
567              // located on the constant buffer, and "UpperBound" is the total number of offsets for the BRX, minus 1.
568  
569              HashSet<Block> visited = new();
570  
571              var ldcLocation = FindFirstRegWrite(visited, new BlockLocation(block, block.OpCodes.Count - 1), brxReg);
572              if (ldcLocation.Block == null || ldcLocation.Block.OpCodes[ldcLocation.Index].Name != InstName.Ldc)
573              {
574                  return (0, 0);
575              }
576  
577              GetOp<InstLdc>(ldcLocation, out var opLdc);
578  
579              if (opLdc.CbufSlot != 1 || opLdc.AddressMode != 0)
580              {
581                  return (0, 0);
582              }
583  
584              var shlLocation = FindFirstRegWrite(visited, ldcLocation, opLdc.SrcA);
585              if (shlLocation.Block == null || !shlLocation.IsImmInst(InstName.Shl))
586              {
587                  return (0, 0);
588              }
589  
590              GetOp<InstShlI>(shlLocation, out var opShl);
591  
592              if (opShl.Imm20 != 2)
593              {
594                  return (0, 0);
595              }
596  
597              var imnmxLocation = FindFirstRegWrite(visited, shlLocation, opShl.SrcA);
598              if (imnmxLocation.Block == null || !imnmxLocation.IsImmInst(InstName.Imnmx))
599              {
600                  return (0, 0);
601              }
602  
603              GetOp<InstImnmxI>(imnmxLocation, out var opImnmx);
604  
605              if (opImnmx.Signed || opImnmx.SrcPred != RegisterConsts.PredicateTrueIndex || opImnmx.SrcPredInv)
606              {
607                  return (0, 0);
608              }
609  
610              return (opLdc.CbufOffset, opImnmx.Imm20 + 1);
611          }
612  
613          private static void GetOp<T>(BlockLocation location, out T op) where T : unmanaged
614          {
615              ulong rawOp = location.Block.OpCodes[location.Index].RawOpCode;
616              op = Unsafe.As<ulong, T>(ref rawOp);
617          }
618  
619          private readonly struct BlockLocation
620          {
621              public Block Block { get; }
622              public int Index { get; }
623  
624              public BlockLocation(Block block, int index)
625              {
626                  Block = block;
627                  Index = index;
628              }
629  
630              public bool IsImmInst(InstName name)
631              {
632                  InstOp op = Block.OpCodes[Index];
633                  return op.Name == name && op.Props.HasFlag(InstProps.Ib);
634              }
635          }
636  
637          private static BlockLocation FindFirstRegWrite(HashSet<Block> visited, BlockLocation location, int regIndex)
638          {
639              Queue<BlockLocation> toVisit = new();
640              toVisit.Enqueue(location);
641              visited.Add(location.Block);
642  
643              while (toVisit.TryDequeue(out var currentLocation))
644              {
645                  Block block = currentLocation.Block;
646                  for (int i = currentLocation.Index - 1; i >= 0; i--)
647                  {
648                      if (WritesToRegister(block.OpCodes[i], regIndex))
649                      {
650                          return new BlockLocation(block, i);
651                      }
652                  }
653  
654                  foreach (Block predecessor in block.Predecessors)
655                  {
656                      if (visited.Add(predecessor))
657                      {
658                          toVisit.Enqueue(new BlockLocation(predecessor, predecessor.OpCodes.Count));
659                      }
660                  }
661              }
662  
663              return new BlockLocation(null, 0);
664          }
665  
666          private static bool WritesToRegister(InstOp op, int regIndex)
667          {
668              // Predicate instruction only ever writes to predicate, so we shouldn't check those.
669              if ((op.Props & (InstProps.Rd | InstProps.Rd2)) == 0)
670              {
671                  return false;
672              }
673  
674              if (op.Props.HasFlag(InstProps.Rd2) && (byte)(op.RawOpCode >> 28) == regIndex)
675              {
676                  return true;
677              }
678  
679              return (byte)op.RawOpCode == regIndex;
680          }
681  
682          private enum MergeType
683          {
684              Brk,
685              Cont,
686              Sync,
687          }
688  
689          private readonly struct PathBlockState
690          {
691              public Block Block { get; }
692  
693              private enum RestoreType
694              {
695                  None,
696                  PopPushOp,
697                  PushBranchOp,
698              }
699  
700              private readonly RestoreType _restoreType;
701  
702              private readonly ulong _restoreValue;
703              private readonly MergeType _restoreMergeType;
704  
705              public bool ReturningFromVisit => _restoreType != RestoreType.None;
706  
707              public PathBlockState(Block block)
708              {
709                  Block = block;
710                  _restoreType = RestoreType.None;
711                  _restoreValue = 0;
712                  _restoreMergeType = default;
713              }
714  
715              public PathBlockState(int oldStackSize)
716              {
717                  Block = null;
718                  _restoreType = RestoreType.PopPushOp;
719                  _restoreValue = (ulong)oldStackSize;
720                  _restoreMergeType = default;
721              }
722  
723              public PathBlockState(ulong syncAddress, MergeType mergeType)
724              {
725                  Block = null;
726                  _restoreType = RestoreType.PushBranchOp;
727                  _restoreValue = syncAddress;
728                  _restoreMergeType = mergeType;
729              }
730  
731              public void RestoreStackState(Stack<(ulong, MergeType)> branchStack)
732              {
733                  if (_restoreType == RestoreType.PushBranchOp)
734                  {
735                      branchStack.Push((_restoreValue, _restoreMergeType));
736                  }
737                  else if (_restoreType == RestoreType.PopPushOp)
738                  {
739                      while (branchStack.Count > (uint)_restoreValue)
740                      {
741                          branchStack.Pop();
742                      }
743                  }
744              }
745          }
746  
747          private static void PropagatePushOp(Dictionary<ulong, Block> blocks, Block currBlock, int pushOpIndex)
748          {
749              PushOpInfo pushOpInfo = currBlock.PushOpCodes[pushOpIndex];
750              InstOp pushOp = pushOpInfo.Op;
751  
752              Block target = blocks[pushOp.GetAbsoluteAddress()];
753  
754              Stack<PathBlockState> workQueue = new();
755              HashSet<Block> visited = new();
756              Stack<(ulong, MergeType)> branchStack = new();
757  
758              void Push(PathBlockState pbs)
759              {
760                  // When block is null, this means we are pushing a restore operation.
761                  // Restore operations are used to undo the work done inside a block
762                  // when we return from it, for example it pops addresses pushed by
763                  // SSY/PBK instructions inside the block, and pushes addresses poped
764                  // by SYNC/BRK.
765                  // For blocks, if it's already visited, we just ignore to avoid going
766                  // around in circles and getting stuck here.
767                  if (pbs.Block == null || !visited.Contains(pbs.Block))
768                  {
769                      workQueue.Push(pbs);
770                  }
771              }
772  
773              Push(new PathBlockState(currBlock));
774  
775              while (workQueue.TryPop(out PathBlockState pbs))
776              {
777                  if (pbs.ReturningFromVisit)
778                  {
779                      pbs.RestoreStackState(branchStack);
780  
781                      continue;
782                  }
783  
784                  Block current = pbs.Block;
785  
786                  // If the block was already processed, we just ignore it, otherwise
787                  // we would push the same child blocks of an already processed block,
788                  // and go around in circles until memory is exhausted.
789                  if (!visited.Add(current))
790                  {
791                      continue;
792                  }
793  
794                  int pushOpsCount = current.PushOpCodes.Count;
795                  if (pushOpsCount != 0)
796                  {
797                      Push(new PathBlockState(branchStack.Count));
798  
799                      for (int index = pushOpIndex; index < pushOpsCount; index++)
800                      {
801                          InstOp currentPushOp = current.PushOpCodes[index].Op;
802                          MergeType pushMergeType = GetMergeTypeFromPush(currentPushOp.Name);
803                          branchStack.Push((currentPushOp.GetAbsoluteAddress(), pushMergeType));
804                      }
805                  }
806  
807                  pushOpIndex = 0;
808  
809                  bool hasNext = current.HasNext();
810                  if (hasNext)
811                  {
812                      Push(new PathBlockState(current.Successors[0]));
813                  }
814  
815                  InstOp lastOp = current.GetLastOp();
816                  if (IsPopBranch(lastOp.Name))
817                  {
818                      MergeType popMergeType = GetMergeTypeFromPop(lastOp.Name);
819  
820                      bool found = true;
821                      ulong targetAddress = 0UL;
822                      MergeType mergeType;
823  
824                      do
825                      {
826                          if (branchStack.Count == 0)
827                          {
828                              found = false;
829                              break;
830                          }
831  
832                          (targetAddress, mergeType) = branchStack.Pop();
833  
834                          // Push the target address (this will be used to push the address
835                          // back into the PBK/PCNT/SSY stack when we return from that block),
836                          Push(new PathBlockState(targetAddress, mergeType));
837                      }
838                      while (mergeType != popMergeType);
839  
840                      // Make sure we found the correct address,
841                      // the push and pop instruction types must match, so:
842                      // - BRK can only consume addresses pushed by PBK.
843                      // - CONT can only consume addresses pushed by PCNT.
844                      // - SYNC can only consume addresses pushed by SSY.
845                      if (found)
846                      {
847                          if (branchStack.Count == 0)
848                          {
849                              // If the entire stack was consumed, then the current pop instruction
850                              // just consumed the address from our push instruction.
851                              if (current.SyncTargets.TryAdd(pushOp.Address, new SyncTarget(pushOpInfo, current.SyncTargets.Count)))
852                              {
853                                  pushOpInfo.Consumers.Add(current, Local());
854                                  target.Predecessors.Add(current);
855                                  current.Successors.Add(target);
856                              }
857                          }
858                          else
859                          {
860                              // Push the block itself into the work queue for processing.
861                              Push(new PathBlockState(blocks[targetAddress]));
862                          }
863                      }
864                  }
865                  else
866                  {
867                      // By adding them in descending order (sorted by address), we process the blocks
868                      // in order (of ascending address), since we work with a LIFO.
869                      foreach (Block possibleTarget in current.Successors.OrderByDescending(x => x.Address))
870                      {
871                          if (!hasNext || possibleTarget != current.Successors[0])
872                          {
873                              Push(new PathBlockState(possibleTarget));
874                          }
875                      }
876                  }
877              }
878          }
879  
880          public static bool IsPopBranch(InstName name)
881          {
882              return name == InstName.Brk || name == InstName.Cont || name == InstName.Sync;
883          }
884  
885          private static MergeType GetMergeTypeFromPush(InstName name)
886          {
887              return name switch
888              {
889                  InstName.Pbk => MergeType.Brk,
890                  InstName.Pcnt => MergeType.Cont,
891                  _ => MergeType.Sync,
892              };
893          }
894  
895          private static MergeType GetMergeTypeFromPop(InstName name)
896          {
897              return name switch
898              {
899                  InstName.Brk => MergeType.Brk,
900                  InstName.Cont => MergeType.Cont,
901                  _ => MergeType.Sync,
902              };
903          }
904      }
905  }