Nova64/core/src/wasm_runner.c

#include "wasm_runner.h"
#include "FreeRTOS.h"
#include "nova64_api.h"
#include "nova64_fs.h"
#include "nova64_internal.h"
#include "task.h"
#include "wasm_export.h"
#include <stdlib.h>

#define MAX_WASM_TASKS 4

typedef enum { SLOT_FREE = 0, SLOT_RUNNING, SLOT_STOPPING } slot_state_t;

typedef struct {
  int id;
  slot_state_t state;
  TaskHandle_t freertos_handle;
  wasm_module_inst_t module_inst;
  wasm_exec_env_t exec_env;
} wasm_task_slot_t;

static wasm_task_slot_t g_wasm_slots[MAX_WASM_TASKS];

char error_buf[128];
uint32_t stack_size = 8092, heap_size = 8092;

bool wasm_runner_initialize(void) {

  for (int i = 0; i < MAX_WASM_TASKS; i++) {
    g_wasm_slots[i].id = i;
    g_wasm_slots[i].state = SLOT_FREE;
    g_wasm_slots[i].freertos_handle = NULL;
    g_wasm_slots[i].module_inst = NULL;
  }

  wasm_runtime_init();
  nova64_register_imports();
  nova64_log("WASM runtime initialized successfully.\n");

  return true;
}

bool wasm_runner_destroy(void) {
  wasm_runtime_destroy();
  return true;
}

uint32_t wasm_runner_get_active_slots() {
  uint32_t result = 0;
  for (int i = 0; i < MAX_WASM_TASKS; i++) {
    if (g_wasm_slots[i].state != SLOT_FREE) {
      result++;
    }
  }
  return result;
}

bool wasm_runner_start_program(
    const char *module_path,
    const wasm_runner_framebuffer_config_t *framebuffer_config) {

  nova64_log("Running WASM module: %s\n", module_path);

  wasm_task_slot_t *slot = NULL;
  for (int i = 0; i < MAX_WASM_TASKS; i++) {
    if (g_wasm_slots[i].state == SLOT_FREE) {
      slot = &g_wasm_slots[i];
      break;
    }
  }

  if (!slot) {
    nova64_log("Maximum number of WASM tasks (%d) reached!\n", MAX_WASM_TASKS);
    return false;
  }

  if (!wasm_runtime_init_thread_env()) {
    nova64_log("Failed to initialize WASM runtime thread environment.\n");
    return false;
  }

  int32_t file_size = nova64_get_file_size(module_path);
  if (file_size <= 0) {
    nova64_log("Failed to determine WASM module size.\n");
    return false;
  }

  unsigned char *buffer = malloc((size_t)file_size);
  if (!buffer) {
    nova64_log("Failed to allocate WASM module buffer.\n");
    return false;
  }

  if (!nova64_read_file_to_buffer(module_path, buffer, (size_t)file_size)) {
    nova64_log("Failed to read WASM module into buffer.\n");
    free(buffer);
    return false;
  }

  wasm_module_t module = wasm_runtime_load(buffer, (uint32_t)file_size,
                                           error_buf, sizeof(error_buf));

  slot->module_inst = wasm_runtime_instantiate(module, stack_size, heap_size,
                                               error_buf, sizeof(error_buf));

  wasm_function_inst_t start_func =
      wasm_runtime_lookup_function(slot->module_inst, "_start");

  if (start_func == NULL) {
    nova64_log("Failed to find '_start' function in WASM module: %s\n",
               wasm_runtime_get_exception(slot->module_inst));
    free(buffer);
    return false;
  }

  slot->exec_env = wasm_runtime_create_exec_env(slot->module_inst, stack_size);

  nova64_log("WASM module execution environment created successfully.\n");

  uint32_t argv[2] = {5, 3}; // Example arguments for the "add" function

  slot->state = SLOT_RUNNING;

  /* Record the FreeRTOS task handle for this slot so the task can be
   * referenced or force-deleted later if necessary. */
  slot->freertos_handle = xTaskGetCurrentTaskHandle();

  if (wasm_runtime_call_wasm(slot->exec_env, start_func, 0, NULL)) {
    nova64_log("WASM module executed successfully.\n");
  } else {
    nova64_log("Failed to execute WASM module: %s\n",
               wasm_runtime_get_exception(slot->module_inst));
  }

  slot->state = SLOT_STOPPING;

  wasm_runtime_destroy_exec_env(slot->exec_env);
  wasm_runtime_deinstantiate(slot->module_inst);
  wasm_runtime_unload(module);
  free(buffer);
  wasm_runtime_destroy_thread_env();

  slot->module_inst = NULL;
  slot->exec_env = NULL;
  slot->state = SLOT_FREE;
  nova64_log("WASM module execution completed and resources cleaned up.\n");
  nova64_log("WASM task slot %d is now free. Deleting task %p...\n", slot->id,
             (void *)xTaskGetCurrentTaskHandle());
  slot->freertos_handle = NULL;
  vTaskDelete(NULL);

  return true;
}

void wasm_runner_terminate(void) {
  nova64_log("Stopping all WASM tasks...\n");

  // 1. Send stop signal to all running WASM instances
  for (int i = 0; i < MAX_WASM_TASKS; i++) {
    if (g_wasm_slots[i].state == SLOT_RUNNING) {
      g_wasm_slots[i].state = SLOT_STOPPING;
      if (g_wasm_slots[i].module_inst) {
        // Thread-safe: Only sets an abort flag internally in WAMR
        wasm_runtime_terminate(g_wasm_slots[i].module_inst);
      }
    }
  }

  // 2. Wait with HOST-Sleep (NATIVE C / C# thread-safe!)
  int timeout_ms = 2000;
  while (timeout_ms > 0) {
    bool any_running = false;
    for (int i = 0; i < MAX_WASM_TASKS; i++) {
      if (g_wasm_slots[i].state != SLOT_FREE) {
        any_running = true;
        break;
      }
    }

    if (!any_running) {
      break; // All WASM tasks have terminated and cleaned up themselves!
    }

    // IMPORTANT: Use Host-Sleep, NEVER use vTaskDelay on the C# thread!
    nova64_host_sleep_ms(10);
    timeout_ms -= 10;
  }

  // 3. If after 2s a task is still hanging: ONLY clean up flags,
  //    NEVER call vTaskDelete from outside!
  for (int i = 0; i < MAX_WASM_TASKS; i++) {
    if (g_wasm_slots[i].state != SLOT_FREE) {
      nova64_log("Warning: Task slot %d failed to exit gracefully.\n", i);
      g_wasm_slots[i].state = SLOT_FREE;
      g_wasm_slots[i].freertos_handle = NULL;
    }
  }

  nova64_log("All WASM tasks successfully stopped.\n");
}