// Integer simple_mul bytecode loading and input/output processes

#include "iree/base/api.h"
#include "iree/hal/api.h"

// Compiled module embedded here to avoid file IO:
#include "samples/simple_vec_mul/simple_int_mul_bytecode_module_dylib_c.h"

const iree_const_byte_span_t load_bytecode_module_data() {
  const struct iree_file_toc_t *module_file_toc =
      samples_simple_vec_mul_simple_int_mul_bytecode_module_dylib_create();
  return iree_make_const_byte_span(module_file_toc->data,
                                   module_file_toc->size);
}

iree_status_t prepare_input_hal_buffer_views(
    iree_hal_device_t *device, const int buffer_length, void **arg0_buffer,
    void **arg1_buffer, iree_hal_buffer_view_t **arg0_buffer_view,
    iree_hal_buffer_view_t **arg1_buffer_view) {
  iree_status_t result = iree_ok_status();
  *arg0_buffer =
      iree_aligned_alloc(sizeof(uint32_t), sizeof(int32_t) * buffer_length);
  if (*arg0_buffer == NULL) {
    result = iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED);
  }
  if (iree_status_is_ok(result)) {
    *arg1_buffer =
        iree_aligned_alloc(sizeof(uint32_t), sizeof(int32_t) * buffer_length);
    if (*arg1_buffer == NULL) {
      result = iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED);
    }
  }

  // Populate initial values
  // arg0 = 0, 0, 1, 1,..., 511
  // arg1 = 0, 1, 2, 3,..., 1023
  for (int i = 0; i < buffer_length; ++i) {
    ((int32_t *)*arg0_buffer)[i] = i >> 1;
    ((int32_t *)*arg1_buffer)[i] = i;
  }
  // Wrap buffers in shaped buffer views.
  // The buffers can be mapped on the CPU and that can also be used
  // on the device. Not all devices support this, but the ones we have now do.
  iree_hal_dim_t shape[1] = {buffer_length};

  iree_hal_memory_type_t input_memory_type =
      IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE;

  result = iree_hal_buffer_view_wrap_or_clone_heap_buffer(
      iree_hal_device_allocator(device), shape, IREE_ARRAYSIZE(shape),
      IREE_HAL_ELEMENT_TYPE_SINT_32, input_memory_type,
      IREE_HAL_MEMORY_ACCESS_READ, IREE_HAL_BUFFER_USAGE_ALL,
      iree_make_byte_span(*arg0_buffer, sizeof(int32_t) * buffer_length),
      iree_allocator_null(), arg0_buffer_view);

  if (iree_status_is_ok(result)) {
    result = iree_hal_buffer_view_wrap_or_clone_heap_buffer(
        iree_hal_device_allocator(device), shape, IREE_ARRAYSIZE(shape),
        IREE_HAL_ELEMENT_TYPE_SINT_32, input_memory_type,
        IREE_HAL_MEMORY_ACCESS_READ, IREE_HAL_BUFFER_USAGE_ALL,
        iree_make_byte_span(*arg1_buffer, sizeof(int32_t) * buffer_length),
        iree_allocator_null(), arg1_buffer_view);
  }
  return result;
}

iree_status_t check_output_data(iree_hal_buffer_mapping_t *mapped_memory) {
  iree_status_t result = iree_ok_status();
  for (int i = 0; i < mapped_memory->contents.data_length / sizeof(int32_t);
       ++i) {
    if (((const int32_t *)mapped_memory->contents.data)[i] != (i >> 1) * i) {
      result = iree_make_status(IREE_STATUS_UNKNOWN, "result mismatches");
      break;
    }
  }
  return result;
}