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/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