zink: correctly set up ubo bindings and buffer indices

var->data.binding is only set for vulkan drivers (though it also will
get incremented after nir_lower_uniforms_to_ubo), so we have to generate
our own values here.

to do this, we iterate backwards over the ubos to account for the
"first" ubos being at the end of the list, and we must also ensure that we
remap the buffer index correctly based on whether we're running our
nir_lower_uniforms_to_ubo pass

note that running nir_lower_uniforms_to_ubo unconditionally would require
us to add a number of checks in this patch for !shader->num_uniforms in
order to properly adjust to the altered instructions which become 1-indexed
instead of 0-indexed when this pass is run with no uniforms present

Reviewed-by: Erik Faye-Lund <erik.faye-lund@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/6981>
diff --git a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
index ae45cbc..98a1937 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
+++ b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
@@ -2333,12 +2333,15 @@
    nir_foreach_shader_out_variable(var, s)
       emit_output(&ctx, var);
 
+
    if (so_info)
       emit_so_info(&ctx, util_last_bit64(s->info.outputs_written), so_info);
-   nir_foreach_variable_with_modes(var, s, nir_var_uniform |
-                                           nir_var_mem_ubo |
-                                           nir_var_mem_ssbo)
-      emit_uniform(&ctx, var);
+   /* we have to reverse iterate to match what's done in zink_compiler.c */
+   foreach_list_typed_reverse(nir_variable, var, node, &s->variables)
+      if (_nir_shader_variable_has_mode(var, nir_var_uniform |
+                                        nir_var_mem_ubo |
+                                        nir_var_mem_ssbo))
+         emit_uniform(&ctx, var);
 
    if (s->info.stage == MESA_SHADER_FRAGMENT) {
       spirv_builder_emit_exec_mode(&ctx.builder, entry_point,
diff --git a/src/gallium/drivers/zink/zink_compiler.c b/src/gallium/drivers/zink/zink_compiler.c
index eca835c..9c7bdce 100644
--- a/src/gallium/drivers/zink/zink_compiler.c
+++ b/src/gallium/drivers/zink/zink_compiler.c
@@ -281,46 +281,59 @@
    }
 
    ret->num_bindings = 0;
-   nir_foreach_variable_with_modes(var, nir, nir_var_uniform |
-                                             nir_var_mem_ubo) {
-      if (var->data.mode == nir_var_mem_ubo) {
-         /* ignore variables being accessed if they aren't the base of the UBO */
-         if (var->data.location)
-            continue;
-         int binding = zink_binding(nir->info.stage,
-                                    VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
-                                    var->data.binding);
-         ret->bindings[ret->num_bindings].index = var->data.binding;
-         ret->bindings[ret->num_bindings].binding = binding;
-         ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
-         ret->num_bindings++;
-      } else {
-         assert(var->data.mode == nir_var_uniform);
-         if (glsl_type_is_sampler(var->type)) {
-            int binding = zink_binding(nir->info.stage,
-                                       VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-                                       var->data.binding);
-            ret->bindings[ret->num_bindings].index = var->data.binding;
-            ret->bindings[ret->num_bindings].binding = binding;
-            ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
-            ret->num_bindings++;
-         } else if (glsl_type_is_array(var->type)) {
-            /* need to unroll possible arrays of arrays before checking type
-             * in order to handle ARB_arrays_of_arrays extension
-             */
-            const struct glsl_type *type = glsl_without_array(var->type);
-            if (!glsl_type_is_sampler(type))
+   uint32_t cur_ubo = 0;
+   /* UBO buffers are zero-indexed, but buffer 0 is always the one created by nir_lower_uniforms_to_ubo,
+    * which means there is no buffer 0 if there are no uniforms
+    */
+   int ubo_index = !nir->num_uniforms;
+   /* need to set up var->data.binding for UBOs, which means we need to start at
+    * the "first" UBO, which is at the end of the list
+    */
+   foreach_list_typed_reverse(nir_variable, var, node, &nir->variables) {
+      if (_nir_shader_variable_has_mode(var, nir_var_uniform |
+                                        nir_var_mem_ubo |
+                                        nir_var_mem_ssbo)) {
+         if (var->data.mode == nir_var_mem_ubo) {
+            /* ignore variables being accessed if they aren't the base of the UBO */
+            if (var->data.location)
                continue;
+            var->data.binding = cur_ubo++;
 
-            unsigned size = glsl_get_aoa_size(var->type);
-            for (int i = 0; i < size; ++i) {
+            int binding = zink_binding(nir->info.stage,
+                                       VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
+                                       var->data.binding);
+            ret->bindings[ret->num_bindings].index = ubo_index++;
+            ret->bindings[ret->num_bindings].binding = binding;
+            ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
+            ret->num_bindings++;
+         } else {
+            assert(var->data.mode == nir_var_uniform);
+            if (glsl_type_is_sampler(var->type)) {
                int binding = zink_binding(nir->info.stage,
                                           VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
-                                          var->data.binding + i);
-               ret->bindings[ret->num_bindings].index = var->data.binding + i;
+                                          var->data.binding);
+               ret->bindings[ret->num_bindings].index = var->data.binding;
                ret->bindings[ret->num_bindings].binding = binding;
                ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
                ret->num_bindings++;
+            } else if (glsl_type_is_array(var->type)) {
+               /* need to unroll possible arrays of arrays before checking type
+                * in order to handle ARB_arrays_of_arrays extension
+                */
+               const struct glsl_type *type = glsl_without_array(var->type);
+               if (!glsl_type_is_sampler(type))
+                  continue;
+
+               unsigned size = glsl_get_aoa_size(var->type);
+               for (int i = 0; i < size; ++i) {
+                  int binding = zink_binding(nir->info.stage,
+                                             VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
+                                             var->data.binding + i);
+                  ret->bindings[ret->num_bindings].index = var->data.binding + i;
+                  ret->bindings[ret->num_bindings].binding = binding;
+                  ret->bindings[ret->num_bindings].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
+                  ret->num_bindings++;
+               }
             }
          }
       }