shader: Abstract breadth searches and use the abstraction

src/shader_recompiler/CMakeLists.txt

@@ -27,6 +27,7 @@ add_library(shader_recompiler STATIC
     frontend/ir/attribute.h
     frontend/ir/basic_block.cpp
     frontend/ir/basic_block.h
+    frontend/ir/breadth_first_search.h
     frontend/ir/condition.cpp
     frontend/ir/condition.h
     frontend/ir/flow_test.cpp

src/shader_recompiler/frontend/ir/breadth_first_search.h

@@ -0,0 +1,57 @@
+// Copyright 2021 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include <optional>
+#include <type_traits>
+#include <queue>
+
+#include <boost/container/small_vector.hpp>
+
+#include "shader_recompiler/frontend/ir/microinstruction.h"
+#include "shader_recompiler/frontend/ir/value.h"
+
+namespace Shader::IR {
+
+template <typename Pred>
+auto BreadthFirstSearch(const Value& value, Pred&& pred)
+    -> std::invoke_result_t<Pred, const Inst*> {
+    if (value.IsImmediate()) {
+        // Nothing to do with immediates
+        return std::nullopt;
+    }
+    // Breadth-first search visiting the right most arguments first
+    // Small vector has been determined from shaders in Super Smash Bros. Ultimate
+    boost::container::small_vector<const Inst*, 2> visited;
+    std::queue<const Inst*> queue;
+    queue.push(value.InstRecursive());
+
+    while (!queue.empty()) {
+        // Pop one instruction from the queue
+        const Inst* const inst{queue.front()};
+        queue.pop();
+        if (const std::optional result = pred(inst)) {
+            // This is the instruction we were looking for
+            return result;
+        }
+        // Visit the right most arguments first
+        for (size_t arg = inst->NumArgs(); arg--;) {
+            const Value arg_value{inst->Arg(arg)};
+            if (arg_value.IsImmediate()) {
+                continue;
+            }
+            // Queue instruction if it hasn't been visited
+            const Inst* const arg_inst{arg_value.InstRecursive()};
+            if (std::ranges::find(visited, arg_inst) == visited.end()) {
+                visited.push_back(arg_inst);
+                queue.push(arg_inst);
+            }
+        }
+    }
+    // SSA tree has been traversed and the result hasn't been found
+    return std::nullopt;
+}
+
+} // namespace Shader::IR

src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp

@@ -12,6 +12,7 @@
 #include <boost/container/small_vector.hpp>
 
 #include "shader_recompiler/frontend/ir/basic_block.h"
+#include "shader_recompiler/frontend/ir/breadth_first_search.h"
 #include "shader_recompiler/frontend/ir/ir_emitter.h"
 #include "shader_recompiler/frontend/ir/microinstruction.h"
 #include "shader_recompiler/ir_opt/passes.h"
@@ -219,15 +220,17 @@ std::optional<LowAddrInfo> TrackLowAddress(IR::Inst* inst) {
     };
 }
 
-/// Tries to get the storage buffer out of a constant buffer read instruction
-std::optional<StorageBufferAddr> TryGetStorageBuffer(const IR::Inst* inst, const Bias* bias) {
+/// Tries to track the storage buffer address used by a global memory instruction
+std::optional<StorageBufferAddr> Track(const IR::Value& value, const Bias* bias) {
+    const auto pred{[bias](const IR::Inst* inst) -> std::optional<StorageBufferAddr> {
         if (inst->Opcode() != IR::Opcode::GetCbufU32) {
             return std::nullopt;
         }
         const IR::Value index{inst->Arg(0)};
         const IR::Value offset{inst->Arg(1)};
         if (!index.IsImmediate()) {
-        // Definitely not a storage buffer if it's read from a non-immediate index
+            // Definitely not a storage buffer if it's read from a
+            // non-immediate index
             return std::nullopt;
         }
         if (!offset.IsImmediate()) {
@@ -239,48 +242,13 @@ std::optional<StorageBufferAddr> TryGetStorageBuffer(const IR::Inst* inst, const
             .offset{offset.U32()},
         };
         if (bias && !MeetsBias(storage_buffer, *bias)) {
-        // We have to blacklist some addresses in case we wrongly point to them
+            // We have to blacklist some addresses in case we wrongly
+            // point to them
             return std::nullopt;
         }
         return storage_buffer;
-}
-
-/// Tries to track the storage buffer address used by a global memory instruction
-std::optional<StorageBufferAddr> Track(const IR::Value& value, const Bias* bias) {
-    if (value.IsImmediate()) {
-        // Nothing to do with immediates
-        return std::nullopt;
-    }
-    // Breadth-first search visiting the right most arguments first
-    // Small vector has been determined from shaders in Super Smash Bros. Ultimate
-    small_vector<const IR::Inst*, 2> visited;
-    std::queue<const IR::Inst*> queue;
-    queue.push(value.InstRecursive());
-
-    while (!queue.empty()) {
-        // Pop one instruction from the queue
-        const IR::Inst* const inst{queue.front()};
-        queue.pop();
-        if (const std::optional<StorageBufferAddr> result = TryGetStorageBuffer(inst, bias)) {
-            // This is the instruction we were looking for
-            return result;
-        }
-        // Visit the right most arguments first
-        for (size_t arg = inst->NumArgs(); arg--;) {
-            const IR::Value arg_value{inst->Arg(arg)};
-            if (arg_value.IsImmediate()) {
-                continue;
-            }
-            // Queue instruction if it hasn't been visited
-            const IR::Inst* const arg_inst{arg_value.InstRecursive()};
-            if (std::ranges::find(visited, arg_inst) == visited.end()) {
-                visited.push_back(arg_inst);
-                queue.push(arg_inst);
-            }
-        }
-    }
-    // SSA tree has been traversed and the origin hasn't been found
-    return std::nullopt;
+    }};
+    return BreadthFirstSearch(value, pred);
 }
 
 /// Collects the storage buffer used by a global memory instruction and the instruction itself

src/shader_recompiler/ir_opt/texture_pass.cpp

@@ -2,13 +2,14 @@
 // Licensed under GPLv2 or any later version
 // Refer to the license.txt file included.
 
+#include <algorithm>
 #include <optional>
 
-#include <boost/container/flat_set.hpp>
 #include <boost/container/small_vector.hpp>
 
 #include "shader_recompiler/environment.h"
 #include "shader_recompiler/frontend/ir/basic_block.h"
+#include "shader_recompiler/frontend/ir/breadth_first_search.h"
 #include "shader_recompiler/frontend/ir/ir_emitter.h"
 #include "shader_recompiler/ir_opt/passes.h"
 #include "shader_recompiler/shader_info.h"
@@ -28,9 +29,6 @@ struct TextureInst {
 
 using TextureInstVector = boost::container::small_vector<TextureInst, 24>;
 
-using VisitedBlocks = boost::container::flat_set<IR::Block*, std::less<IR::Block*>,
-                                                 boost::container::small_vector<IR::Block*, 2>>;
-
 IR::Opcode IndexedInstruction(const IR::Inst& inst) {
     switch (inst.Opcode()) {
     case IR::Opcode::BindlessImageSampleImplicitLod:
@@ -101,14 +99,10 @@ bool IsTextureInstruction(const IR::Inst& inst) {
     return IndexedInstruction(inst) != IR::Opcode::Void;
 }
 
-std::optional<ConstBufferAddr> Track(IR::Block* block, const IR::Value& value,
-                                     VisitedBlocks& visited) {
-    if (value.IsImmediate()) {
-        // Immediates can't be a storage buffer
+std::optional<ConstBufferAddr> TryGetConstBuffer(const IR::Inst* inst) {
+    if (inst->Opcode() != IR::Opcode::GetCbufU32) {
         return std::nullopt;
     }
-    const IR::Inst* const inst{value.InstRecursive()};
-    if (inst->Opcode() == IR::Opcode::GetCbufU32) {
     const IR::Value index{inst->Arg(0)};
     const IR::Value offset{inst->Arg(1)};
     if (!index.IsImmediate()) {
@@ -124,34 +118,16 @@ std::optional<ConstBufferAddr> Track(IR::Block* block, const IR::Value& value,
         .index{index.U32()},
         .offset{offset.U32()},
     };
-    }
-    // Reversed loops are more likely to find the right result
-    for (size_t arg = inst->NumArgs(); arg--;) {
-        IR::Block* inst_block{block};
-        if (inst->Opcode() == IR::Opcode::Phi) {
-            // If we are going through a phi node, mark the current block as visited
-            visited.insert(block);
-            // and skip already visited blocks to avoid looping forever
-            IR::Block* const phi_block{inst->PhiBlock(arg)};
-            if (visited.contains(phi_block)) {
-                // Already visited, skip
-                continue;
-            }
-            inst_block = phi_block;
-        }
-        const std::optional storage_buffer{Track(inst_block, inst->Arg(arg), visited)};
-        if (storage_buffer) {
-            return *storage_buffer;
-        }
-    }
-    return std::nullopt;
+}
+
+std::optional<ConstBufferAddr> Track(const IR::Value& value) {
+    return IR::BreadthFirstSearch(value, TryGetConstBuffer);
 }
 
 TextureInst MakeInst(Environment& env, IR::Block* block, IR::Inst& inst) {
     ConstBufferAddr addr;
     if (IsBindless(inst)) {
-        VisitedBlocks visited;
-        const std::optional<ConstBufferAddr> track_addr{Track(block, inst.Arg(0), visited)};
+        const std::optional<ConstBufferAddr> track_addr{Track(inst.Arg(0))};
         if (!track_addr) {
             throw NotImplementedException("Failed to track bindless texture constant buffer");
         }