/*
 * Copyright 2018 Collabora Ltd.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * on the rights to use, copy, modify, merge, publish, distribute, sub
 * license, and/or sell copies of the Software, and to permit persons to whom
 * the Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice (including the next
 * paragraph) shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL
 * THE AUTHOR(S) AND/OR THEIR SUPPLIERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 * USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#include "spirv_builder.h"

#include "util/macros.h"
#include "util/u_bitcast.h"
#include "util/u_memory.h"
#include "util/hash_table.h"

#include <stdbool.h>
#include <inttypes.h>
#include <string.h>

static bool
spirv_buffer_grow(struct spirv_buffer *b, size_t needed)
{
   size_t new_room = MAX3(64, (b->room * 3) / 2, needed);

   uint32_t *new_words = realloc(b->words, new_room * sizeof(uint32_t));
   if (!new_words)
      return false;

   b->words = new_words;
   b->room = new_room;
   return true;
}

static inline bool
spirv_buffer_prepare(struct spirv_buffer *b, size_t needed)
{
   needed += b->num_words;
   if (b->room >= b->num_words + needed)
      return true;

   return spirv_buffer_grow(b, needed);
}

static inline void
spirv_buffer_emit_word(struct spirv_buffer *b, uint32_t word)
{
   assert(b->num_words < b->room);
   b->words[b->num_words++] = word;
}

static int
spirv_buffer_emit_string(struct spirv_buffer *b, const char *str)
{
   int pos = 0;
   uint32_t word = 0;
   while (str[pos] != '\0') {
      word |= str[pos] << (8 * (pos % 4));
      if (++pos % 4 == 0) {
         spirv_buffer_prepare(b, 1);
         spirv_buffer_emit_word(b, word);
         word = 0;
      }
   }

   spirv_buffer_prepare(b, 1);
   spirv_buffer_emit_word(b, word);

   return 1 + pos / 4;
}

void
spirv_builder_emit_cap(struct spirv_builder *b, SpvCapability cap)
{
   spirv_buffer_prepare(&b->capabilities, 2);
   spirv_buffer_emit_word(&b->capabilities, SpvOpCapability | (2 << 16));
   spirv_buffer_emit_word(&b->capabilities, cap);
}

void
spirv_builder_emit_source(struct spirv_builder *b, SpvSourceLanguage lang,
                          uint32_t version)
{
   spirv_buffer_prepare(&b->debug_names, 3);
   spirv_buffer_emit_word(&b->debug_names, SpvOpSource | (3 << 16));
   spirv_buffer_emit_word(&b->debug_names, lang);
   spirv_buffer_emit_word(&b->debug_names, version);
}

void
spirv_builder_emit_mem_model(struct spirv_builder *b,
                             SpvAddressingModel addr_model,
                             SpvMemoryModel mem_model)
{
   spirv_buffer_prepare(&b->memory_model, 3);
   spirv_buffer_emit_word(&b->memory_model, SpvOpMemoryModel | (3 << 16));
   spirv_buffer_emit_word(&b->memory_model, addr_model);
   spirv_buffer_emit_word(&b->memory_model, mem_model);
}

void
spirv_builder_emit_entry_point(struct spirv_builder *b,
                               SpvExecutionModel exec_model, SpvId entry_point,
                               const char *name, const SpvId interfaces[],
                               size_t num_interfaces)
{
   size_t pos = b->entry_points.num_words;
   spirv_buffer_prepare(&b->entry_points, 3);
   spirv_buffer_emit_word(&b->entry_points, SpvOpEntryPoint);
   spirv_buffer_emit_word(&b->entry_points, exec_model);
   spirv_buffer_emit_word(&b->entry_points, entry_point);
   int len = spirv_buffer_emit_string(&b->entry_points, name);
   b->entry_points.words[pos] |= (3 + len + num_interfaces) << 16;
   spirv_buffer_prepare(&b->entry_points, num_interfaces);
   for (int i = 0; i < num_interfaces; ++i)
        spirv_buffer_emit_word(&b->entry_points, interfaces[i]);
}

void
spirv_builder_emit_exec_mode(struct spirv_builder *b, SpvId entry_point,
                             SpvExecutionMode exec_mode)
{
   spirv_buffer_prepare(&b->exec_modes, 3);
   spirv_buffer_emit_word(&b->exec_modes, SpvOpExecutionMode | (3 << 16));
   spirv_buffer_emit_word(&b->exec_modes, entry_point);
   spirv_buffer_emit_word(&b->exec_modes, exec_mode);
}

void
spirv_builder_emit_name(struct spirv_builder *b, SpvId target,
                        const char *name)
{
   size_t pos = b->debug_names.num_words;
   spirv_buffer_prepare(&b->debug_names, 2);
   spirv_buffer_emit_word(&b->debug_names, SpvOpName);
   spirv_buffer_emit_word(&b->debug_names, target);
   int len = spirv_buffer_emit_string(&b->debug_names, name);
   b->debug_names.words[pos] |= (2 + len) << 16;
}

static void
emit_decoration(struct spirv_builder *b, SpvId target,
                SpvDecoration decoration, const uint32_t extra_operands[],
                size_t num_extra_operands)
{
   int words = 3 + num_extra_operands;
   spirv_buffer_prepare(&b->decorations, words);
   spirv_buffer_emit_word(&b->decorations, SpvOpDecorate | (words << 16));
   spirv_buffer_emit_word(&b->decorations, target);
   spirv_buffer_emit_word(&b->decorations, decoration);
   for (int i = 0; i < num_extra_operands; ++i)
      spirv_buffer_emit_word(&b->decorations, extra_operands[i]);
}

void
spirv_builder_emit_decoration(struct spirv_builder *b, SpvId target,
                              SpvDecoration decoration)
{
   emit_decoration(b, target, decoration, NULL, 0);
}

void
spirv_builder_emit_location(struct spirv_builder *b, SpvId target,
                            uint32_t location)
{
   uint32_t args[] = { location };
   emit_decoration(b, target, SpvDecorationLocation, args, ARRAY_SIZE(args));
}

void
spirv_builder_emit_component(struct spirv_builder *b, SpvId target,
                             uint32_t component)
{
   uint32_t args[] = { component };
   emit_decoration(b, target, SpvDecorationComponent, args, ARRAY_SIZE(args));
}

void
spirv_builder_emit_builtin(struct spirv_builder *b, SpvId target,
                           SpvBuiltIn builtin)
{
   uint32_t args[] = { builtin };
   emit_decoration(b, target, SpvDecorationBuiltIn, args, ARRAY_SIZE(args));
}

void
spirv_builder_emit_descriptor_set(struct spirv_builder *b, SpvId target,
                                  uint32_t descriptor_set)
{
   uint32_t args[] = { descriptor_set };
   emit_decoration(b, target, SpvDecorationDescriptorSet, args,
                   ARRAY_SIZE(args));
}

void
spirv_builder_emit_binding(struct spirv_builder *b, SpvId target,
                           uint32_t binding)
{
   uint32_t args[] = { binding };
   emit_decoration(b, target, SpvDecorationBinding, args, ARRAY_SIZE(args));
}

void
spirv_builder_emit_array_stride(struct spirv_builder *b, SpvId target,
                                uint32_t stride)
{
   uint32_t args[] = { stride };
   emit_decoration(b, target, SpvDecorationArrayStride, args, ARRAY_SIZE(args));
}

void
spirv_builder_emit_index(struct spirv_builder *b, SpvId target, int index)
{
   uint32_t args[] = { index };
   emit_decoration(b, target, SpvDecorationIndex, args, ARRAY_SIZE(args));
}

static void
emit_member_decoration(struct spirv_builder *b, SpvId target, uint32_t member,
                       SpvDecoration decoration, const uint32_t extra_operands[],
                       size_t num_extra_operands)
{
   int words = 4 + num_extra_operands;
   spirv_buffer_prepare(&b->decorations, words);
   spirv_buffer_emit_word(&b->decorations,
                          SpvOpMemberDecorate | (words << 16));
   spirv_buffer_emit_word(&b->decorations, target);
   spirv_buffer_emit_word(&b->decorations, member);
   spirv_buffer_emit_word(&b->decorations, decoration);
   for (int i = 0; i < num_extra_operands; ++i)
      spirv_buffer_emit_word(&b->decorations, extra_operands[i]);
}

void
spirv_builder_emit_member_offset(struct spirv_builder *b, SpvId target,
                          uint32_t member, uint32_t offset)
{
   uint32_t args[] = { offset };
   emit_member_decoration(b, target, member, SpvDecorationOffset,
                          args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_emit_undef(struct spirv_builder *b, SpvId result_type)
{
   SpvId result = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->instructions, 3);
   spirv_buffer_emit_word(&b->instructions, SpvOpUndef | (3 << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   return result;
}

void
spirv_builder_function(struct spirv_builder *b, SpvId result,
                       SpvId return_type,
                       SpvFunctionControlMask function_control,
                       SpvId function_type)
{
   spirv_buffer_prepare(&b->instructions, 5);
   spirv_buffer_emit_word(&b->instructions, SpvOpFunction | (5 << 16));
   spirv_buffer_emit_word(&b->instructions, return_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, function_control);
   spirv_buffer_emit_word(&b->instructions, function_type);
}

void
spirv_builder_function_end(struct spirv_builder *b)
{
   spirv_buffer_prepare(&b->instructions, 1);
   spirv_buffer_emit_word(&b->instructions, SpvOpFunctionEnd | (1 << 16));
}

void
spirv_builder_label(struct spirv_builder *b, SpvId label)
{
   spirv_buffer_prepare(&b->instructions, 2);
   spirv_buffer_emit_word(&b->instructions, SpvOpLabel | (2 << 16));
   spirv_buffer_emit_word(&b->instructions, label);
}

void
spirv_builder_return(struct spirv_builder *b)
{
   spirv_buffer_prepare(&b->instructions, 1);
   spirv_buffer_emit_word(&b->instructions, SpvOpReturn | (1 << 16));
}

SpvId
spirv_builder_emit_load(struct spirv_builder *b, SpvId result_type,
                        SpvId pointer)
{
   return spirv_builder_emit_unop(b, SpvOpLoad, result_type, pointer);
}

void
spirv_builder_emit_store(struct spirv_builder *b, SpvId pointer, SpvId object)
{
   spirv_buffer_prepare(&b->instructions, 3);
   spirv_buffer_emit_word(&b->instructions, SpvOpStore | (3 << 16));
   spirv_buffer_emit_word(&b->instructions, pointer);
   spirv_buffer_emit_word(&b->instructions, object);
}

SpvId
spirv_builder_emit_access_chain(struct spirv_builder *b, SpvId result_type,
                                SpvId base, const SpvId indexes[],
                                size_t num_indexes)
{
   SpvId result = spirv_builder_new_id(b);

   int words = 4 + num_indexes;
   spirv_buffer_prepare(&b->instructions, words);
   spirv_buffer_emit_word(&b->instructions, SpvOpAccessChain | (words << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, base);
   for (int i = 0; i < num_indexes; ++i)
      spirv_buffer_emit_word(&b->instructions, indexes[i]);
   return result;
}


SpvId
spirv_builder_emit_unop(struct spirv_builder *b, SpvOp op, SpvId result_type,
                        SpvId operand)
{
   SpvId result = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->instructions, 4);
   spirv_buffer_emit_word(&b->instructions, op | (4 << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, operand);
   return result;
}

SpvId
spirv_builder_emit_binop(struct spirv_builder *b, SpvOp op, SpvId result_type,
                         SpvId operand0, SpvId operand1)
{
   SpvId result = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->instructions, 5);
   spirv_buffer_emit_word(&b->instructions, op | (5 << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, operand0);
   spirv_buffer_emit_word(&b->instructions, operand1);
   return result;
}

SpvId
spirv_builder_emit_triop(struct spirv_builder *b, SpvOp op, SpvId result_type,
                         SpvId operand0, SpvId operand1, SpvId operand2)
{
   SpvId result = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->instructions, 6);
   spirv_buffer_emit_word(&b->instructions, op | (6 << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, operand0);
   spirv_buffer_emit_word(&b->instructions, operand1);
   spirv_buffer_emit_word(&b->instructions, operand2);
   return result;
}

SpvId
spirv_builder_emit_composite_extract(struct spirv_builder *b, SpvId result_type,
                                     SpvId composite, const uint32_t indexes[],
                                     size_t num_indexes)
{
   SpvId result = spirv_builder_new_id(b);

   assert(num_indexes > 0);
   int words = 4 + num_indexes;
   spirv_buffer_prepare(&b->instructions, words);
   spirv_buffer_emit_word(&b->instructions,
                          SpvOpCompositeExtract | (words << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, composite);
   for (int i = 0; i < num_indexes; ++i)
      spirv_buffer_emit_word(&b->instructions, indexes[i]);
   return result;
}

SpvId
spirv_builder_emit_composite_construct(struct spirv_builder *b,
                                       SpvId result_type,
                                       const SpvId constituents[],
                                       size_t num_constituents)
{
   SpvId result = spirv_builder_new_id(b);

   assert(num_constituents > 0);
   int words = 3 + num_constituents;
   spirv_buffer_prepare(&b->instructions, words);
   spirv_buffer_emit_word(&b->instructions,
                          SpvOpCompositeConstruct | (words << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   for (int i = 0; i < num_constituents; ++i)
      spirv_buffer_emit_word(&b->instructions, constituents[i]);
   return result;
}

SpvId
spirv_builder_emit_vector_shuffle(struct spirv_builder *b, SpvId result_type,
                                  SpvId vector_1, SpvId vector_2,
                                  const uint32_t components[],
                                  size_t num_components)
{
   SpvId result = spirv_builder_new_id(b);

   assert(num_components > 0);
   int words = 5 + num_components;
   spirv_buffer_prepare(&b->instructions, words);
   spirv_buffer_emit_word(&b->instructions, SpvOpVectorShuffle | (words << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, vector_1);
   spirv_buffer_emit_word(&b->instructions, vector_2);
   for (int i = 0; i < num_components; ++i)
      spirv_buffer_emit_word(&b->instructions, components[i]);
   return result;
}

void
spirv_builder_emit_branch(struct spirv_builder *b, SpvId label)
{
   spirv_buffer_prepare(&b->instructions, 2);
   spirv_buffer_emit_word(&b->instructions, SpvOpBranch | (2 << 16));
   spirv_buffer_emit_word(&b->instructions, label);
}

void
spirv_builder_emit_selection_merge(struct spirv_builder *b, SpvId merge_block,
                                   SpvSelectionControlMask selection_control)
{
   spirv_buffer_prepare(&b->instructions, 3);
   spirv_buffer_emit_word(&b->instructions, SpvOpSelectionMerge | (3 << 16));
   spirv_buffer_emit_word(&b->instructions, merge_block);
   spirv_buffer_emit_word(&b->instructions, selection_control);
}

void
spirv_builder_loop_merge(struct spirv_builder *b, SpvId merge_block,
                         SpvId cont_target, SpvLoopControlMask loop_control)
{
   spirv_buffer_prepare(&b->instructions, 4);
   spirv_buffer_emit_word(&b->instructions, SpvOpLoopMerge | (4 << 16));
   spirv_buffer_emit_word(&b->instructions, merge_block);
   spirv_buffer_emit_word(&b->instructions, cont_target);
   spirv_buffer_emit_word(&b->instructions, loop_control);
}

void
spirv_builder_emit_branch_conditional(struct spirv_builder *b, SpvId condition,
                                      SpvId true_label, SpvId false_label)
{
   spirv_buffer_prepare(&b->instructions, 4);
   spirv_buffer_emit_word(&b->instructions, SpvOpBranchConditional | (4 << 16));
   spirv_buffer_emit_word(&b->instructions, condition);
   spirv_buffer_emit_word(&b->instructions, true_label);
   spirv_buffer_emit_word(&b->instructions, false_label);
}

SpvId
spirv_builder_emit_phi(struct spirv_builder *b, SpvId result_type,
                       size_t num_vars, size_t *position)
{
   SpvId result = spirv_builder_new_id(b);

   assert(num_vars > 0);
   int words = 3 + 2 * num_vars;
   spirv_buffer_prepare(&b->instructions, words);
   spirv_buffer_emit_word(&b->instructions, SpvOpPhi | (words << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   *position = b->instructions.num_words;
   for (int i = 0; i < 2 * num_vars; ++i)
      spirv_buffer_emit_word(&b->instructions, 0);
   return result;
}

void
spirv_builder_set_phi_operand(struct spirv_builder *b, size_t position,
                              size_t index, SpvId variable, SpvId parent)
{
   b->instructions.words[position + index * 2 + 0] = variable;
   b->instructions.words[position + index * 2 + 1] = parent;
}

void
spirv_builder_emit_kill(struct spirv_builder *b)
{
   spirv_buffer_prepare(&b->instructions, 1);
   spirv_buffer_emit_word(&b->instructions, SpvOpKill | (1 << 16));
}

SpvId
spirv_builder_emit_image_sample(struct spirv_builder *b,
                                SpvId result_type,
                                SpvId sampled_image,
                                SpvId coordinate,
                                bool proj,
                                SpvId lod,
                                SpvId bias,
                                SpvId dref,
                                SpvId dx,
                                SpvId dy,
                                SpvId offset)
{
   SpvId result = spirv_builder_new_id(b);

   int opcode = SpvOpImageSampleImplicitLod;
   int operands = 5;
   if (proj)
      opcode += SpvOpImageSampleProjImplicitLod - SpvOpImageSampleImplicitLod;
   if (lod || (dx && dy))
      opcode += SpvOpImageSampleExplicitLod - SpvOpImageSampleImplicitLod;
   if (dref) {
      opcode += SpvOpImageSampleDrefImplicitLod - SpvOpImageSampleImplicitLod;
      operands++;
   }

   SpvImageOperandsMask operand_mask = SpvImageOperandsMaskNone;
   SpvId extra_operands[5];
   int num_extra_operands = 0;
   if (bias) {
      extra_operands[++num_extra_operands] = bias;
      operand_mask |= SpvImageOperandsBiasMask;
   }
   if (lod) {
      extra_operands[++num_extra_operands] = lod;
      operand_mask |= SpvImageOperandsLodMask;
   } else if (dx && dy) {
      extra_operands[++num_extra_operands] = dx;
      extra_operands[++num_extra_operands] = dy;
      operand_mask |= SpvImageOperandsGradMask;
   }
   if (offset) {
      extra_operands[++num_extra_operands] = offset;
      operand_mask |= SpvImageOperandsOffsetMask;
   }

   /* finalize num_extra_operands / extra_operands */
   if (num_extra_operands > 0) {
      extra_operands[0] = operand_mask;
      num_extra_operands++;
   }

   spirv_buffer_prepare(&b->instructions, operands + num_extra_operands);
   spirv_buffer_emit_word(&b->instructions, opcode | ((operands + num_extra_operands) << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, sampled_image);
   spirv_buffer_emit_word(&b->instructions, coordinate);
   if (dref)
      spirv_buffer_emit_word(&b->instructions, dref);
   for (int i = 0; i < num_extra_operands; ++i)
      spirv_buffer_emit_word(&b->instructions, extra_operands[i]);
   return result;
}

SpvId
spirv_builder_emit_image(struct spirv_builder *b, SpvId result_type,
                         SpvId sampled_image)
{
   SpvId result = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->instructions, 4);
   spirv_buffer_emit_word(&b->instructions, SpvOpImage | (4 << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, sampled_image);
   return result;
}

SpvId
spirv_builder_emit_image_fetch(struct spirv_builder *b,
                               SpvId result_type,
                               SpvId image,
                               SpvId coordinate,
                               SpvId lod)
{
   SpvId result = spirv_builder_new_id(b);

   SpvId extra_operands[2];
   int num_extra_operands = 0;
   if (lod) {
      extra_operands[0] = SpvImageOperandsLodMask;
      extra_operands[1] = lod;
      num_extra_operands = 2;
   }

   spirv_buffer_prepare(&b->instructions, 5 + num_extra_operands);
   spirv_buffer_emit_word(&b->instructions, SpvOpImageFetch |
                          ((5 + num_extra_operands) << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, image);
   spirv_buffer_emit_word(&b->instructions, coordinate);
   for (int i = 0; i < num_extra_operands; ++i)
      spirv_buffer_emit_word(&b->instructions, extra_operands[i]);
   return result;
}

SpvId
spirv_builder_emit_image_query_size(struct spirv_builder *b,
                                    SpvId result_type,
                                    SpvId image,
                                    SpvId lod)
{
   int opcode = SpvOpImageQuerySize;
   int words = 4;
   if (lod) {
      words++;
      opcode = SpvOpImageQuerySizeLod;
   }

   SpvId result = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->instructions, words);
   spirv_buffer_emit_word(&b->instructions, opcode | (words << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, image);

   if (lod)
      spirv_buffer_emit_word(&b->instructions, lod);

   return result;
}

SpvId
spirv_builder_emit_ext_inst(struct spirv_builder *b, SpvId result_type,
                            SpvId set, uint32_t instruction,
                            const SpvId *args, size_t num_args)
{
   SpvId result = spirv_builder_new_id(b);

   int words = 5 + num_args;
   spirv_buffer_prepare(&b->instructions, words);
   spirv_buffer_emit_word(&b->instructions, SpvOpExtInst | (words << 16));
   spirv_buffer_emit_word(&b->instructions, result_type);
   spirv_buffer_emit_word(&b->instructions, result);
   spirv_buffer_emit_word(&b->instructions, set);
   spirv_buffer_emit_word(&b->instructions, instruction);
   for (int i = 0; i < num_args; ++i)
      spirv_buffer_emit_word(&b->instructions, args[i]);
   return result;
}

struct spirv_type {
   SpvOp op;
   uint32_t args[8];
   size_t num_args;

   SpvId type;
};

static uint32_t
non_aggregate_type_hash(const void *arg)
{
   const struct spirv_type *type = arg;

   uint32_t hash = _mesa_fnv32_1a_offset_bias;
   hash = _mesa_fnv32_1a_accumulate(hash, type->op);
   hash = _mesa_fnv32_1a_accumulate_block(hash, type->args, sizeof(uint32_t) *
                                          type->num_args);
   return hash;
}

static bool
non_aggregate_type_equals(const void *a, const void *b)
{
   const struct spirv_type *ta = a, *tb = b;

   if (ta->op != tb->op)
      return false;

   assert(ta->num_args == tb->num_args);
   return memcmp(ta->args, tb->args, sizeof(uint32_t) * ta->num_args) == 0;
}

static SpvId
get_type_def(struct spirv_builder *b, SpvOp op, const uint32_t args[],
             size_t num_args)
{
   /* According to the SPIR-V specification:
    *
    *   "Two different type <id>s form, by definition, two different types. It
    *    is valid to declare multiple aggregate type <id>s having the same
    *    opcode and operands. This is to allow multiple instances of aggregate
    *    types with the same structure to be decorated differently. (Different
    *    decorations are not required; two different aggregate type <id>s are
    *    allowed to have identical declarations and decorations, and will still
    *    be two different types.) Non-aggregate types are different: It is
    *    invalid to declare multiple type <id>s for the same scalar, vector, or
    *    matrix type. That is, non-aggregate type declarations must all have
    *    different opcodes or operands. (Note that non-aggregate types cannot
    *    be decorated in ways that affect their type.)"
    *
    *  ..so, we need to prevent the same non-aggregate type to be re-defined
    *  with a new <id>. We do this by putting the definitions in a hash-map, so
    *  we can easily look up and reuse them.
    */

   struct spirv_type key;
   assert(num_args <= ARRAY_SIZE(key.args));
   key.op = op;
   memcpy(&key.args, args, sizeof(uint32_t) * num_args);
   key.num_args = num_args;

   struct hash_entry *entry;
   if (b->types) {
      entry = _mesa_hash_table_search(b->types, &key);
      if (entry)
         return ((struct spirv_type *)entry->data)->type;
   } else {
      b->types = _mesa_hash_table_create(NULL, non_aggregate_type_hash,
                                         non_aggregate_type_equals);
      assert(b->types);
   }

   struct spirv_type *type = CALLOC_STRUCT(spirv_type);
   if (!type)
      return 0;

   type->op = op;
   memcpy(&type->args, args, sizeof(uint32_t) * num_args);
   type->num_args = num_args;

   type->type = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->types_const_defs, 2 + num_args);
   spirv_buffer_emit_word(&b->types_const_defs, op | ((2 + num_args) << 16));
   spirv_buffer_emit_word(&b->types_const_defs, type->type);
   for (int i = 0; i < num_args; ++i)
      spirv_buffer_emit_word(&b->types_const_defs, args[i]);

   entry = _mesa_hash_table_insert(b->types, type, type);
   assert(entry);

   return ((struct spirv_type *)entry->data)->type;
}

SpvId
spirv_builder_type_void(struct spirv_builder *b)
{
   return get_type_def(b, SpvOpTypeVoid, NULL, 0);
}

SpvId
spirv_builder_type_bool(struct spirv_builder *b)
{
   return get_type_def(b, SpvOpTypeBool, NULL, 0);
}

SpvId
spirv_builder_type_int(struct spirv_builder *b, unsigned width)
{
   uint32_t args[] = { width, 1 };
   return get_type_def(b, SpvOpTypeInt, args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_type_uint(struct spirv_builder *b, unsigned width)
{
   uint32_t args[] = { width, 0 };
   return get_type_def(b, SpvOpTypeInt, args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_type_float(struct spirv_builder *b, unsigned width)
{
   uint32_t args[] = { width };
   return get_type_def(b, SpvOpTypeFloat, args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_type_image(struct spirv_builder *b, SpvId sampled_type,
                         SpvDim dim, bool depth, bool arrayed, bool ms,
                         unsigned sampled, SpvImageFormat image_format)
{
   assert(sampled < 3);
   uint32_t args[] = {
      sampled_type, dim, depth ? 1 : 0, arrayed ? 1 : 0, ms ? 1 : 0, sampled,
      image_format
   };
   return get_type_def(b, SpvOpTypeImage, args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_type_sampled_image(struct spirv_builder *b, SpvId image_type)
{
   uint32_t args[] = { image_type };
   return get_type_def(b, SpvOpTypeSampledImage, args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_type_pointer(struct spirv_builder *b,
                           SpvStorageClass storage_class, SpvId type)
{
   uint32_t args[] = { storage_class, type };
   return get_type_def(b, SpvOpTypePointer, args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_type_vector(struct spirv_builder *b, SpvId component_type,
                          unsigned component_count)
{
   assert(component_count > 1);
   uint32_t args[] = { component_type, component_count };
   return get_type_def(b, SpvOpTypeVector, args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_type_array(struct spirv_builder *b, SpvId component_type,
                         SpvId length)
{
   SpvId type = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->types_const_defs, 4);
   spirv_buffer_emit_word(&b->types_const_defs, SpvOpTypeArray | (4 << 16));
   spirv_buffer_emit_word(&b->types_const_defs, type);
   spirv_buffer_emit_word(&b->types_const_defs, component_type);
   spirv_buffer_emit_word(&b->types_const_defs, length);
   return type;
}

SpvId
spirv_builder_type_struct(struct spirv_builder *b, const SpvId member_types[],
                          size_t num_member_types)
{
   int words = 2 + num_member_types;
   SpvId type = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->types_const_defs, words);
   spirv_buffer_emit_word(&b->types_const_defs, SpvOpTypeStruct | (words << 16));
   spirv_buffer_emit_word(&b->types_const_defs, type);
   for (int i = 0; i < num_member_types; ++i)
      spirv_buffer_emit_word(&b->types_const_defs, member_types[i]);
   return type;
}

SpvId
spirv_builder_type_function(struct spirv_builder *b, SpvId return_type,
                            const SpvId parameter_types[],
                            size_t num_parameter_types)
{
   int words = 3 + num_parameter_types;
   SpvId type = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->types_const_defs, words);
   spirv_buffer_emit_word(&b->types_const_defs, SpvOpTypeFunction | (words << 16));
   spirv_buffer_emit_word(&b->types_const_defs, type);
   spirv_buffer_emit_word(&b->types_const_defs, return_type);
   for (int i = 0; i < num_parameter_types; ++i)
      spirv_buffer_emit_word(&b->types_const_defs, parameter_types[i]);
   return type;
}

struct spirv_const {
   SpvOp op, type;
   uint32_t args[8];
   size_t num_args;

   SpvId result;
};

static uint32_t
const_hash(const void *arg)
{
   const struct spirv_const *key = arg;

   uint32_t hash = _mesa_fnv32_1a_offset_bias;
   hash = _mesa_fnv32_1a_accumulate(hash, key->op);
   hash = _mesa_fnv32_1a_accumulate(hash, key->type);
   hash = _mesa_fnv32_1a_accumulate_block(hash, key->args, sizeof(uint32_t) *
                                          key->num_args);
   return hash;
}

static bool
const_equals(const void *a, const void *b)
{
   const struct spirv_const *ca = a, *cb = b;

   if (ca->op != cb->op ||
       ca->type != cb->type)
      return false;

   assert(ca->num_args == cb->num_args);
   return memcmp(ca->args, cb->args, sizeof(uint32_t) * ca->num_args) == 0;
}

static SpvId
get_const_def(struct spirv_builder *b, SpvOp op, SpvId type,
              const uint32_t args[], size_t num_args)
{
   struct spirv_const key;
   assert(num_args <= ARRAY_SIZE(key.args));
   key.op = op;
   key.type = type;
   memcpy(&key.args, args, sizeof(uint32_t) * num_args);
   key.num_args = num_args;

   struct hash_entry *entry;
   if (b->consts) {
      entry = _mesa_hash_table_search(b->consts, &key);
      if (entry)
         return ((struct spirv_const *)entry->data)->result;
   } else {
      b->consts = _mesa_hash_table_create(NULL, const_hash, const_equals);
      assert(b->consts);
   }

   struct spirv_const *cnst = CALLOC_STRUCT(spirv_const);
   if (!cnst)
      return 0;

   cnst->op = op;
   cnst->type = type;
   memcpy(&cnst->args, args, sizeof(uint32_t) * num_args);
   cnst->num_args = num_args;

   cnst->result = spirv_builder_new_id(b);
   spirv_buffer_prepare(&b->types_const_defs, 3 + num_args);
   spirv_buffer_emit_word(&b->types_const_defs, op | ((3 + num_args) << 16));
   spirv_buffer_emit_word(&b->types_const_defs, type);
   spirv_buffer_emit_word(&b->types_const_defs, cnst->result);
   for (int i = 0; i < num_args; ++i)
      spirv_buffer_emit_word(&b->types_const_defs, args[i]);

   entry = _mesa_hash_table_insert(b->consts, cnst, cnst);
   assert(entry);

   return ((struct spirv_const *)entry->data)->result;
}

SpvId
spirv_builder_const_bool(struct spirv_builder *b, bool val)
{
   return get_const_def(b, val ? SpvOpConstantTrue : SpvOpConstantFalse,
                        spirv_builder_type_bool(b), NULL, 0);
}

SpvId
spirv_builder_const_int(struct spirv_builder *b, int width, int32_t val)
{
   assert(width <= 32);
   uint32_t args[] = { val };
   return get_const_def(b, SpvOpConstant, spirv_builder_type_int(b, width),
                        args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_const_uint(struct spirv_builder *b, int width, uint32_t val)
{
   assert(width <= 32);
   uint32_t args[] = { val };
   return get_const_def(b, SpvOpConstant, spirv_builder_type_uint(b, width),
                        args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_const_float(struct spirv_builder *b, int width, float val)
{
   assert(width <= 32);
   uint32_t args[] = { u_bitcast_f2u(val) };
   return get_const_def(b, SpvOpConstant, spirv_builder_type_float(b, width),
                        args, ARRAY_SIZE(args));
}

SpvId
spirv_builder_const_composite(struct spirv_builder *b, SpvId result_type,
                              const SpvId constituents[],
                              size_t num_constituents)
{
   return get_const_def(b, SpvOpConstantComposite, result_type,
                        (const uint32_t *)constituents,
                        num_constituents);
}

SpvId
spirv_builder_emit_var(struct spirv_builder *b, SpvId type,
                       SpvStorageClass storage_class)
{
   assert(storage_class != SpvStorageClassGeneric);
   struct spirv_buffer *buf = storage_class != SpvStorageClassFunction ?
                              &b->types_const_defs : &b->instructions;

   SpvId ret = spirv_builder_new_id(b);
   spirv_buffer_prepare(buf, 4);
   spirv_buffer_emit_word(buf, SpvOpVariable | (4 << 16));
   spirv_buffer_emit_word(buf, type);
   spirv_buffer_emit_word(buf, ret);
   spirv_buffer_emit_word(buf, storage_class);
   return ret;
}

SpvId
spirv_builder_import(struct spirv_builder *b, const char *name)
{
   SpvId result = spirv_builder_new_id(b);
   size_t pos = b->imports.num_words;
   spirv_buffer_prepare(&b->imports, 2);
   spirv_buffer_emit_word(&b->imports, SpvOpExtInstImport);
   spirv_buffer_emit_word(&b->imports, result);
   int len = spirv_buffer_emit_string(&b->imports, name);
   b->imports.words[pos] |= (2 + len) << 16;
   return result;
}

size_t
spirv_builder_get_num_words(struct spirv_builder *b)
{
   const size_t header_size = 5;
   return header_size +
          b->capabilities.num_words +
          b->imports.num_words +
          b->memory_model.num_words +
          b->entry_points.num_words +
          b->exec_modes.num_words +
          b->debug_names.num_words +
          b->decorations.num_words +
          b->types_const_defs.num_words +
          b->instructions.num_words;
}

size_t
spirv_builder_get_words(struct spirv_builder *b, uint32_t *words,
                        size_t num_words)
{
   assert(num_words >= spirv_builder_get_num_words(b));

   size_t written  = 0;
   words[written++] = SpvMagicNumber;
   words[written++] = 0x00010000;
   words[written++] = 0;
   words[written++] = b->prev_id + 1;
   words[written++] = 0;

   const struct spirv_buffer *buffers[] = {
      &b->capabilities,
      &b->imports,
      &b->memory_model,
      &b->entry_points,
      &b->exec_modes,
      &b->debug_names,
      &b->decorations,
      &b->types_const_defs,
      &b->instructions
   };

   for (int i = 0; i < ARRAY_SIZE(buffers); ++i) {
      const struct spirv_buffer *buffer = buffers[i];
      for (int j = 0; j < buffer->num_words; ++j)
         words[written++] = buffer->words[j];
   }

   assert(written == spirv_builder_get_num_words(b));
   return written;
}