nir: Add a new helper for iterating phi sources leaving a block

This takes the same callback as nir_foreach_src except it walks all phi
sources which leave a given block.

Reviewed-by: Alyssa Rosenzweig <[email protected]>
Reviewed-by: Rob Clark <[email protected]>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5094>

3 files changed, 30 insertions, 15 deletions

diff --git a/src/compiler/nir/nir.c b/src/compiler/nir/nir.c
index 0f64c458535..7c1d159c3bd 100644
--- a/src/compiler/nir/nir.c
+++ b/src/compiler/nir/nir.c
@@ -1256,6 +1256,32 @@ nir_foreach_src(nir_instr *instr, nir_foreach_src_cb cb, void *state)
    return nir_foreach_dest(instr, visit_dest_indirect, &dest_state);
 }
 
+bool
+nir_foreach_phi_src_leaving_block(nir_block *block,
+                                  nir_foreach_src_cb cb,
+                                  void *state)
+{
+   for (unsigned i = 0; i < ARRAY_SIZE(block->successors); i++) {
+      if (block->successors[i] == NULL)
+         continue;
+
+      nir_foreach_instr(instr, block->successors[i]) {
+         if (instr->type != nir_instr_type_phi)
+            break;
+
+         nir_phi_instr *phi = nir_instr_as_phi(instr);
+         nir_foreach_phi_src(phi_src, phi) {
+            if (phi_src->pred == block) {
+               if (!cb(&phi_src->src, state))
+                  return false;
+            }
+         }
+      }
+   }
+
+   return true;
+}
+
 nir_const_value
 nir_const_value_for_float(double f, unsigned bit_size)
 {
diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index ca799c540a1..ba9c030df9b 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -3554,6 +3554,9 @@ bool nir_foreach_ssa_def(nir_instr *instr, nir_foreach_ssa_def_cb cb,
                          void *state);
 bool nir_foreach_dest(nir_instr *instr, nir_foreach_dest_cb cb, void *state);
 bool nir_foreach_src(nir_instr *instr, nir_foreach_src_cb cb, void *state);
+bool nir_foreach_phi_src_leaving_block(nir_block *instr,
+                                       nir_foreach_src_cb cb,
+                                       void *state);
 
 nir_const_value *nir_src_as_const_value(nir_src src);
 
diff --git a/src/compiler/nir/nir_lower_regs_to_ssa.c b/src/compiler/nir/nir_lower_regs_to_ssa.c
index 027c5db504c..e045d97bac7 100644
--- a/src/compiler/nir/nir_lower_regs_to_ssa.c
+++ b/src/compiler/nir/nir_lower_regs_to_ssa.c
@@ -273,21 +273,7 @@ nir_lower_regs_to_ssa_impl(nir_function_impl *impl)
        * loops, a phi source may be a back-edge so we have to handle it as if
        * it were one of the last instructions in the predecessor block.
        */
-      for (unsigned i = 0; i < ARRAY_SIZE(block->successors); i++) {
-         if (block->successors[i] == NULL)
-            continue;
-
-         nir_foreach_instr(instr, block->successors[i]) {
-            if (instr->type != nir_instr_type_phi)
-               break;
-
-            nir_phi_instr *phi = nir_instr_as_phi(instr);
-            nir_foreach_phi_src(phi_src, phi) {
-               if (phi_src->pred == block)
-                  rewrite_src(&phi_src->src, &state);
-            }
-         }
-      }
+      nir_foreach_phi_src_leaving_block(block, rewrite_src, &state);
    }
 
    nir_phi_builder_finish(phi_build);