// Copyright 2020 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>

#include <memory>

#include "iree/base/api.h"
#include "iree/vm/buffer.h"
#include "iree/vm/context.h"
#include "iree/vm/instance.h"
#include "iree/vm/module.h"
#include "iree/vm/native_module.h"
#include "iree/vm/native_module_cc.h"
#include "iree/vm/ref.h"
#include "iree/vm/shims.h"
#include "iree/vm/stack.h"

// Wrapper for calling the import functions with type (i32)->i32.
// NOTE: we should have some common ones prebuilt or can generate and rely on
// LTO to strip duplicates across the entire executable.
// TODO(benvanik): generate/export these shims/call functions in stack.h.
static iree_status_t call_import_i32_i32(iree_vm_stack_t* stack,
                                         const iree_vm_function_t* import,
                                         int32_t arg0, int32_t* out_ret0) {
  iree_vm_function_call_t call;
  call.function = *import;
  call.arguments = iree_make_byte_span(&arg0, sizeof(arg0));
  call.results = iree_make_byte_span(out_ret0, sizeof(*out_ret0));
  return import->module->begin_call(import->module, stack, call);
}

typedef iree_status_t (*call_i32_i32_t)(iree_vm_stack_t* stack,
                                        void* module_ptr, void* module_state,
                                        int32_t arg0, int32_t* out_ret0);

// Wrapper for calling a |target_fn| C function from the VM ABI.
// It's optional to bounce through like this; if the function can more
// efficiently directly access the arguments from the |call| then it can do so.
// This approach is most useful when the function may also be exported/used by
// non-VM code or may be internally referenced using a target-specific ABI.
// TODO(benvanik): generate/export these shims/call functions in stack.h.
static iree_status_t call_shim_i32_i32(iree_vm_stack_t* stack,
                                       iree_vm_native_function_flags_t flags,
                                       iree_byte_span_t args_storage,
                                       iree_byte_span_t rets_storage,
                                       call_i32_i32_t target_fn, void* module,
                                       void* module_state) {
  // We can use structs to allow compiler-controlled indexing optimizations,
  // though this won't work for variadic cases.
  // TODO(benvanik): packed attributes.
  typedef struct {
    int32_t arg0;
  } args_t;
  typedef struct {
    int32_t ret0;
  } results_t;

  const args_t* args = (const args_t*)args_storage.data;
  results_t* results = (results_t*)rets_storage.data;

  // For simple cases like this (zero or 1 result) we can tail-call.
  return target_fn(stack, module, module_state, args->arg0, &results->ret0);
}

//===----------------------------------------------------------------------===//
// module_a
//===----------------------------------------------------------------------===//
// This simple stateless module exports two functions that can be imported by
// other modules or called directly by the user. When no imports, custom types,
// or per-context state is required this simplifies module definitions.
//
// module_b below imports these functions and demonstrates a more complex module
// with state.

typedef struct module_a_t module_a_t;
typedef struct module_a_state_t module_a_state_t;

// vm.import private @module_a.add_1(%arg0 : i32) -> i32
static iree_status_t module_a_add_1(iree_vm_stack_t* stack, module_a_t* module,
                                    module_a_state_t* module_state,
                                    int32_t arg0, int32_t* out_ret0) {
  // Add 1 to arg0 and return.
  *out_ret0 = arg0 + 1;
  return iree_ok_status();
}

// vm.import private @module_a.sub_1(%arg0 : i32) -> i32
static iree_status_t module_a_sub_1(iree_vm_stack_t* stack, module_a_t* module,
                                    module_a_state_t* module_state,
                                    int32_t arg0, int32_t* out_ret0) {
  // Sub 1 to arg0 and return. Fail if < 0.
  *out_ret0 = arg0 - 1;
  return iree_ok_status();
}

static const iree_vm_native_export_descriptor_t module_a_exports_[] = {
    {IREE_SV("add_1"), IREE_SV("0i_i"), 0, NULL},
    {IREE_SV("sub_1"), IREE_SV("0i_i"), 0, NULL},
};
static const iree_vm_native_function_ptr_t module_a_funcs_[] = {
    {(iree_vm_native_function_shim_t)call_shim_i32_i32,
     (iree_vm_native_function_target_t)module_a_add_1},
    {(iree_vm_native_function_shim_t)call_shim_i32_i32,
     (iree_vm_native_function_target_t)module_a_sub_1},
};
static_assert(IREE_ARRAYSIZE(module_a_funcs_) ==
                  IREE_ARRAYSIZE(module_a_exports_),
              "function pointer table must be 1:1 with exports");
static const iree_vm_native_module_descriptor_t module_a_descriptor_ = {
    /*name=*/IREE_SV("module_a"),
    /*version=*/0,
    /*attr_count=*/0,
    /*attrs=*/NULL,
    /*dependency_count=*/0,
    /*dependencies=*/NULL,
    /*import_count=*/0,
    /*imports=*/NULL,
    /*export_count=*/IREE_ARRAYSIZE(module_a_exports_),
    /*exports=*/module_a_exports_,
    /*function_count=*/IREE_ARRAYSIZE(module_a_funcs_),
    /*functions=*/module_a_funcs_,
};

static iree_status_t module_a_create(iree_vm_instance_t* instance,
                                     iree_allocator_t allocator,
                                     iree_vm_module_t** out_module) {
  // NOTE: this module has neither shared or per-context module state.
  iree_vm_module_t interface;
  IREE_RETURN_IF_ERROR(iree_vm_module_initialize(&interface, NULL));
  return iree_vm_native_module_create(&interface, &module_a_descriptor_,
                                      instance, allocator, out_module);
}

//===----------------------------------------------------------------------===//
// module_b
//===----------------------------------------------------------------------===//
// A more complex module that holds state for resolved types (shared across
// all instances), imported functions (stored per-context), per-context user
// data, and reflection metadata.

typedef struct module_b_t module_b_t;
typedef struct module_b_state_t module_b_state_t;

// Stores shared state across all instances of the module.
// This should generally be treated as read-only and if mutation is possible
// then users must synchronize themselves.
typedef struct module_b_t {
  // Allocator the module must be freed with and that can be used for any other
  // shared dynamic allocations.
  iree_allocator_t allocator;
  // Resolved types; these never change once queried and are safe to store on
  // the shared structure to avoid needing to look them up again.
  iree_vm_ref_type_t types[1];
} module_b_t;

// Stores per-context state; at the minimum imports, but possibly other user
// state data. No synchronization is required as the VM will not call functions
// with the same state from multiple threads concurrently.
typedef struct module_b_state_t {
  // Allocator the state must be freed with and that can be used for any other
  // per-context dynamic allocations.
  iree_allocator_t allocator;
  // Resolved import functions matching 1:1 with the module import descriptors.
  iree_vm_function_t imports[2];
  // Example user data stored per-state.
  int counter;
} module_b_state_t;

// Frees the shared module; by this point all per-context states have been
// freed and no more shared data is required.
static void IREE_API_PTR module_b_destroy(void* self) {
  module_b_t* module = (module_b_t*)self;
  iree_allocator_free(module->allocator, module);
}

// Allocates per-context state, which stores resolved import functions and any
// other non-shared user state.
static iree_status_t IREE_API_PTR
module_b_alloc_state(void* self, iree_allocator_t allocator,
                     iree_vm_module_state_t** out_module_state) {
  module_b_state_t* state = NULL;
  IREE_RETURN_IF_ERROR(
      iree_allocator_malloc(allocator, sizeof(*state), (void**)&state));
  memset(state, 0, sizeof(*state));
  state->allocator = allocator;
  *out_module_state = (iree_vm_module_state_t*)state;
  return iree_ok_status();
}

// Frees the per-context state.
static void IREE_API_PTR
module_b_free_state(void* self, iree_vm_module_state_t* module_state) {
  module_b_state_t* state = (module_b_state_t*)module_state;
  iree_allocator_free(state->allocator, state);
}

// Clones the module state and retains resources by-reference.
static iree_status_t IREE_API_PTR module_b_fork_state(
    void* self, iree_vm_module_state_t* parent_state,
    iree_allocator_t allocator, iree_vm_module_state_t** out_child_state) {
  module_b_state_t* child_state = NULL;
  IREE_RETURN_IF_ERROR(iree_allocator_malloc(allocator, sizeof(*child_state),
                                             (void**)&child_state));
  // Copy resolved imports and the counter value.
  memcpy(child_state, parent_state, sizeof(*child_state));
  // Reassign the allocator used.
  child_state->allocator = allocator;
  *out_child_state = (iree_vm_module_state_t*)child_state;
  return iree_ok_status();
}

// Called once per import function so the module can store the function ref.
static iree_status_t IREE_API_PTR module_b_resolve_import(
    void* self, iree_vm_module_state_t* module_state, iree_host_size_t ordinal,
    const iree_vm_function_t* function,
    const iree_vm_function_signature_t* signature) {
  module_b_state_t* state = (module_b_state_t*)module_state;
  state->imports[ordinal] = *function;
  return iree_ok_status();
}

// Our actual function. Here we directly access the registers but one could also
// use this as a trampoline into user code with a native signature (such as
// fetching the args, calling the function as a normal C function, and stashing
// back the results).
//
// vm.import private @module_b.entry(%arg0 : i32) -> i32
static iree_status_t module_b_entry(iree_vm_stack_t* stack, module_b_t* module,
                                    module_b_state_t* module_state,
                                    int32_t arg0, int32_t* out_ret0) {
  // NOTE: if we needed to use ref types here we have them under module->types.
  assert(module->types[0]);

  // Call module_a.add_1.
  IREE_RETURN_IF_ERROR(
      call_import_i32_i32(stack, &module_state->imports[0], arg0, &arg0));

  // Increment per-context state (persists across calls). No need for a mutex as
  // only one thread can be using the per-context state at a time.
  module_state->counter += arg0;
  int32_t ret0 = module_state->counter;

  // Call module_a.sub_1.
  IREE_RETURN_IF_ERROR(
      call_import_i32_i32(stack, &module_state->imports[1], ret0, &ret0));

  *out_ret0 = ret0;
  return iree_ok_status();
}

// Table of exported function pointers. Note that this table could be read-only
// (like here) or shared/per-context to allow exposing different functions based
// on versions, access rights, etc.
static const iree_vm_native_function_ptr_t module_b_funcs_[] = {
    {(iree_vm_native_function_shim_t)call_shim_i32_i32,
     (iree_vm_native_function_target_t)module_b_entry},
};

static const iree_vm_native_import_descriptor_t module_b_imports_[] = {
    {IREE_VM_NATIVE_IMPORT_REQUIRED, IREE_SV("module_a.add_1")},
    {IREE_VM_NATIVE_IMPORT_REQUIRED, IREE_SV("module_a.sub_1")},
};
static_assert(IREE_ARRAYSIZE(((module_b_state_t*)NULL)->imports) ==
                  IREE_ARRAYSIZE(module_b_imports_),
              "import storage must be able to hold all imports");
static const iree_string_pair_t module_b_entry_attrs_[] = {
    {{IREE_SV("key1")}, {IREE_SV("value1")}},
};
static const iree_vm_native_export_descriptor_t module_b_exports_[] = {
    {IREE_SV("entry"), IREE_SV("0i_i"), IREE_ARRAYSIZE(module_b_entry_attrs_),
     module_b_entry_attrs_},
};
static_assert(IREE_ARRAYSIZE(module_b_funcs_) ==
                  IREE_ARRAYSIZE(module_b_exports_),
              "function pointer table must be 1:1 with exports");
static const iree_vm_native_module_descriptor_t module_b_descriptor_ = {
    /*name=*/IREE_SV("module_b"),
    /*version=*/0,
    /*attr_count=*/0,
    /*attrs=*/NULL,
    /*dependency_count=*/0,
    /*dependencies=*/NULL,
    /*import_count=*/IREE_ARRAYSIZE(module_b_imports_),
    /*imports=*/module_b_imports_,
    /*export_count=*/IREE_ARRAYSIZE(module_b_exports_),
    /*exports=*/module_b_exports_,
    /*function_count=*/IREE_ARRAYSIZE(module_b_funcs_),
    /*functions=*/module_b_funcs_,
};

static iree_status_t module_b_create(iree_vm_instance_t* instance,
                                     iree_allocator_t allocator,
                                     iree_vm_module_t** out_module) {
  // Allocate shared module state.
  module_b_t* module = NULL;
  IREE_RETURN_IF_ERROR(
      iree_allocator_malloc(allocator, sizeof(*module), (void**)&module));
  memset(module, 0, sizeof(*module));
  module->allocator = allocator;

  // Resolve types used by the module once so that we can share it across all
  // instances of the module. Depending on the types here can be somewhat risky
  // as it can lead to ordering issues. If possible resolving types on module
  // state is better as all dependent modules are guaranteed to have been
  // loaded.
  module->types[0] = iree_vm_instance_lookup_type(
      instance, iree_make_cstring_view("vm.buffer"));
  if (!module->types[0]) {
    iree_allocator_free(allocator, module);
    return iree_make_status(
        IREE_STATUS_NOT_FOUND,
        "required type vm.buffer not registered with the type system");
  }

  // Setup the interface with the functions we implement ourselves. Any function
  // we omit will be handled by the base native module.
  iree_vm_module_t interface;
  iree_status_t status = iree_vm_module_initialize(&interface, module);
  if (!iree_status_is_ok(status)) {
    iree_allocator_free(allocator, module);
    return status;
  }
  interface.destroy = module_b_destroy;
  interface.alloc_state = module_b_alloc_state;
  interface.free_state = module_b_free_state;
  interface.fork_state = module_b_fork_state;
  interface.resolve_import = module_b_resolve_import;
  return iree_vm_native_module_create(&interface, &module_b_descriptor_,
                                      instance, allocator, out_module);
}

//===----------------------------------------------------------------------===//
// module_c_align
//===----------------------------------------------------------------------===//
// C implementation to test alignment-sensitive parameter unpacking and result
// packing patterns. This module implements the same functions as module_cpp
// using the C native module API.

typedef struct module_c_align_t {
  iree_allocator_t allocator;
} module_c_align_t;

typedef struct module_c_align_state_t {
  iree_allocator_t allocator;
  int counter;
} module_c_align_state_t;

static void IREE_API_PTR module_c_align_destroy(void* self) {
  module_c_align_t* module = (module_c_align_t*)self;
  iree_allocator_free(module->allocator, module);
}

static iree_status_t IREE_API_PTR
module_c_align_alloc_state(void* self, iree_allocator_t allocator,
                           iree_vm_module_state_t** out_module_state) {
  module_c_align_state_t* state = NULL;
  IREE_RETURN_IF_ERROR(
      iree_allocator_malloc(allocator, sizeof(*state), (void**)&state));
  memset(state, 0, sizeof(*state));
  state->allocator = allocator;
  *out_module_state = (iree_vm_module_state_t*)state;
  return iree_ok_status();
}

static void IREE_API_PTR
module_c_align_free_state(void* self, iree_vm_module_state_t* module_state) {
  module_c_align_state_t* state = (module_c_align_state_t*)module_state;
  iree_allocator_free(state->allocator, state);
}

// (i32) -> i32 - basic entry point.
IREE_VM_ABI_EXPORT(module_c_align_entry, module_c_align_state_t, i, i) {
  state->counter += args->i0 + 1;
  rets->i0 = state->counter - 1;
  return iree_ok_status();
}

// (i32, ref<buffer>) -> i32 - ref after i32 triggers alignment.
IREE_VM_ABI_EXPORT(module_c_align_mixed_i32_ref, module_c_align_state_t, ir,
                   i) {
  iree_host_size_t buf_len =
      iree_vm_ref_is_null(&args->r1)
          ? 0
          : iree_vm_buffer_length((iree_vm_buffer_t*)args->r1.ptr);
  rets->i0 = args->i0 + (int32_t)buf_len;
  return iree_ok_status();
}

// (ref<buffer>, i32, ref<buffer>) -> i32 - ref/i32/ref pattern.
IREE_VM_ABI_EXPORT(module_c_align_mixed_ref_i32_ref, module_c_align_state_t,
                   rir, i) {
  iree_host_size_t len1 =
      iree_vm_ref_is_null(&args->r0)
          ? 0
          : iree_vm_buffer_length((iree_vm_buffer_t*)args->r0.ptr);
  iree_host_size_t len2 =
      iree_vm_ref_is_null(&args->r2)
          ? 0
          : iree_vm_buffer_length((iree_vm_buffer_t*)args->r2.ptr);
  rets->i0 = (int32_t)len1 + args->i1 + (int32_t)len2;
  return iree_ok_status();
}

// (i32, i64) -> i64 - i64 after i32 triggers alignment.
IREE_VM_ABI_EXPORT(module_c_align_mixed_i32_i64, module_c_align_state_t, iI,
                   I) {
  rets->i0 = args->i0 + args->i1;
  return iree_ok_status();
}

// (i32, i32, i32, ref<buffer>) -> i32 - ref after 3 i32s (12 bytes).
IREE_VM_ABI_EXPORT(module_c_align_mixed_i32x3_ref, module_c_align_state_t, iiir,
                   i) {
  iree_host_size_t buf_len =
      iree_vm_ref_is_null(&args->r3)
          ? 0
          : iree_vm_buffer_length((iree_vm_buffer_t*)args->r3.ptr);
  rets->i0 = args->i0 + args->i1 + args->i2 + (int32_t)buf_len;
  return iree_ok_status();
}

// (i64, i32) -> i32 - i32 after i64.
IREE_VM_ABI_EXPORT(module_c_align_mixed_i64_i32, module_c_align_state_t, Ii,
                   i) {
  rets->i0 = (int32_t)args->i0 + args->i1;
  return iree_ok_status();
}

// (i32, i64, i32) -> i64 - i64 sandwiched between i32s.
IREE_VM_ABI_EXPORT(module_c_align_mixed_i32_i64_i32, module_c_align_state_t,
                   iIi, I) {
  rets->i0 = args->i0 + args->i1 + args->i2;
  return iree_ok_status();
}

// Exports must be sorted alphabetically by name.
static const iree_vm_native_export_descriptor_t module_c_align_exports_[] = {
    {IREE_SV("entry"), IREE_SV("0i_i"), 0, NULL},
    {IREE_SV("mixed_i32_i64"), IREE_SV("0iI_I"), 0, NULL},
    {IREE_SV("mixed_i32_i64_i32"), IREE_SV("0iIi_I"), 0, NULL},
    {IREE_SV("mixed_i32_ref"), IREE_SV("0ir_i"), 0, NULL},
    {IREE_SV("mixed_i32x3_ref"), IREE_SV("0iiir_i"), 0, NULL},
    {IREE_SV("mixed_i64_i32"), IREE_SV("0Ii_i"), 0, NULL},
    {IREE_SV("mixed_ref_i32_ref"), IREE_SV("0rir_i"), 0, NULL},
};

static const iree_vm_native_function_ptr_t module_c_align_funcs_[] = {
    {(iree_vm_native_function_shim_t)iree_vm_shim_i_i,
     (iree_vm_native_function_target_t)module_c_align_entry},
    {(iree_vm_native_function_shim_t)iree_vm_shim_iI_I,
     (iree_vm_native_function_target_t)module_c_align_mixed_i32_i64},
    {(iree_vm_native_function_shim_t)iree_vm_shim_iIi_I,
     (iree_vm_native_function_target_t)module_c_align_mixed_i32_i64_i32},
    {(iree_vm_native_function_shim_t)iree_vm_shim_ir_i,
     (iree_vm_native_function_target_t)module_c_align_mixed_i32_ref},
    {(iree_vm_native_function_shim_t)iree_vm_shim_iiir_i,
     (iree_vm_native_function_target_t)module_c_align_mixed_i32x3_ref},
    {(iree_vm_native_function_shim_t)iree_vm_shim_Ii_i,
     (iree_vm_native_function_target_t)module_c_align_mixed_i64_i32},
    {(iree_vm_native_function_shim_t)iree_vm_shim_rir_i,
     (iree_vm_native_function_target_t)module_c_align_mixed_ref_i32_ref},
};

static_assert(IREE_ARRAYSIZE(module_c_align_funcs_) ==
                  IREE_ARRAYSIZE(module_c_align_exports_),
              "function pointer table must be 1:1 with exports");

static const iree_vm_native_module_descriptor_t module_c_align_descriptor_ = {
    /*name=*/IREE_SV("module_c"),
    /*version=*/0,
    /*attr_count=*/0,
    /*attrs=*/NULL,
    /*dependency_count=*/0,
    /*dependencies=*/NULL,
    /*import_count=*/0,
    /*imports=*/NULL,
    /*export_count=*/IREE_ARRAYSIZE(module_c_align_exports_),
    /*exports=*/module_c_align_exports_,
    /*function_count=*/IREE_ARRAYSIZE(module_c_align_funcs_),
    /*functions=*/module_c_align_funcs_,
};

static iree_status_t module_c_align_create(iree_vm_instance_t* instance,
                                           iree_allocator_t allocator,
                                           iree_vm_module_t** out_module) {
  module_c_align_t* module = NULL;
  IREE_RETURN_IF_ERROR(
      iree_allocator_malloc(allocator, sizeof(*module), (void**)&module));
  memset(module, 0, sizeof(*module));
  module->allocator = allocator;

  iree_vm_module_t interface;
  iree_status_t status = iree_vm_module_initialize(&interface, module);
  if (!iree_status_is_ok(status)) {
    iree_allocator_free(allocator, module);
    return status;
  }
  interface.destroy = module_c_align_destroy;
  interface.alloc_state = module_c_align_alloc_state;
  interface.free_state = module_c_align_free_state;
  return iree_vm_native_module_create(&interface, &module_c_align_descriptor_,
                                      instance, allocator, out_module);
}

//===----------------------------------------------------------------------===//
// module_cpp
//===----------------------------------------------------------------------===//
// C++ implementation using native_module_cc.h to test alignment-sensitive
// parameter unpacking and result packing patterns. This module implements
// functions with mixed type signatures that can trigger alignment issues.

namespace {

struct ModuleCppState final {
  int counter = 0;

  // Same signature as C module_b.entry: (i32) -> i32
  iree::StatusOr<int32_t> Entry(int32_t arg0) {
    counter += arg0 + 1;  // Equivalent to add_1
    return counter - 1;   // Equivalent to sub_1
  }

  // (i32, ref<buffer>) -> i32 - tests ref after i32
  iree::StatusOr<int32_t> MixedI32Ref(int32_t arg0,
                                      iree::vm::ref<iree_vm_buffer_t> buf) {
    // Return arg0 + buffer length (or 0 if null).
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf.get()) : 0;
    return arg0 + static_cast<int32_t>(buf_len);
  }

  // (ref<buffer>, i32, ref<buffer>) -> i32 - tests ref after i32 after ref
  iree::StatusOr<int32_t> MixedRefI32Ref(iree::vm::ref<iree_vm_buffer_t> buf1,
                                         int32_t arg0,
                                         iree::vm::ref<iree_vm_buffer_t> buf2) {
    iree_host_size_t len1 = buf1 ? iree_vm_buffer_length(buf1.get()) : 0;
    iree_host_size_t len2 = buf2 ? iree_vm_buffer_length(buf2.get()) : 0;
    return static_cast<int32_t>(len1) + arg0 + static_cast<int32_t>(len2);
  }

  // (i32, i64) -> i64 - tests i64 after i32
  iree::StatusOr<int64_t> MixedI32I64(int32_t a, int64_t b) { return a + b; }

  // (i32, i32, i32, ref<buffer>) -> i32 - tests ref after 3 i32s (12 bytes)
  iree::StatusOr<int32_t> MixedI32x3Ref(int32_t a, int32_t b, int32_t c,
                                        iree::vm::ref<iree_vm_buffer_t> buf) {
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf.get()) : 0;
    return a + b + c + static_cast<int32_t>(buf_len);
  }

  // (i64, i32) -> i32 - tests i32 after i64
  iree::StatusOr<int32_t> MixedI64I32(int64_t a, int32_t b) {
    return static_cast<int32_t>(a) + b;
  }

  // (i32, i64, i32) -> i64 - tests i64 sandwiched between i32s
  iree::StatusOr<int64_t> MixedI32I64I32(int32_t a, int64_t b, int32_t c) {
    return a + b + c;
  }

  //===--------------------------------------------------------------------===//
  // Borrowed pointer (T*) variants - tests ParamUnpack<T*> alignment
  //===--------------------------------------------------------------------===//
  // These use raw pointers instead of vm::ref<T> to test the borrowed pointer
  // unpacking path which had a missing alignment bug.

  // (i32, buffer*) -> i32 - borrowed ref after i32 (4 bytes, needs alignment)
  iree::StatusOr<int32_t> BorrowedI32Ref(int32_t arg0, iree_vm_buffer_t* buf) {
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf) : 0;
    return arg0 + static_cast<int32_t>(buf_len);
  }

  // (i32, i32, buffer*) -> i32 - borrowed ref after 8 bytes (already aligned)
  iree::StatusOr<int32_t> BorrowedI32x2Ref(int32_t a, int32_t b,
                                           iree_vm_buffer_t* buf) {
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf) : 0;
    return a + b + static_cast<int32_t>(buf_len);
  }

  // (i32, i32, i32, buffer*) -> i32 - borrowed ref after 12 bytes (needs align)
  // This is the pattern that crashed in e2e tests!
  iree::StatusOr<int32_t> BorrowedI32x3Ref(int32_t a, int32_t b, int32_t c,
                                           iree_vm_buffer_t* buf) {
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf) : 0;
    return a + b + c + static_cast<int32_t>(buf_len);
  }

  // (i32, i32, i32, i32, buffer*) -> i32 - borrowed ref after 16 bytes
  // (aligned)
  iree::StatusOr<int32_t> BorrowedI32x4Ref(int32_t a, int32_t b, int32_t c,
                                           int32_t d, iree_vm_buffer_t* buf) {
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf) : 0;
    return a + b + c + d + static_cast<int32_t>(buf_len);
  }

  // (i32, i32, i32, i32, i32, buffer*) -> i32 - after 20 bytes (needs align)
  iree::StatusOr<int32_t> BorrowedI32x5Ref(int32_t a, int32_t b, int32_t c,
                                           int32_t d, int32_t e,
                                           iree_vm_buffer_t* buf) {
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf) : 0;
    return a + b + c + d + e + static_cast<int32_t>(buf_len);
  }

  // (buffer*, i32, buffer*) -> i32 - borrowed ref/i32/ref pattern
  iree::StatusOr<int32_t> BorrowedRefI32Ref(iree_vm_buffer_t* buf1,
                                            int32_t arg0,
                                            iree_vm_buffer_t* buf2) {
    iree_host_size_t len1 = buf1 ? iree_vm_buffer_length(buf1) : 0;
    iree_host_size_t len2 = buf2 ? iree_vm_buffer_length(buf2) : 0;
    return static_cast<int32_t>(len1) + arg0 + static_cast<int32_t>(len2);
  }

  // (buffer*, buffer*) -> i32 - two consecutive borrowed refs
  iree::StatusOr<int32_t> BorrowedRefRef(iree_vm_buffer_t* buf1,
                                         iree_vm_buffer_t* buf2) {
    iree_host_size_t len1 = buf1 ? iree_vm_buffer_length(buf1) : 0;
    iree_host_size_t len2 = buf2 ? iree_vm_buffer_length(buf2) : 0;
    return static_cast<int32_t>(len1) + static_cast<int32_t>(len2);
  }

  // (i64, buffer*) -> i32 - borrowed ref after i64 (already aligned)
  iree::StatusOr<int32_t> BorrowedI64Ref(int64_t a, iree_vm_buffer_t* buf) {
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf) : 0;
    return static_cast<int32_t>(a) + static_cast<int32_t>(buf_len);
  }

  // (i32, i64, buffer*) -> i32 - borrowed ref after i32+i64
  iree::StatusOr<int32_t> BorrowedI32I64Ref(int32_t a, int64_t b,
                                            iree_vm_buffer_t* buf) {
    iree_host_size_t buf_len = buf ? iree_vm_buffer_length(buf) : 0;
    return a + static_cast<int32_t>(b) + static_cast<int32_t>(buf_len);
  }

  // (buffer*, i64, buffer*) -> i32 - i64 between two borrowed refs
  iree::StatusOr<int32_t> BorrowedRefI64Ref(iree_vm_buffer_t* buf1, int64_t a,
                                            iree_vm_buffer_t* buf2) {
    iree_host_size_t len1 = buf1 ? iree_vm_buffer_length(buf1) : 0;
    iree_host_size_t len2 = buf2 ? iree_vm_buffer_length(buf2) : 0;
    return static_cast<int32_t>(len1) + static_cast<int32_t>(a) +
           static_cast<int32_t>(len2);
  }
};

static const iree::vm::NativeFunction<ModuleCppState> kModuleCppFunctions[] = {
    // Owned ref (vm::ref<T>) functions.
    iree::vm::MakeNativeFunction("entry", &ModuleCppState::Entry),
    iree::vm::MakeNativeFunction("mixed_i32_ref", &ModuleCppState::MixedI32Ref),
    iree::vm::MakeNativeFunction("mixed_ref_i32_ref",
                                 &ModuleCppState::MixedRefI32Ref),
    iree::vm::MakeNativeFunction("mixed_i32_i64", &ModuleCppState::MixedI32I64),
    iree::vm::MakeNativeFunction("mixed_i32x3_ref",
                                 &ModuleCppState::MixedI32x3Ref),
    iree::vm::MakeNativeFunction("mixed_i64_i32", &ModuleCppState::MixedI64I32),
    iree::vm::MakeNativeFunction("mixed_i32_i64_i32",
                                 &ModuleCppState::MixedI32I64I32),
    // Borrowed pointer (T*) functions - tests ParamUnpack<T*> alignment.
    iree::vm::MakeNativeFunction("borrowed_i32_ref",
                                 &ModuleCppState::BorrowedI32Ref),
    iree::vm::MakeNativeFunction("borrowed_i32x2_ref",
                                 &ModuleCppState::BorrowedI32x2Ref),
    iree::vm::MakeNativeFunction("borrowed_i32x3_ref",
                                 &ModuleCppState::BorrowedI32x3Ref),
    iree::vm::MakeNativeFunction("borrowed_i32x4_ref",
                                 &ModuleCppState::BorrowedI32x4Ref),
    iree::vm::MakeNativeFunction("borrowed_i32x5_ref",
                                 &ModuleCppState::BorrowedI32x5Ref),
    iree::vm::MakeNativeFunction("borrowed_ref_i32_ref",
                                 &ModuleCppState::BorrowedRefI32Ref),
    iree::vm::MakeNativeFunction("borrowed_ref_ref",
                                 &ModuleCppState::BorrowedRefRef),
    iree::vm::MakeNativeFunction("borrowed_i64_ref",
                                 &ModuleCppState::BorrowedI64Ref),
    iree::vm::MakeNativeFunction("borrowed_i32_i64_ref",
                                 &ModuleCppState::BorrowedI32I64Ref),
    iree::vm::MakeNativeFunction("borrowed_ref_i64_ref",
                                 &ModuleCppState::BorrowedRefI64Ref),
};

class ModuleCpp final : public iree::vm::NativeModule<ModuleCppState> {
 public:
  using iree::vm::NativeModule<ModuleCppState>::NativeModule;

 protected:
  iree::StatusOr<std::unique_ptr<ModuleCppState>> CreateState(
      iree_allocator_t allocator) override {
    return std::make_unique<ModuleCppState>();
  }

  iree::StatusOr<std::unique_ptr<ModuleCppState>> ForkState(
      ModuleCppState* parent_state, iree_allocator_t allocator) override {
    auto child_state = std::make_unique<ModuleCppState>();
    child_state->counter = parent_state->counter;
    return child_state;
  }
};

}  // namespace

static iree_status_t module_cpp_create(iree_vm_instance_t* instance,
                                       iree_allocator_t allocator,
                                       iree_vm_module_t** out_module) {
  auto module = std::make_unique<ModuleCpp>(
      "module_cpp", /*version=*/0, instance, allocator,
      iree::span<const iree::vm::NativeFunction<ModuleCppState>>{
          kModuleCppFunctions});
  *out_module = module.release()->interface();
  return iree_ok_status();
}
