knut-bifrost 0.3.0

//! 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",
        );
    }
}