nir: Add a find_variable_with_[driver_]location helper

We've hand-rolled this loop 10 places and those are just the ones I
found easily.

Reviewed-by: Eric Anholt <eric@anholt.net>
Reviewed-by: Iago Toral Quiroga <itoral@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5966>

src/amd/vulkan/radv_shader.c

@@ -586,14 +586,12 @@ type_size_vec4(const struct glsl_type *type, bool bindless)
 static nir_variable *
 find_layer_in_var(nir_shader *nir)
 {
-	nir_foreach_shader_in_variable(var, nir) {
-		if (var->data.location == VARYING_SLOT_LAYER) {
-			return var;
-		}
-	}
-
 	nir_variable *var =
-		nir_variable_create(nir, nir_var_shader_in, glsl_int_type(), "layer id");
+		nir_find_variable_with_location(nir, nir_var_shader_in, VARYING_SLOT_LAYER);
+	if (var != NULL)
+		return var;
+
+	var = nir_variable_create(nir, nir_var_shader_in, glsl_int_type(), "layer id");
 	var->data.location = VARYING_SLOT_LAYER;
 	var->data.interpolation = INTERP_MODE_FLAT;
 	return var;

src/compiler/nir/nir.c

@@ -171,6 +171,32 @@ nir_local_variable_create(nir_function_impl *impl,
    return var;
 }
 
+nir_variable *
+nir_find_variable_with_location(nir_shader *shader,
+                                nir_variable_mode mode,
+                                unsigned location)
+{
+   assert(util_bitcount(mode) == 1 && mode != nir_var_function_temp);
+   nir_foreach_variable_with_modes(var, shader, mode) {
+      if (var->data.location == location)
+         return var;
+   }
+   return NULL;
+}
+
+nir_variable *
+nir_find_variable_with_driver_location(nir_shader *shader,
+                                       nir_variable_mode mode,
+                                       unsigned location)
+{
+   assert(util_bitcount(mode) == 1 && mode != nir_var_function_temp);
+   nir_foreach_variable_with_modes(var, shader, mode) {
+      if (var->data.driver_location == location)
+         return var;
+   }
+   return NULL;
+}
+
 nir_function *
 nir_function_create(nir_shader *shader, const char *name)
 {

src/compiler/nir/nir.h

@@ -3308,6 +3308,14 @@ nir_variable *nir_local_variable_create(nir_function_impl *impl,
                                         const struct glsl_type *type,
                                         const char *name);
 
+nir_variable *nir_find_variable_with_location(nir_shader *shader,
+                                              nir_variable_mode mode,
+                                              unsigned location);
+
+nir_variable *nir_find_variable_with_driver_location(nir_shader *shader,
+                                                     nir_variable_mode mode,
+                                                     unsigned location);
+
 /** creates a function and adds it to the shader's list of functions */
 nir_function *nir_function_create(nir_shader *shader, const char *name);
 

src/compiler/nir/nir_lower_bitmap.c

@@ -55,16 +55,9 @@
 static nir_variable *
 get_texcoord(nir_shader *shader)
 {
-   nir_variable *texcoord = NULL;
-
-   /* find gl_TexCoord, if it exists: */
-   nir_foreach_shader_in_variable(var, shader) {
-      if (var->data.location == VARYING_SLOT_TEX0) {
-         texcoord = var;
-         break;
-      }
-   }
-
+   nir_variable *texcoord =
+      nir_find_variable_with_location(shader, nir_var_shader_in,
+                                      VARYING_SLOT_TEX0);
    /* otherwise create it: */
    if (texcoord == NULL) {
       texcoord = nir_variable_create(shader,

src/compiler/nir/nir_lower_input_attachments.c

@@ -27,14 +27,14 @@
 static nir_ssa_def *
 load_frag_coord(nir_builder *b)
 {
-   nir_foreach_shader_in_variable(var, b->shader) {
-      if (var->data.location == VARYING_SLOT_POS)
-         return nir_load_var(b, var);
+   nir_variable *pos =
+      nir_find_variable_with_location(b->shader, nir_var_shader_in,
+                                      VARYING_SLOT_POS);
+   if (pos == NULL) {
+      pos = nir_variable_create(b->shader, nir_var_shader_in,
+                                glsl_vec4_type(), NULL);
+      pos->data.location = VARYING_SLOT_POS;
    }
-
-   nir_variable *pos = nir_variable_create(b->shader, nir_var_shader_in,
-                                           glsl_vec4_type(), NULL);
-   pos->data.location = VARYING_SLOT_POS;
    /**
     * From Vulkan spec:
     *   "The OriginLowerLeft execution mode must not be used; fragment entry

src/compiler/nir/nir_lower_point_size_mov.c

@@ -71,13 +71,9 @@ nir_lower_point_size_mov(nir_shader *shader,
    assert(shader->info.stage != MESA_SHADER_FRAGMENT &&
           shader->info.stage != MESA_SHADER_COMPUTE);
 
-   nir_variable *out = NULL;
-   nir_foreach_shader_out_variable(var, shader) {
-      if (var->data.location == VARYING_SLOT_PSIZ) {
-         out = var;
-         break;
-      }
-   }
+   nir_variable *out =
+      nir_find_variable_with_location(shader, nir_var_shader_out,
+                                      VARYING_SLOT_PSIZ);
 
    lower_impl(nir_shader_get_entrypoint(shader), pointsize_state_tokens,
               out);

src/compiler/nir/nir_lower_two_sided_color.c

@@ -66,19 +66,20 @@ create_input(nir_shader *shader, gl_varying_slot slot,
 static nir_variable *
 create_face_input(nir_shader *shader)
 {
-   nir_foreach_shader_in_variable(var, shader) {
-      if (var->data.location == VARYING_SLOT_FACE)
-         return var;
+   nir_variable *var =
+      nir_find_variable_with_location(shader, nir_var_shader_in,
+                                      VARYING_SLOT_FACE);
+
+   if (var == NULL) {
+      var = nir_variable_create(shader, nir_var_shader_in,
+                                glsl_bool_type(), "gl_FrontFacing");
+
+      var->data.driver_location = shader->num_inputs++;
+      var->data.index = 0;
+      var->data.location = VARYING_SLOT_FACE;
+      var->data.interpolation = INTERP_MODE_FLAT;
    }
 
-   nir_variable *var = nir_variable_create(shader, nir_var_shader_in,
-                                           glsl_bool_type(), "gl_FrontFacing");
-
-   var->data.driver_location = shader->num_inputs++;
-   var->data.index = 0;
-   var->data.location = VARYING_SLOT_FACE;
-   var->data.interpolation = INTERP_MODE_FLAT;
-
    return var;
 }
 

src/compiler/nir/nir_range_analysis.c

@@ -1141,11 +1141,8 @@ search_phi_bcsel(nir_ssa_scalar scalar, nir_ssa_scalar *buf, unsigned buf_size,
 static nir_variable *
 lookup_input(nir_shader *shader, unsigned driver_location)
 {
-   nir_foreach_shader_in_variable(var, shader) {
-      if (driver_location == var->data.driver_location)
-         return var;
-   }
-   return NULL;
+   return nir_find_variable_with_driver_location(shader, nir_var_shader_in,
+                                                 driver_location);
 }
 
 uint32_t

src/gallium/drivers/iris/iris_program.c

@@ -273,14 +273,8 @@ iris_fix_edge_flags(nir_shader *nir)
       return false;
    }
 
-   nir_variable *var = NULL;
-   nir_foreach_shader_out_variable(v, nir) {
-      if (v->data.location == VARYING_SLOT_EDGE) {
-         var = v;
-         break;
-      }
-   }
-
+   nir_variable *var = nir_find_variable_with_location(nir, nir_var_shader_out,
+                                                       VARYING_SLOT_EDGE);
    if (!var) {
       nir_shader_preserve_all_metadata(nir);
       return false;

src/gallium/drivers/vc4/vc4_nir_lower_io.c

@@ -239,13 +239,9 @@ vc4_nir_lower_fs_input(struct vc4_compile *c, nir_builder *b,
                 return;
         }
 
-        nir_variable *input_var = NULL;
-        nir_foreach_shader_in_variable(var, c->s) {
-                if (var->data.driver_location == nir_intrinsic_base(intr)) {
-                        input_var = var;
-                        break;
-                }
-        }
+        nir_variable *input_var =
+                nir_find_variable_with_driver_location(c->s, nir_var_shader_in,
+                                                       nir_intrinsic_base(intr));
         assert(input_var);
 
         int comp = nir_intrinsic_component(intr);
@@ -290,13 +286,9 @@ static void
 vc4_nir_lower_output(struct vc4_compile *c, nir_builder *b,
                      nir_intrinsic_instr *intr)
 {
-        nir_variable *output_var = NULL;
-        nir_foreach_shader_out_variable(var, c->s) {
-                if (var->data.driver_location == nir_intrinsic_base(intr)) {
-                        output_var = var;
-                        break;
-                }
-        }
+        nir_variable *output_var =
+                nir_find_variable_with_driver_location(c->s, nir_var_shader_out,
+                                                       nir_intrinsic_base(intr));
         assert(output_var);
 
         if (c->stage == QSTAGE_COORD &&

src/intel/compiler/brw_nir_lower_alpha_to_coverage.c

@@ -88,15 +88,8 @@ brw_nir_lower_alpha_to_coverage(nir_shader *shader)
    assert(shader->info.stage == MESA_SHADER_FRAGMENT);
 
    /* Bail out early if we don't have gl_SampleMask */
-   bool is_sample_mask = false;
-   nir_foreach_shader_out_variable(var, shader) {
-      if (var->data.location == FRAG_RESULT_SAMPLE_MASK) {
-         is_sample_mask = true;
-         break;
-      }
-   }
-
-   if (!is_sample_mask)
+   if (!nir_find_variable_with_location(shader, nir_var_shader_out,
+                                        FRAG_RESULT_SAMPLE_MASK))
       return;
 
    nir_foreach_function(function, shader) {
@@ -115,16 +108,9 @@ brw_nir_lower_alpha_to_coverage(nir_shader *shader)
 
                switch (intr->intrinsic) {
                case nir_intrinsic_store_output:
-                  nir_foreach_shader_out_variable(var, shader) {
-                     int drvloc = var->data.driver_location;
-                     if (nir_intrinsic_base(intr) == drvloc) {
-                        out = var;
-                        break;
-                     }
-                  }
-
-                  if (out->data.mode != nir_var_shader_out)
-                     continue;
+                  out = nir_find_variable_with_driver_location(shader, nir_var_shader_out,
+                                                               nir_intrinsic_base(intr));
+                  assert(out->data.mode == nir_var_shader_out);
 
                   /* save gl_SampleMask instruction pointer */
                   if (out->data.location == FRAG_RESULT_SAMPLE_MASK) {