nir: Add and use a nir_instr_init_src() helper

This helper exists for a very tiny set of use-cases but it's better to
have the helper live in nir.c than hand-roll it elsewhere.

Reviewed-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24729>

src/compiler/nir/nir.c

@@ -797,8 +797,7 @@ nir_tex_instr_add_src(nir_tex_instr *tex,
    tex->src = new_srcs;
 
    tex->src[tex->num_srcs].src_type = src_type;
-   nir_instr_rewrite_src(&tex->instr, &tex->src[tex->num_srcs].src,
-                         nir_src_for_ssa(src));
+   nir_instr_init_src(&tex->instr, &tex->src[tex->num_srcs].src, src);
    tex->num_srcs++;
 }
 
@@ -1476,6 +1475,13 @@ nir_instr_rewrite_src(nir_instr *instr, nir_src *src, nir_src new_src)
    src_add_all_uses(src, instr, NULL);
 }
 
+void
+nir_instr_init_src(nir_instr *instr, nir_src *src, nir_def *def)
+{
+   *src = nir_src_for_ssa(def);
+   src_add_all_uses(src, instr, NULL);
+}
+
 void
 nir_instr_clear_src(nir_instr *instr, nir_src *src)
 {

src/compiler/nir/nir.h

@@ -4421,6 +4421,23 @@ nir_instr_rewrite_src_ssa(ASSERTED nir_instr *instr,
 
 void nir_instr_rewrite_src(nir_instr *instr, nir_src *src, nir_src new_src);
 
+/** Initialize a nir_src
+ *
+ * This is almost never the helper you want to use.  This helper assumes that
+ * the source is uninitialized garbage and blasts over it without doing any
+ * tear-down the existing source, including removing it from uses lists.
+ * Using this helper on a source that currently exists in any uses list will
+ * result in linked list corruption.  It also assumes that the instruction is
+ * currently live in the IR and adds the source to the uses list for the given
+ * nir_def as part of setup.
+ *
+ * This is pretty much only useful for adding sources to extant instructions
+ * or manipulating parallel copy instructions as part of out-of-SSA.
+ *
+ * When in doubt, use nir_src_rewrite() instead.
+ */
+void nir_instr_init_src(nir_instr *instr, nir_src *src, nir_def *def);
+
 /** Clear a nir_src
  *
  * This helper clears a nir_src by removing it from any uses lists and

src/compiler/nir/nir_from_ssa.c

@@ -399,14 +399,19 @@ isolate_phi_nodes_block(nir_shader *shader, nir_block *block, void *dead_ctx)
 
          nir_parallel_copy_entry *entry = rzalloc(dead_ctx,
                                                   nir_parallel_copy_entry);
-         entry->src_is_reg = false;
+
          entry->dest_is_reg = false;
          nir_def_init(&pcopy->instr, &entry->dest.def,
                       phi->def.num_components, phi->def.bit_size);
          entry->dest.def.divergent = nir_src_is_divergent(src->src);
-         exec_list_push_tail(&pcopy->entries, &entry->node);
 
-         nir_instr_rewrite_src(&pcopy->instr, &entry->src, src->src);
+         /* We're adding a source to a live instruction so we need to use
+          * nir_instr_init_src()
+          */
+         entry->src_is_reg = false;
+         nir_instr_init_src(&pcopy->instr, &entry->src, src->src.ssa);
+
+         exec_list_push_tail(&pcopy->entries, &entry->node);
 
          nir_instr_rewrite_src(&phi->instr, &src->src,
                                nir_src_for_ssa(&entry->dest.def));
@@ -414,18 +419,25 @@ isolate_phi_nodes_block(nir_shader *shader, nir_block *block, void *dead_ctx)
 
       nir_parallel_copy_entry *entry = rzalloc(dead_ctx,
                                                nir_parallel_copy_entry);
-      entry->src_is_reg = false;
+
       entry->dest_is_reg = false;
       nir_def_init(&block_pcopy->instr, &entry->dest.def,
                    phi->def.num_components, phi->def.bit_size);
       entry->dest.def.divergent = phi->def.divergent;
+
+      nir_def_rewrite_uses(&phi->def, &entry->dest.def);
+
+      /* We're adding a source to a live instruction so we need to use
+       * nir_instr_init_src().
+       *
+       * Note that we do this after we've rewritten all uses of the phi to
+       * entry->def, ensuring that entry->src will be the only remaining use
+       * of the phi.
+       */
+      entry->src_is_reg = false;
+      nir_instr_init_src(&block_pcopy->instr, &entry->src, &phi->def);
+
       exec_list_push_tail(&block_pcopy->entries, &entry->node);
-
-      nir_def_rewrite_uses(&phi->def,
-                           &entry->dest.def);
-
-      nir_instr_rewrite_src(&block_pcopy->instr, &entry->src,
-                            nir_src_for_ssa(&phi->def));
    }
 
    return true;
@@ -716,16 +728,17 @@ resolve_registers_impl(nir_function_impl *impl, struct from_ssa_state *state)
 
             nir_foreach_parallel_copy_entry(entry, pcopy) {
                assert(!entry->dest_is_reg);
-               assert(nir_def_is_unused(&entry->dest.def));
 
                /* Parallel copy destinations will always be registers */
                nir_def *reg = reg_for_ssa_def(&entry->dest.def, state);
                assert(reg != NULL);
 
+               /* We're switching from the nir_def to the nir_src in the dest
+                * union so we need to use nir_instr_init_src() here.
+                */
+               assert(nir_def_is_unused(&entry->dest.def));
                entry->dest_is_reg = true;
-               entry->dest.reg = NIR_SRC_INIT;
-               nir_instr_rewrite_src(&pcopy->instr, &entry->dest.reg,
-                                     nir_src_for_ssa(reg));
+               nir_instr_init_src(&pcopy->instr, &entry->dest.reg, reg);
             }
 
             nir_foreach_parallel_copy_entry(entry, pcopy) {