nir/nir: use nir_foreach_ssa_src()

1 files changed, 48 insertions, 10 deletions

diff --git a/src/glsl/nir/nir.c b/src/glsl/nir/nir.c
index e89a0eddf4..db8a605e34 100644
--- a/src/glsl/nir/nir.c
+++ b/src/glsl/nir/nir.c
@@ -1238,9 +1238,19 @@ add_use_cb(nir_src *src, void *state)
 {
    nir_instr *instr = state;
 
-   struct set *uses_set = src->is_ssa ? src->ssa->uses : src->reg.reg->uses;
+   if (src->is_ssa)
+      return true;
+
+   _mesa_set_add(src->reg.reg->uses, instr);
+
+   return true;
+}
+
+static bool add_ssa_use_cb(nir_ssa_def **src, void *state)
+{
+   nir_instr *instr = state;
 
-   _mesa_set_add(uses_set, instr);
+   _mesa_set_add((*src)->uses, instr);
 
    return true;
 }
@@ -1275,6 +1285,7 @@ static void
 add_defs_uses(nir_instr *instr)
 {
    nir_foreach_src(instr, add_use_cb, instr);
+   nir_foreach_ssa_src(instr, add_ssa_use_cb, instr);
    nir_foreach_dest(instr, add_reg_def_cb, instr);
    nir_foreach_ssa_def(instr, add_ssa_def_cb, instr);
 }
@@ -1383,8 +1394,23 @@ remove_use_cb(nir_src *src, void *state)
 {
    nir_instr *instr = state;
 
-   struct set *uses_set = src->is_ssa ? src->ssa->uses : src->reg.reg->uses;
+   if (src->is_ssa)
+      return true;
 
+   struct set *uses_set = src->reg.reg->uses;
+   struct set_entry *entry = _mesa_set_search(uses_set, instr);
+   if (entry)
+      _mesa_set_remove(uses_set, entry);
+
+   return true;
+}
+
+static bool
+remove_ssa_use_cb(nir_ssa_def **src, void *state)
+{
+   nir_instr *instr = state;
+
+   struct set *uses_set = (*src)->uses;
    struct set_entry *entry = _mesa_set_search(uses_set, instr);
    if (entry)
       _mesa_set_remove(uses_set, entry);
@@ -1414,6 +1440,7 @@ remove_defs_uses(nir_instr *instr)
 {
    nir_foreach_dest(instr, remove_def_cb, instr);
    nir_foreach_src(instr, remove_use_cb, instr);
+   nir_foreach_ssa_src(instr, remove_ssa_use_cb, instr);
 }
 
 void nir_instr_remove(nir_instr *instr)
@@ -1807,15 +1834,11 @@ nir_srcs_equal(nir_src src1, nir_src src2)
 }
 
 static bool
-src_does_not_use_def(nir_src *src, void *void_def)
+src_does_not_use_def(nir_ssa_def **src, void *void_def)
 {
    nir_ssa_def *def = void_def;
 
-   if (src->is_ssa) {
-      return src->ssa != def;
-   } else {
-      return true;
-   }
+   return *src != def;
 }
 
 static bool
@@ -1836,7 +1859,8 @@ nir_instr_rewrite_src(nir_instr *instr, nir_src *src, nir_src new_src)
    if (src->is_ssa) {
       nir_ssa_def *old_ssa = src->ssa;
       *src = new_src;
-      if (old_ssa && nir_foreach_src(instr, src_does_not_use_def, old_ssa)) {
+      if (old_ssa && nir_foreach_ssa_src(instr, src_does_not_use_def,
+                                         old_ssa)) {
          struct set_entry *entry = _mesa_set_search(old_ssa->uses, instr);
          assert(entry);
          _mesa_set_remove(old_ssa->uses, entry);
@@ -1912,6 +1936,18 @@ ssa_def_rewrite_uses_src(nir_src *src, void *void_state)
    return true;
 }
 
+static bool
+ssa_def_rewrite_uses_ssa_src(nir_ssa_def **src, void *void_state)
+{
+   struct ssa_def_rewrite_state *state = void_state;
+
+   if (*src == state->old)
+      *src = state->new_src.ssa;
+
+   return true;
+}
+
+
 void
 nir_ssa_def_rewrite_uses(nir_ssa_def *def, nir_src new_src, void *mem_ctx)
 {
@@ -1937,6 +1973,8 @@ nir_ssa_def_rewrite_uses(nir_ssa_def *def, nir_src new_src, void *mem_ctx)
 
       _mesa_set_remove(def->uses, entry);
       nir_foreach_src(instr, ssa_def_rewrite_uses_src, &state);
+      if (new_src.is_ssa)
+         nir_foreach_ssa_src(instr, ssa_def_rewrite_uses_ssa_src, &state);
       _mesa_set_add(new_uses, instr);
    }