nir: Add and use a nir_variable_list_for_mode helper

We also add a new list iterator which takes a modes bitfield and
automatically figures out which list to use.  In the future, this
iterator will work for multiple modes but today it assumes a single mode
thanks to the behavior of nir_variable_list_for_mode.  This also doesn't
work for function_temp variables.

Reviewed-by: Gert Wollny <gert.wollny@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5966>

3 files changed, 46 insertions, 25 deletions

diff --git a/src/compiler/nir/nir.c b/src/compiler/nir/nir.c
index c33cdc69032..3d0fd42b316 100644
--- a/src/compiler/nir/nir.c
+++ b/src/compiler/nir/nir.c
@@ -104,56 +104,58 @@ nir_reg_remove(nir_register *reg)
    exec_node_remove(&reg->node);
 }
 
-void
-nir_shader_add_variable(nir_shader *shader, nir_variable *var)
+struct exec_list *
+nir_variable_list_for_mode(nir_shader *shader, nir_variable_mode mode)
 {
-   switch (var->data.mode) {
-   case nir_num_variable_modes:
-   case nir_var_all:
-      assert(!"invalid mode");
-      break;
-
+   switch (mode) {
    case nir_var_function_temp:
       assert(!"nir_shader_add_variable cannot be used for local variables");
-      break;
+      return NULL;
 
    case nir_var_shader_temp:
-      exec_list_push_tail(&shader->globals, &var->node);
-      break;
+      return &shader->globals;
 
    case nir_var_shader_in:
-      exec_list_push_tail(&shader->inputs, &var->node);
-      break;
+      return &shader->inputs;
 
    case nir_var_shader_out:
-      exec_list_push_tail(&shader->outputs, &var->node);
-      break;
+      return &shader->outputs;
 
    case nir_var_uniform:
    case nir_var_mem_ubo:
    case nir_var_mem_ssbo:
-      exec_list_push_tail(&shader->uniforms, &var->node);
-      break;
+      return &shader->uniforms;
 
    case nir_var_mem_shared:
       assert(gl_shader_stage_is_compute(shader->info.stage));
-      exec_list_push_tail(&shader->shared, &var->node);
-      break;
+      return &shader->shared;
 
    case nir_var_mem_global:
       assert(!"nir_shader_add_variable cannot be used for global memory");
-      break;
+      return NULL;
 
    case nir_var_system_value:
-      exec_list_push_tail(&shader->system_values, &var->node);
-      break;
+      return &shader->system_values;
 
    case nir_var_mem_push_const:
       assert(!"nir_var_push_constant is not supposed to be used for variables");
-      break;
+      return NULL;
+
+   default:
+      assert(!"invalid mode");
+      return NULL;
    }
 }
 
+void
+nir_shader_add_variable(nir_shader *shader, nir_variable *var)
+{
+   struct exec_list *var_list =
+      nir_variable_list_for_mode(shader, var->data.mode);
+   if (var_list)
+      exec_list_push_tail(var_list, &var->node);
+}
+
 nir_variable *
 nir_variable_create(nir_shader *shader, nir_variable_mode mode,
                     const struct glsl_type *type, const char *name)
diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index b7fba8b59a4..7c04a773fbc 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -618,12 +618,29 @@ typedef struct nir_variable {
    struct nir_variable_data *members;
 } nir_variable;
 
+
+static inline bool
+_nir_shader_variable_has_mode(nir_variable *var, unsigned modes)
+{
+   /* This isn't a shader variable */
+   assert(!(modes & nir_var_function_temp));
+   return var->data.mode & modes;
+}
+
 #define nir_foreach_variable(var, var_list) \
    foreach_list_typed(nir_variable, var, node, var_list)
 
 #define nir_foreach_variable_safe(var, var_list) \
    foreach_list_typed_safe(nir_variable, var, node, var_list)
 
+#define nir_foreach_variable_with_modes(var, shader, modes) \
+   nir_foreach_variable(var, nir_variable_list_for_mode(shader, modes)) \
+      if (_nir_shader_variable_has_mode(var, modes))
+
+#define nir_foreach_variable_with_modes_safe(var, shader, modes) \
+   nir_foreach_variable_safe(var, nir_variable_list_for_mode(shader, modes)) \
+      if (_nir_shader_variable_has_mode(var, modes))
+
 #define nir_foreach_shader_in_variable(var, shader) \
    nir_foreach_variable(var, &(shader)->inputs)
 
@@ -3283,6 +3300,9 @@ nir_register *nir_local_reg_create(nir_function_impl *impl);
 
 void nir_reg_remove(nir_register *reg);
 
+struct exec_list *
+nir_variable_list_for_mode(nir_shader *shader, nir_variable_mode mode);
+
 /** Adds a variable to the appropriate list in nir_shader */
 void nir_shader_add_variable(nir_shader *shader, nir_variable *var);
 
diff --git a/src/compiler/nir/nir_linking_helpers.c b/src/compiler/nir/nir_linking_helpers.c
index d3a249edd82..8047d122ffe 100644
--- a/src/compiler/nir/nir_linking_helpers.c
+++ b/src/compiler/nir/nir_linking_helpers.c
@@ -128,8 +128,7 @@ nir_remove_unused_io_vars(nir_shader *shader,
    uint64_t *used;
 
    assert(mode == nir_var_shader_in || mode == nir_var_shader_out);
-   struct exec_list *var_list =
-      mode == nir_var_shader_in ? &shader->inputs : &shader->outputs;
+   struct exec_list *var_list = nir_variable_list_for_mode(shader, mode);
 
    nir_foreach_variable_safe(var, var_list) {
       if (var->data.patch)