crates/knut-bifrost/src/wasm.rs
//! WASM transform plugins (#6): a [`Transform`] whose logic lives in a sandboxed
//! WebAssembly module instead of compile-time Rust — **Arrow IPC in, Arrow IPC
//! out** — run by the pure-Rust [`wasmi`] interpreter (no C, no JIT). Lets a
//! transform be written in any language that compiles to wasm and ship
//! out-of-line, while still producing the same [`GraphMutation`] the fabric
//! consumes.
//!
//! ## Guest ABI
//!
//! The module must export:
//! - `memory` — linear memory.
//! - `alloc(len: i32) -> i32` — reserve `len` bytes, return the offset. The host
//! writes the input Arrow-IPC stream there.
//! - `transform(ptr: i32, len: i32) -> i64` — read the input IPC at `[ptr, ptr+len)`,
//! produce an output Arrow-IPC stream in memory, and return it **packed** as
//! `(out_ptr as i64) << 32 | (out_len as i64)`.
//!
//! Each output record batch becomes one [`VertexSet`] under the configured label
//! and id column. An identity guest (echo the input) yields the input rows as
//! vertices; a real guest filters/derives/reshapes first.
//!
//! ```no_run
//! # #[cfg(feature = "wasm")]
//! # fn demo() -> knut_bifrost::Result<()> {
//! use knut_bifrost::wasm::WasmTransform;
//! let t = WasmTransform::from_file("transform.wasm", "Person", "_gar_id")?;
//! # let _ = t; Ok(()) }
//! ```
use std::path::Path;
use arrow_array::RecordBatch;
use arrow_schema::SchemaRef;
use crate::error::{BifrostError, Result};
use crate::mutation::{GraphMutation, VertexSet};
use crate::transform::{TableContext, Transform};
/// A [`Transform`] backed by a WebAssembly module (Arrow IPC in/out).
pub struct WasmTransform {
engine: wasmi::Engine,
module: wasmi::Module,
label: String,
id_column: String,
}
impl WasmTransform {
/// Build from raw wasm bytes. The module is validated + compiled once; each
/// `apply` instantiates a fresh [`wasmi::Store`] (so calls don't share state).
pub fn from_bytes(
wasm: &[u8],
label: impl Into<String>,
id_column: impl Into<String>,
) -> Result<Self> {
let engine = wasmi::Engine::default();
let module = wasmi::Module::new(&engine, wasm)
.map_err(|e| BifrostError::Other(format!("wasm compile: {e}")))?;
Ok(Self {
engine,
module,
label: label.into(),
id_column: id_column.into(),
})
}
/// Build from a `.wasm` file on disk.
pub fn from_file(
path: impl AsRef<Path>,
label: impl Into<String>,
id_column: impl Into<String>,
) -> Result<Self> {
let wasm = std::fs::read(path.as_ref())
.map_err(|e| BifrostError::Other(format!("read wasm {:?}: {e}", path.as_ref())))?;
Self::from_bytes(&wasm, label, id_column)
}
/// Run the guest over one Arrow-IPC input stream, returning its IPC output.
fn run(&self, input: &[u8]) -> Result<Vec<u8>> {
use wasmi::{Linker, Store};
let wasm_err = |m: String| BifrostError::Other(format!("wasm: {m}"));
let mut store = Store::new(&self.engine, ());
let linker = Linker::<()>::new(&self.engine);
// `instantiate(..).and_then(|pre| pre.start(..))` is precisely what
// `instantiate_and_start` does, and the two-step form is deprecated — same
// behaviour, including running the guest's `start` function.
let instance = linker
.instantiate_and_start(&mut store, &self.module)
.map_err(|e| wasm_err(format!("instantiate: {e}")))?;
let memory = instance
.get_memory(&store, "memory")
.ok_or_else(|| wasm_err("module exports no `memory`".into()))?;
let alloc = instance
.get_typed_func::<i32, i32>(&store, "alloc")
.map_err(|e| wasm_err(format!("no `alloc(i32)->i32`: {e}")))?;
let transform = instance
.get_typed_func::<(i32, i32), i64>(&store, "transform")
.map_err(|e| wasm_err(format!("no `transform(i32,i32)->i64`: {e}")))?;
let len = i32::try_from(input.len())
.map_err(|_| wasm_err("input too large for wasm32".into()))?;
let ptr = alloc
.call(&mut store, len)
.map_err(|e| wasm_err(format!("alloc: {e}")))?;
memory
.write(&mut store, ptr as usize, input)
.map_err(|e| wasm_err(format!("write input: {e}")))?;
let packed = transform
.call(&mut store, (ptr, len))
.map_err(|e| wasm_err(format!("transform: {e}")))?;
let out_ptr = ((packed >> 32) & 0xFFFF_FFFF) as usize;
let out_len = (packed & 0xFFFF_FFFF) as usize;
let mut out = vec![0u8; out_len];
memory
.read(&store, out_ptr, &mut out)
.map_err(|e| wasm_err(format!("read output: {e}")))?;
Ok(out)
}
}
impl Transform for WasmTransform {
fn apply(&self, _ctx: &TableContext, batches: &[RecordBatch]) -> Result<GraphMutation> {
let Some(first) = batches.first() else {
return Ok(GraphMutation::default());
};
let input = write_ipc(first.schema(), batches)?;
let output = self.run(&input)?;
let out_batches = read_ipc(&output)?;
let vertices = out_batches
.into_iter()
.map(|batch| VertexSet {
label: self.label.clone(),
id_column: self.id_column.clone(),
batch,
})
.collect();
Ok(GraphMutation {
edge_deletes: Vec::new(),
vertices,
..Default::default()
})
}
}
/// Serialize record batches to an Arrow IPC **stream**.
fn write_ipc(schema: SchemaRef, batches: &[RecordBatch]) -> Result<Vec<u8>> {
let mut buf = Vec::new();
{
let mut w = arrow_ipc::writer::StreamWriter::try_new(&mut buf, &schema)
.map_err(|e| BifrostError::Other(format!("ipc writer: {e}")))?;
for b in batches {
w.write(b)
.map_err(|e| BifrostError::Other(format!("ipc write: {e}")))?;
}
w.finish()
.map_err(|e| BifrostError::Other(format!("ipc finish: {e}")))?;
}
Ok(buf)
}
/// Read an Arrow IPC **stream** back into record batches.
fn read_ipc(bytes: &[u8]) -> Result<Vec<RecordBatch>> {
let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None)
.map_err(|e| BifrostError::Other(format!("ipc reader: {e}")))?;
reader
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|e| BifrostError::Other(format!("ipc read: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use arrow_array::{Int64Array, StringArray};
use arrow_schema::{DataType, Field, Schema};
/// An identity guest in WAT: `alloc` bump-allocates from a heap cursor, and
/// `transform` returns the input buffer unchanged (echo). Exercises the full
/// host path — IPC serialize → guest memory → call → read back → VertexSet —
/// without a separate wasm build. (A real guest would parse the IPC and
/// reshape it; that's a documented `wasm32` example.)
const IDENTITY_WAT: &str = r#"
(module
(memory (export "memory") 32)
(global $heap (mut i32) (i32.const 65536))
(func (export "alloc") (param $len i32) (result i32)
(local $p i32)
(local.set $p (global.get $heap))
(global.set $heap (i32.add (global.get $heap) (local.get $len)))
(local.get $p))
(func (export "transform") (param $ptr i32) (param $len i32) (result i64)
(i64.or
(i64.shl (i64.extend_i32_u (local.get $ptr)) (i64.const 32))
(i64.extend_i32_u (local.get $len)))))
"#;
fn sample() -> RecordBatch {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("_gar_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec!["a", "b", "c"])),
],
)
.unwrap()
}
#[test]
fn wasm_identity_transform_round_trips_to_vertices() {
let wasm = wat::parse_str(IDENTITY_WAT).expect("wat compiles");
let t = WasmTransform::from_bytes(&wasm, "Person", "_gar_id").unwrap();
let batch = sample();
let mutation = t
.apply(&TableContext { table: "t".into() }, &[batch.clone()])
.unwrap();
assert_eq!(
mutation.vertices.len(),
1,
"one output batch → one vertex set"
);
let vs = &mutation.vertices[0];
assert_eq!(vs.label, "Person");
assert_eq!(vs.id_column, "_gar_id");
assert_eq!(vs.batch.num_rows(), 3);
assert_eq!(
vs.batch, batch,
"identity guest echoes the rows through the IPC bridge"
);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"wasm::wasm_identity_transform_round_trips_to_vertices",
vs.label == "Person" && vs.batch.num_rows() == 3 && vs.batch == batch,
&format!("identity guest echoed {} rows via IPC", vs.batch.num_rows()),
);
}
#[test]
fn missing_export_is_a_clean_error() {
// A module with no `transform` export → a descriptive error, not a panic.
let wasm = wat::parse_str(r#"(module (memory (export "memory") 1))"#).unwrap();
let t = WasmTransform::from_bytes(&wasm, "X", "_gar_id").unwrap();
let err = t
.apply(&TableContext { table: "t".into() }, &[sample()])
.unwrap_err();
let msg = format!("{err}");
assert!(msg.contains("alloc"), "got: {err}");
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"wasm::missing_export_is_a_clean_error",
msg.contains("alloc"),
"missing export → descriptive error, not panic",
);
}
}