nir: Add a function for rewriting the condition of an if statement

Reviewed-by: Connor Abbott <cwabbott0@gmail.com>

diff --git a/src/glsl/nir/nir.c b/src/glsl/nir/nir.c
index 4cc074b808f68907cf0febb318255442a71d1dc6..b8f5dd491a4eaed3071645a7154f594aa57a7c98 100644 (file)
--- a/src/glsl/nir/nir.c
+++ b/src/glsl/nir/nir.c
@@ -1894,6 +1894,28 @@ nir_instr_rewrite_src(nir_instr *instr, nir_src *src, nir_src new_src)
    }
 }
 
+void
+nir_if_rewrite_condition(nir_if *if_stmt, nir_src new_src)
+{
+   for (nir_src *src = &if_stmt->condition; src;
+        src = src->is_ssa ? NULL : src->reg.indirect) {
+      struct set *uses = src->is_ssa ? src->ssa->if_uses
+                                     : src->reg.reg->if_uses;
+      struct set_entry *entry = _mesa_set_search(uses, if_stmt);
+      assert(entry);
+      _mesa_set_remove(uses, entry);
+   }
+
+   if_stmt->condition = new_src;
+
+   for (nir_src *src = &if_stmt->condition; src;
+        src = src->is_ssa ? NULL : src->reg.indirect) {
+      struct set *uses = src->is_ssa ? src->ssa->if_uses
+                                     : src->reg.reg->if_uses;
+      _mesa_set_add(uses, if_stmt);
+   }
+}
+
 void
 nir_ssa_def_init(nir_instr *instr, nir_ssa_def *def,
                  unsigned num_components, const char *name)

diff --git a/src/glsl/nir/nir.h b/src/glsl/nir/nir.h
index a174666e0e4b31206d098d96154f80db85e1886a..aaf1c572ebadf7218edcf858912e21e89d50baf9 100644 (file)
--- a/src/glsl/nir/nir.h
+++ b/src/glsl/nir/nir.h
@@ -1549,6 +1549,7 @@ bool nir_foreach_src(nir_instr *instr, nir_foreach_src_cb cb, void *state);
 nir_const_value *nir_src_as_const_value(nir_src src);
 bool nir_srcs_equal(nir_src src1, nir_src src2);
 void nir_instr_rewrite_src(nir_instr *instr, nir_src *src, nir_src new_src);
+void nir_if_rewrite_condition(nir_if *if_stmt, nir_src new_src);
 
 void nir_ssa_dest_init(nir_instr *instr, nir_dest *dest,
                        unsigned num_components, const char *name);