reifydb_extension/function/
wasm.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4//! WASM scalar function implementation that executes WebAssembly modules as scalar functions
5
6use reifydb_core::value::column::{Column, columns::Columns, data::ColumnData};
7use reifydb_routine::function::{Function, FunctionCapability, FunctionContext, FunctionInfo, error::FunctionError};
8use reifydb_sdk::marshal::wasm::{marshal_columns_to_bytes, unmarshal_columns_from_bytes};
9use reifydb_type::{fragment::Fragment, value::r#type::Type};
10
11use crate::loader::wasm::invoke_wasm_module;
12
13/// WASM scalar function that loads and executes a `.wasm` module.
14///
15/// Each WASM module must export:
16/// - `alloc(size: i32) -> i32` — allocate `size` bytes, return pointer
17/// - `dealloc(ptr: i32, size: i32)` — free memory
18/// - `scalar(input_ptr: i32, input_len: i32) -> i32` — pointer to output (first 4 bytes at output pointer = output
19///   length as LE u32)
20///
21/// Input: the context's `columns` marshalled as flat binary.
22/// Output: flat binary representing a single-column `Columns`, from which
23///   the first column's `ColumnData` is extracted.
24pub struct WasmScalarFunction {
25	info: FunctionInfo,
26	wasm_bytes: Vec<u8>,
27}
28
29impl WasmScalarFunction {
30	pub fn new(name: impl Into<String>, wasm_bytes: Vec<u8>) -> Self {
31		let name = name.into();
32		Self {
33			info: FunctionInfo::new(&name),
34			wasm_bytes,
35		}
36	}
37
38	pub fn name(&self) -> &str {
39		&self.info.name
40	}
41
42	fn err(&self, reason: impl Into<String>) -> FunctionError {
43		FunctionError::ExecutionFailed {
44			function: Fragment::internal(&self.info.name),
45			reason: reason.into(),
46		}
47	}
48}
49
50// SAFETY: WasmScalarFunction only holds inert data (name + bytes).
51// A fresh Engine is created per invocation, so no shared mutable state.
52unsafe impl Send for WasmScalarFunction {}
53unsafe impl Sync for WasmScalarFunction {}
54
55impl Function for WasmScalarFunction {
56	fn info(&self) -> &FunctionInfo {
57		&self.info
58	}
59
60	fn capabilities(&self) -> &[FunctionCapability] {
61		&[FunctionCapability::Scalar]
62	}
63
64	fn return_type(&self, _input_types: &[Type]) -> Type {
65		Type::Any
66	}
67
68	fn execute(&self, ctx: &FunctionContext, args: &Columns) -> Result<Columns, FunctionError> {
69		let input_bytes = marshal_columns_to_bytes(args);
70		let label = format!("WASM scalar function '{}'", self.info.name);
71
72		let output_bytes = invoke_wasm_module(&self.wasm_bytes, "scalar", &input_bytes, &label)
73			.map_err(|e| self.err(e.to_string()))?;
74
75		// Unmarshal as Columns and extract the first column's data
76		let output_columns = unmarshal_columns_from_bytes(&output_bytes);
77
78		match output_columns.first() {
79			Some(col) => {
80				let data = col.data().clone();
81				Ok(Columns::new(vec![Column::new(ctx.fragment.clone(), data)]))
82			}
83			None => {
84				let data = ColumnData::none_typed(Type::Any, args.row_count());
85				Ok(Columns::new(vec![Column::new(ctx.fragment.clone(), data)]))
86			}
87		}
88	}
89}
reifydb_extension/function/wasm.rs

reifydb_extension/function/
wasm.rs
