reifydb-extension 0.4.9

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2025 ReifyDB

//! WASM procedure implementation that executes WebAssembly modules as stored procedures

use postcard::to_stdvec;
use reifydb_core::value::column::columns::Columns;
use reifydb_routine::procedure::{Procedure, context::ProcedureContext, error::ProcedureError};
use reifydb_sdk::{error::FFIError, marshal::wasm::unmarshal_columns_from_bytes};
use reifydb_transaction::transaction::Transaction;
use reifydb_type::error::Error;

use crate::{error::ExtensionError, loader::wasm::invoke_wasm_module};

fn ext_err(err: ExtensionError) -> ProcedureError {
	ProcedureError::Wrapped(Box::new(Error::from(FFIError::Other(err.to_string()))))
}

/// WASM procedure that loads and executes a `.wasm` module.
///
/// Each WASM module must export:
/// - `alloc(size: i32) -> i32` — allocate `size` bytes, return pointer
/// - `dealloc(ptr: i32, size: i32)` — free memory
/// - `procedure(params_ptr: i32, params_len: i32) -> i32` — pointer to output (first 4 bytes at output pointer = output
///   length as LE u32)
pub struct WasmProcedure {
	name: String,
	wasm_bytes: Vec<u8>,
}

impl WasmProcedure {
	pub fn new(name: impl Into<String>, wasm_bytes: Vec<u8>) -> Self {
		Self {
			name: name.into(),
			wasm_bytes,
		}
	}

	pub fn name(&self) -> &str {
		&self.name
	}
}

// SAFETY: WasmProcedure only holds inert data (name + bytes).
// A fresh Engine is created per invocation, so no shared mutable state.
unsafe impl Send for WasmProcedure {}
unsafe impl Sync for WasmProcedure {}

impl Procedure for WasmProcedure {
	fn call(&self, ctx: &ProcedureContext, _tx: &mut Transaction<'_>) -> Result<Columns, ProcedureError> {
		let params_bytes = to_stdvec(ctx.params).map_err(|e| {
			ext_err(ExtensionError::Invocation(format!(
				"WASM procedure '{}' failed to serialize params: {}",
				self.name, e
			)))
		})?;

		let label = format!("WASM procedure '{}'", self.name);
		let output_bytes =
			invoke_wasm_module(&self.wasm_bytes, "procedure", &params_bytes, &label).map_err(ext_err)?;

		Ok(unmarshal_columns_from_bytes(&output_bytes))
	}
}