nir: Properly clean up nir_src/dest indirects

Now that they're no longer ralloc'd, we have to be much more careful
about indirects.  We have to make sure every time a source or
destination is overwritten, its indirect (if any) is freed.  We also
have to choose a memory ownership convention for the rewrite functions.
Assuming that they will be called with the source from some other
instruction, we choose to always make a copy of the indirect (if any).
It's the responsibility of the caller to ensure its copy of the indirect
is freed.

Unfortunately, all this extra logic is going to make
nir_instr_rewrite/move_src/dest more expensive because they now have
all the logic of nir_src/dest_copy instead of a simple struct
assignment.  Fortunately, the vast majority of rewrite calls are done by
nir_ssa_def_rewrite_uses which is an SSA-only fast-path.

Fixes: 879a569884 "nir: Switch from ralloc to malloc for NIR instructions."
Reviewed-by: Emma Anholt <emma@anholt.net>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/12884>

src/compiler/nir/nir.c

@@ -339,11 +339,36 @@ nir_function_create(nir_shader *shader, const char *name)
    return func;
 }
 
+static bool src_has_indirect(nir_src *src)
+{
+   return !src->is_ssa && src->reg.indirect;
+}
+
+static void src_free_indirects(nir_src *src)
+{
+   if (src_has_indirect(src)) {
+      assert(src->reg.indirect->is_ssa || !src->reg.indirect->reg.indirect);
+      free(src->reg.indirect);
+      src->reg.indirect = NULL;
+   }
+}
+
+static void dest_free_indirects(nir_dest *dest)
+{
+   if (!dest->is_ssa && dest->reg.indirect) {
+      assert(dest->reg.indirect->is_ssa || !dest->reg.indirect->reg.indirect);
+      free(dest->reg.indirect);
+      dest->reg.indirect = NULL;
+   }
+}
+
 /* NOTE: if the instruction you are copying a src to is already added
  * to the IR, use nir_instr_rewrite_src() instead.
  */
 void nir_src_copy(nir_src *dest, const nir_src *src)
 {
+   src_free_indirects(dest);
+
    dest->is_ssa = src->is_ssa;
    if (src->is_ssa) {
       dest->ssa = src->ssa;
@@ -364,6 +389,8 @@ void nir_dest_copy(nir_dest *dest, const nir_dest *src)
    /* Copying an SSA definition makes no sense whatsoever. */
    assert(!src->is_ssa);
 
+   dest_free_indirects(dest);
+
    dest->is_ssa = false;
 
    dest->reg.base_offset = src->reg.base_offset;
@@ -1120,30 +1147,22 @@ void nir_instr_remove_v(nir_instr *instr)
    }
 }
 
-static bool nir_instr_free_src_indirects(nir_src *src, void *state)
+static bool free_src_indirects_cb(nir_src *src, void *state)
 {
-   if (!src->is_ssa && src->reg.indirect) {
-      assert(src->reg.indirect->is_ssa || !src->reg.indirect->reg.indirect);
-      free(src->reg.indirect);
-      src->reg.indirect = NULL;
-   }
+   src_free_indirects(src);
    return true;
 }
 
-static bool nir_instr_free_dest_indirects(nir_dest *dest, void *state)
+static bool free_dest_indirects_cb(nir_dest *dest, void *state)
 {
-   if (!dest->is_ssa && dest->reg.indirect) {
-      assert(dest->reg.indirect->is_ssa || !dest->reg.indirect->reg.indirect);
-      free(dest->reg.indirect);
-      dest->reg.indirect = NULL;
-   }
+   dest_free_indirects(dest);
    return true;
 }
 
 void nir_instr_free(nir_instr *instr)
 {
-   nir_foreach_src(instr, nir_instr_free_src_indirects, NULL);
-   nir_foreach_dest(instr, nir_instr_free_dest_indirects, NULL);
+   nir_foreach_src(instr, free_src_indirects_cb, NULL);
+   nir_foreach_dest(instr, free_dest_indirects_cb, NULL);
 
    switch (instr->type) {
    case nir_instr_type_tex:
@@ -1531,7 +1550,7 @@ nir_instr_rewrite_src(nir_instr *instr, nir_src *src, nir_src new_src)
    assert(!src_is_valid(src) || src->parent_instr == instr);
 
    src_remove_all_uses(src);
-   *src = new_src;
+   nir_src_copy(src, &new_src);
    src_add_all_uses(src, instr, NULL);
 }
 
@@ -1541,6 +1560,7 @@ nir_instr_move_src(nir_instr *dest_instr, nir_src *dest, nir_src *src)
    assert(!src_is_valid(dest) || dest->parent_instr == dest_instr);
 
    src_remove_all_uses(dest);
+   src_free_indirects(dest);
    src_remove_all_uses(src);
    *dest = *src;
    *src = NIR_SRC_INIT;
@@ -1554,7 +1574,7 @@ nir_if_rewrite_condition(nir_if *if_stmt, nir_src new_src)
    assert(!src_is_valid(src) || src->parent_if == if_stmt);
 
    src_remove_all_uses(src);
-   *src = new_src;
+   nir_src_copy(src, &new_src);
    src_add_all_uses(src, NULL, if_stmt);
 }