// Copyright 2019 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

#ifndef IREE_VM_NATIVE_MODULE_CC_H_
#define IREE_VM_NATIVE_MODULE_CC_H_

#include <cstring>
#include <functional>
#include <memory>

#include "iree/base/api.h"
#include "iree/base/internal/span.h"
#include "iree/vm/instance.h"
#include "iree/vm/module.h"
#include "iree/vm/native_module_packing.h"  // IWYU pragma: export
#include "iree/vm/stack.h"

#ifndef __cplusplus
#error "This header is meant for use with C++ module implementations."
#endif  // __cplusplus

namespace iree {
namespace vm {

// A native module as exported to the VM dynamic module linking API.
// This allows easy wrapping of C++ module implementations and removes a
// majority of the boilerplate required with marshaling args/results out/in of
// the VM via the ABI.
//
// Functions are defined on the State type as member functions returning either
// Status or StatusOr. Arguments are passed as primitive types (int32_t),
// wrapped ref objects (vm::ref<my_type_t>&), or some nesting of std::array,
// std::tuple, and std::span to match fixed-length arrays of the same type,
// tuples of mixed types, or dynamic arrays (variadic arguments). Results may be
// returned as either their type or an std::tuple/std::array of types.
//
// Usage:
//   // Per-context module state that must only be thread-compatible.
//   // Define
//   struct MyState final {
//     StatusOr<std::tuple<int32_t, int32_t>> MyMethod1(vm::ref<my_type_t> t);
//   };
//
//   // Table of functions mapped to their name in the IR.
//   static const vm::NativeFunction<MyState> kMyFunctions[] = {
//     vm::MakeNativeFunction("my_method_1", &MyState::MyMethod1),
//   };
//
//   // The outer module wrapper shared across contexts.
//   // Must be thread-safe.
//   struct MyModule : public NativeModule<MyState> {
//     StatusOr<std::unique_ptr<MyState>> CreateState(iree_allocator_t) {
//       // You could pass in thread-safe shared resources to MyState.
//       return std::make_unique<MyState>();
//     }
//   };
//
//   // Creates the module and exposes it as a C interface.
//   // Ownership transfers to the caller.
//   iree_vm_module_t* create_my_module(iree_allocator_t allocator) {
//     return std::make_unique<MyModule>("my_module", allocator,
//         std::span{kCustomModuleFunctions}).release()->interface();
//   }
template <typename State>
class NativeModule {
 public:
  NativeModule(const char* name, uint32_t version, iree_vm_instance_t* instance,
               iree_allocator_t allocator,
               iree::span<const NativeFunction<State>> dispatch_table)
      : name_(name),
        version_(version),
        instance_(instance),
        allocator_(allocator),
        dispatch_table_(dispatch_table) {
    iree_vm_instance_retain(instance);
    IREE_CHECK_OK(iree_vm_module_initialize(&interface_, this));
    interface_.destroy = NativeModule::ModuleDestroy;
    interface_.name = NativeModule::ModuleName;
    interface_.signature = NativeModule::ModuleSignature;
    // TODO(benvanik): get_module_attr
    interface_.enumerate_dependencies =
        NativeModule::ModuleEnumerateDependencies;
    interface_.lookup_function = NativeModule::ModuleLookupFunction;
    interface_.get_function = NativeModule::ModuleGetFunction;
    // TODO(benvanik): get_function_attr
    interface_.alloc_state = NativeModule::ModuleAllocState;
    interface_.free_state = NativeModule::ModuleFreeState;
    interface_.fork_state = NativeModule::ModuleForkState;
    interface_.resolve_import = NativeModule::ModuleResolveImport;
    interface_.notify = NativeModule::ModuleNotify;
    interface_.begin_call = NativeModule::ModuleBeginCall;
    // TODO(benvanik): resume_call
  }

  virtual ~NativeModule() { iree_vm_instance_release(instance_); }

  iree_vm_instance_t* instance() const { return instance_; }

  // C API module interface bound to this NativeModule instance.
  iree_vm_module_t* interface() { return &interface_; }

 protected:
  // Enumerates module dependencies by issuing |callback| for each dependency.
  virtual Status EnumerateDependencies(
      std::function<Status(const iree_vm_module_dependency_t*)> callback) {
    return OkStatus();
  }

  // Creates a new per-context module State holder.
  virtual StatusOr<std::unique_ptr<State>> CreateState(
      iree_allocator_t allocator) = 0;

  // Forks a parent state into a new per-context module State holder.
  // Anything that should be shared between the states should be retained by
  // reference.
  virtual StatusOr<std::unique_ptr<State>> ForkState(
      State* parent_state, iree_allocator_t allocator) = 0;

  // Notifies the module a signal has been raised.
  virtual Status Notify(State* state, iree_vm_signal_t signal) {
    return OkStatus();
  }

 private:
  static NativeModule* FromModulePointer(void* self) {
    return reinterpret_cast<NativeModule*>(self);
  }
  static State* FromStatePointer(void* self) {
    return reinterpret_cast<State*>(self);
  }

  static void ModuleDestroy(void* self) { delete FromModulePointer(self); }

  static iree_string_view_t ModuleName(void* self) {
    auto* module = FromModulePointer(self);
    return iree_make_cstring_view(module->name_);
  }

  static iree_vm_module_signature_t ModuleSignature(void* self) {
    auto* module = FromModulePointer(self);
    iree_vm_module_signature_t signature = {0};
    signature.version = module->version_;
    signature.attr_count = 0;
    signature.import_function_count = 0;
    signature.export_function_count = module->dispatch_table_.size();
    signature.internal_function_count = 0;
    return signature;
  }

  static iree_status_t ModuleEnumerateDependencies(
      void* self, iree_vm_module_dependency_callback_t callback,
      void* user_data) {
    auto* module = FromModulePointer(self);
    auto callback_fn =
        [callback, user_data](const iree_vm_module_dependency_t* dependency) {
          return Status(callback(user_data, dependency));
        };
    IREE_RETURN_IF_ERROR(module->EnumerateDependencies(std::move(callback_fn)));
    return OkStatus();
  }

  static iree_status_t ModuleGetFunction(
      void* self, iree_vm_function_linkage_t linkage, iree_host_size_t ordinal,
      iree_vm_function_t* out_function, iree_string_view_t* out_name,
      iree_vm_function_signature_t* out_signature) {
    if (out_function) {
      std::memset(out_function, 0, sizeof(*out_function));
    }
    if (out_name) {
      out_name->data = nullptr;
      out_name->size = 0;
    }
    if (out_signature) {
      std::memset(out_signature, 0, sizeof(*out_signature));
    }
    auto* module = FromModulePointer(self);
    if (IREE_UNLIKELY(ordinal > module->dispatch_table_.size())) {
      return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                              "function out of bounds: 0 < %" PRIhsz
                              " < %" PRIhsz,
                              ordinal, module->dispatch_table_.size());
    }
    const auto& dispatch_function = module->dispatch_table_[ordinal];
    if (out_function) {
      out_function->module = module->interface();
      out_function->linkage = IREE_VM_FUNCTION_LINKAGE_EXPORT;
      out_function->ordinal = static_cast<uint16_t>(ordinal);
    }
    if (out_name) {
      *out_name = dispatch_function.name;
    }
    if (out_signature) {
      out_signature->calling_convention = dispatch_function.cconv;
    }
    return iree_ok_status();
  }

  static iree_status_t ModuleLookupFunction(
      void* self, iree_vm_function_linkage_t linkage, iree_string_view_t name,
      const iree_vm_function_signature_t* expected_signature,
      iree_vm_function_t* out_function) {
    IREE_ASSERT_ARGUMENT(out_function);
    std::memset(out_function, 0, sizeof(*out_function));
    if (IREE_UNLIKELY(!name.data || !name.size)) {
      return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                              "function name empty");
    }

    auto* module = FromModulePointer(self);
    out_function->module = module->interface();
    out_function->linkage = IREE_VM_FUNCTION_LINKAGE_EXPORT;
    for (int i = 0; i < module->dispatch_table_.size(); ++i) {
      if (iree_string_view_equal(name, module->dispatch_table_[i].name)) {
        out_function->ordinal = i;
        return iree_ok_status();
      }
    }

    return iree_make_status(IREE_STATUS_NOT_FOUND, "function %.*s not exported",
                            (int)name.size, name.data);
  }

  static iree_status_t ModuleAllocState(
      void* self, iree_allocator_t allocator,
      iree_vm_module_state_t** out_module_state) {
    IREE_ASSERT_ARGUMENT(out_module_state);
    *out_module_state = nullptr;

    auto* module = FromModulePointer(self);
    IREE_ASSIGN_OR_RETURN(auto module_state, module->CreateState(allocator));

    *out_module_state =
        reinterpret_cast<iree_vm_module_state_t*>(module_state.release());
    return iree_ok_status();
  }

  static void ModuleFreeState(void* self,
                              iree_vm_module_state_t* module_state) {
    if (module_state) delete FromStatePointer(module_state);
  }

  static iree_status_t ModuleForkState(
      void* self, iree_vm_module_state_t* parent_state,
      iree_allocator_t allocator, iree_vm_module_state_t** out_child_state) {
    IREE_ASSERT_ARGUMENT(out_child_state);
    *out_child_state = nullptr;

    // Ignore cases where there is no state required.
    if (!parent_state) {
      return iree_ok_status();
    }

    auto* module = FromModulePointer(self);
    IREE_ASSIGN_OR_RETURN(
        auto child_state,
        module->ForkState(reinterpret_cast<State*>(parent_state), allocator));

    *out_child_state =
        reinterpret_cast<iree_vm_module_state_t*>(child_state.release());
    return iree_ok_status();
  }

  static iree_status_t ModuleResolveImport(
      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) {
    return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
                            "C++ API does not support imports");
  }

  static iree_status_t ModuleNotify(void* self,
                                    iree_vm_module_state_t* module_state,
                                    iree_vm_signal_t signal) {
    auto* module = FromModulePointer(self);
    return module->Notify(FromStatePointer(module_state), signal);
  }

  static iree_status_t ModuleBeginCall(void* self, iree_vm_stack_t* stack,
                                       iree_vm_function_call_t call) {
    auto* module = FromModulePointer(self);
    if (IREE_UNLIKELY(call.function.ordinal >=
                      module->dispatch_table_.size())) {
      return iree_make_status(
          IREE_STATUS_INVALID_ARGUMENT,
          "function ordinal out of bounds: 0 < %u < %" PRIhsz,
          call.function.ordinal, module->dispatch_table_.size());
    }
    const auto& info = module->dispatch_table_[call.function.ordinal];

    // NOTE: VM stack is currently unused. We could stash things here for the
    // debugger or use it for coroutine state.
    iree_host_size_t frame_size = 0;

    iree_vm_stack_frame_t* callee_frame = NULL;
    IREE_RETURN_IF_ERROR(iree_vm_stack_function_enter(
        stack, &call.function, IREE_VM_STACK_FRAME_NATIVE, frame_size,
        /*frame_cleanup_fn=*/nullptr, &callee_frame));

    auto* state = FromStatePointer(callee_frame->module_state);
    iree_status_t status = info.call(info.ptr, state, stack, call);
    if (IREE_UNLIKELY(!iree_status_is_ok(status))) {
      status = iree_status_annotate_f(
          status, "while invoking C++ function %s.%.*s", module->name_,
          (int)info.name.size, info.name.data);
      return status;
    }

    return iree_vm_stack_function_leave(stack);
  }

  const char* name_;
  uint32_t version_;
  iree_vm_instance_t* instance_;
  const iree_allocator_t allocator_;
  iree_vm_module_t interface_;

  const iree::span<const NativeFunction<State>> dispatch_table_;
};

}  // namespace vm
}  // namespace iree

#endif  // IREE_VM_NATIVE_MODULE_CC_H_
