nordisk/knut

crates/knut-bifrost/wasm-guest/src/lib.rs

//! A **real** `wasm32-unknown-unknown` guest for the `WasmTransform` host in
//! `crates/knut-bifrost/src/wasm.rs`. Arrow IPC **stream** in, Arrow IPC stream
//! out — proving arrow-on-wasm works, and a template for users' own guests.
//!
//! ## Transform logic
//!
//! Read the input record batches, **keep only the rows where the `_gar_id`
//! column is even**, and write the surviving rows back as an Arrow IPC stream.
//! `_gar_id` is read as `Int64` (the column the host turns into vertex ids).
//!
//! ## Guest ABI (must match the host)
//!
//! - `memory` — linear memory (exported automatically for `cdylib`).
//! - `alloc(len: i32) -> i32` — reserve `len` bytes, return the offset; the host
//!   writes the input IPC stream there. Implemented as a `Vec`-leak.
//! - `transform(ptr: i32, len: i32) -> i64` — read the IPC at `[ptr, ptr+len)`,
//!   filter, write an IPC stream into a leaked buffer, and return it packed as
//!   `(out_ptr as i64) << 32 | (out_len as i64)`.
//!
//! Build: `cargo build --target wasm32-unknown-unknown --release`.

#![no_std]

extern crate alloc;

use alloc::vec::Vec;
use core::mem;

use arrow_array::cast::AsArray;
use arrow_array::types::Int64Type;
use arrow_array::{BooleanArray, RecordBatch};
use arrow_ipc::reader::StreamReader;
use arrow_ipc::writer::StreamWriter;
use arrow_select::filter::filter_record_batch;

/// Reserve `len` bytes in guest memory and return the offset. The host writes
/// the input Arrow-IPC stream at the returned pointer.
///
/// Implemented as a `Vec`-leak: allocate a `Vec<u8>` of the requested length,
/// hand out its pointer, and forget it so the allocation outlives the call.
/// `transform` reconstructs nothing from it — it only reads through the raw
/// pointer the host passes back — so the bytes simply stay live for the call's
/// duration (a fresh `wasmi::Store` per `apply` means leaks don't accumulate
/// across host invocations).
#[no_mangle]
pub extern "C" fn alloc(len: i32) -> i32 {
    let len = len.max(0) as usize;
    let mut buf: Vec<u8> = Vec::with_capacity(len);
    // SAFETY: capacity == len; the bytes are about to be overwritten by the host.
    unsafe { buf.set_len(len) };
    let ptr = buf.as_mut_ptr();
    mem::forget(buf);
    ptr as i32
}

/// Read the input IPC stream, keep only even-`_gar_id` rows, write the result as
/// an IPC stream into a leaked buffer, and return `(out_ptr << 32) | out_len`.
///
/// On any error (bad IPC, missing/wrong-typed `_gar_id`, write failure) the
/// guest writes an **empty but valid** IPC stream (schema only, no batches) when
/// it can, or returns `(0, 0)` if even that fails — the host reads an empty
/// stream as zero output batches rather than trapping.
#[no_mangle]
pub extern "C" fn transform(ptr: i32, len: i32) -> i64 {
    let input: &[u8] = unsafe { core::slice::from_raw_parts(ptr as *const u8, len.max(0) as usize) };
    let out = run(input).unwrap_or_default();
    let out_len = out.len() as i64;
    let out_ptr = out.as_ptr() as i64;
    mem::forget(out); // hand ownership to the host; it reads then drops the Store.
    (out_ptr << 32) | (out_len & 0xFFFF_FFFF)
}

/// The actual transform: IPC → filter even `_gar_id` → IPC. Returns `None` on
/// any failure so the caller can fall back to `(0, 0)`.
fn run(input: &[u8]) -> Option<Vec<u8>> {
    let reader = StreamReader::try_new(input, None).ok()?;
    let schema = reader.schema();

    let mut out: Vec<u8> = Vec::new();
    {
        let mut writer = StreamWriter::try_new(&mut out, &schema).ok()?;
        for batch in reader {
            let batch = batch.ok()?;
            let filtered = filter_even_gar_id(&batch)?;
            writer.write(&filtered).ok()?;
        }
        writer.finish().ok()?;
    }
    Some(out)
}

/// Build a boolean mask `_gar_id % 2 == 0` and filter the batch with it. Nulls
/// in `_gar_id` are treated as non-matching (dropped). Returns `None` if there
/// is no `Int64` `_gar_id` column.
fn filter_even_gar_id(batch: &RecordBatch) -> Option<RecordBatch> {
    let col = batch.column_by_name("_gar_id")?;
    let ids = col.as_primitive_opt::<Int64Type>()?;

    let mask: BooleanArray = ids
        .iter()
        .map(|v| Some(v.map(|x| x % 2 == 0).unwrap_or(false)))
        .collect();

    filter_record_batch(batch, &mask).ok()
}