ffi/src/lib.rs

#![allow(clippy::missing_safety_doc)]

use paimon_ftindex_core::io::{ReadRequest, SeekRead, SeekWrite};
use paimon_ftindex_core::{
    FullTextIndexConfig, FullTextIndexReader, FullTextIndexWriter, FullTextQuery, TokenizerConfig,
};
use std::cell::RefCell;
use std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::io;
use std::os::raw::{c_char, c_int, c_void};
use std::panic::{self, AssertUnwindSafe};
use std::{ptr, slice};

thread_local! {
    static LAST_ERROR: RefCell<Option<CString>> = const { RefCell::new(None) };
}

fn set_error(msg: impl Into<String>) {
    let msg = msg.into().replace('\0', "\\0");
    LAST_ERROR.with(|e| {
        *e.borrow_mut() = CString::new(msg).ok();
    });
}

fn panic_message(e: &Box<dyn std::any::Any + Send>) -> String {
    if let Some(s) = e.downcast_ref::<String>() {
        format!("native panic: {s}")
    } else if let Some(s) = e.downcast_ref::<&str>() {
        format!("native panic: {s}")
    } else {
        "native panic: unknown".to_string()
    }
}

fn ffi_status<F>(f: F) -> c_int
where
    F: FnOnce() -> Result<(), String>,
{
    match panic::catch_unwind(AssertUnwindSafe(f)) {
        Ok(Ok(())) => 0,
        Ok(Err(e)) => {
            set_error(e);
            -1
        }
        Err(e) => {
            set_error(panic_message(&e));
            -1
        }
    }
}

fn ffi_ptr<T, F>(f: F) -> *mut T
where
    F: FnOnce() -> Result<*mut T, String>,
{
    match panic::catch_unwind(AssertUnwindSafe(f)) {
        Ok(Ok(value)) => value,
        Ok(Err(e)) => {
            set_error(e);
            ptr::null_mut()
        }
        Err(e) => {
            set_error(panic_message(&e));
            ptr::null_mut()
        }
    }
}

#[no_mangle]
pub extern "C" fn paimon_ftindex_last_error() -> *const c_char {
    LAST_ERROR.with(|e| match &*e.borrow() {
        Some(msg) => msg.as_ptr(),
        None => ptr::null(),
    })
}

#[repr(C)]
pub struct PaimonFtindexOutputFile {
    pub ctx: *mut c_void,
    pub write_fn: Option<unsafe extern "C" fn(*mut c_void, *const u8, usize) -> c_int>,
    pub flush_fn: Option<unsafe extern "C" fn(*mut c_void) -> c_int>,
}

struct FfiOutputFile {
    raw: PaimonFtindexOutputFile,
}

unsafe impl Send for FfiOutputFile {}

impl SeekWrite for FfiOutputFile {
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        let write_fn = self
            .raw
            .write_fn
            .ok_or_else(|| io::Error::other("write_fn is null"))?;
        let status = unsafe { write_fn(self.raw.ctx, buf.as_ptr(), buf.len()) };
        if status == 0 {
            Ok(())
        } else {
            Err(io::Error::other("write callback failed"))
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        if let Some(flush_fn) = self.raw.flush_fn {
            let status = unsafe { flush_fn(self.raw.ctx) };
            if status != 0 {
                return Err(io::Error::other("flush callback failed"));
            }
        }
        Ok(())
    }
}

#[repr(C)]
pub struct PaimonFtindexInputFile {
    pub ctx: *mut c_void,
    pub read_at_fn: Option<unsafe extern "C" fn(*mut c_void, u64, *mut u8, usize) -> c_int>,
}

struct FfiInputFile {
    raw: PaimonFtindexInputFile,
}

unsafe impl Send for FfiInputFile {}

impl SeekRead for FfiInputFile {
    fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()> {
        let read_at_fn = self
            .raw
            .read_at_fn
            .ok_or_else(|| io::Error::other("read_at_fn is null"))?;
        for range in ranges {
            let status = unsafe {
                read_at_fn(
                    self.raw.ctx,
                    range.pos,
                    range.buf.as_mut_ptr(),
                    range.buf.len(),
                )
            };
            if status != 0 {
                return Err(io::Error::other(format!(
                    "read_at callback failed at offset {} length {}",
                    range.pos,
                    range.buf.len()
                )));
            }
        }
        Ok(())
    }
}

pub struct PaimonFtindexWriterHandle {
    inner: FullTextIndexWriter,
}

pub struct PaimonFtindexReaderHandle {
    inner: FullTextIndexReader<FfiInputFile>,
}

#[no_mangle]
pub unsafe extern "C" fn paimon_ftindex_writer_open(
    keys: *const *const c_char,
    values: *const *const c_char,
    len: usize,
) -> *mut PaimonFtindexWriterHandle {
    ffi_ptr(|| {
        let options = options_from_raw(keys, values, len)?;
        let tokenizer = TokenizerConfig::from_options(&options).map_err(|e| e.to_string())?;
        let config = FullTextIndexConfig::new().tokenizer(tokenizer);
        let writer = FullTextIndexWriter::new(config).map_err(|e| e.to_string())?;
        Ok(Box::into_raw(Box::new(PaimonFtindexWriterHandle {
            inner: writer,
        })))
    })
}

#[no_mangle]
pub unsafe extern "C" fn paimon_ftindex_writer_add_document(
    writer: *mut PaimonFtindexWriterHandle,
    row_id: i64,
    text: *const c_char,
) -> c_int {
    ffi_status(|| {
        let writer = require_mut(writer, "writer")?;
        let text = cstr_to_string(text, "text")?;
        writer
            .inner
            .add_document(row_id, text)
            .map_err(|e| e.to_string())
    })
}

#[no_mangle]
pub unsafe extern "C" fn paimon_ftindex_writer_write_index(
    writer: *mut PaimonFtindexWriterHandle,
    output: PaimonFtindexOutputFile,
) -> c_int {
    ffi_status(|| {
        let writer = require_mut(writer, "writer")?;
        let mut output = FfiOutputFile { raw: output };
        writer.inner.write(&mut output).map_err(|e| e.to_string())
    })
}

#[no_mangle]
pub unsafe extern "C" fn paimon_ftindex_writer_free(writer: *mut PaimonFtindexWriterHandle) {
    if !writer.is_null() {
        drop(Box::from_raw(writer));
    }
}

#[no_mangle]
pub unsafe extern "C" fn paimon_ftindex_reader_open(
    input: PaimonFtindexInputFile,
) -> *mut PaimonFtindexReaderHandle {
    ffi_ptr(|| {
        let input = FfiInputFile { raw: input };
        let reader = FullTextIndexReader::open(input).map_err(|e| e.to_string())?;
        Ok(Box::into_raw(Box::new(PaimonFtindexReaderHandle {
            inner: reader,
        })))
    })
}

#[no_mangle]
pub unsafe extern "C" fn paimon_ftindex_reader_search_json(
    reader: *mut PaimonFtindexReaderHandle,
    query_json: *const c_char,
    limit: usize,
    row_ids: *mut i64,
    scores: *mut f32,
    capacity: usize,
    result_len: *mut usize,
) -> c_int {
    ffi_status(|| {
        let reader = require_mut(reader, "reader")?;
        if result_len.is_null() {
            return Err("result_len is null".to_string());
        }
        let query_json = cstr_to_string(query_json, "query_json")?;
        let query = FullTextQuery::from_json(&query_json).map_err(|e| e.to_string())?;
        let result = reader
            .inner
            .search(query, limit)
            .map_err(|e| e.to_string())?;
        unsafe {
            *result_len = result.row_ids.len();
        }
        if result.row_ids.len() > capacity {
            return Err(format!(
                "result capacity {} is smaller than result length {}",
                capacity,
                result.row_ids.len()
            ));
        }
        if !result.row_ids.is_empty() {
            if row_ids.is_null() {
                return Err("row_ids is null".to_string());
            }
            if scores.is_null() {
                return Err("scores is null".to_string());
            }
            unsafe {
                ptr::copy_nonoverlapping(result.row_ids.as_ptr(), row_ids, result.row_ids.len());
                ptr::copy_nonoverlapping(result.scores.as_ptr(), scores, result.scores.len());
            }
        }
        Ok(())
    })
}

#[no_mangle]
pub unsafe extern "C" fn paimon_ftindex_reader_free(reader: *mut PaimonFtindexReaderHandle) {
    if !reader.is_null() {
        drop(Box::from_raw(reader));
    }
}

unsafe fn options_from_raw(
    keys: *const *const c_char,
    values: *const *const c_char,
    len: usize,
) -> Result<HashMap<String, String>, String> {
    if len == 0 {
        return Ok(HashMap::new());
    }
    if keys.is_null() {
        return Err("keys is null".to_string());
    }
    if values.is_null() {
        return Err("values is null".to_string());
    }
    let keys = slice::from_raw_parts(keys, len);
    let values = slice::from_raw_parts(values, len);
    let mut options = HashMap::with_capacity(len);
    for i in 0..len {
        let key = cstr_to_string(keys[i], "option key")?;
        let value = cstr_to_string(values[i], "option value")?;
        options.insert(key, value);
    }
    Ok(options)
}

unsafe fn require_mut<'a, T>(ptr: *mut T, name: &str) -> Result<&'a mut T, String> {
    ptr.as_mut().ok_or_else(|| format!("{name} is null"))
}

unsafe fn cstr_to_string(ptr: *const c_char, name: &str) -> Result<String, String> {
    if ptr.is_null() {
        return Err(format!("{name} is null"));
    }
    CStr::from_ptr(ptr)
        .to_str()
        .map(|s| s.to_string())
        .map_err(|e| format!("{name} is not valid UTF-8: {e}"))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::ffi::CString;

    unsafe extern "C" fn write_vec(ctx: *mut c_void, buf: *const u8, len: usize) -> c_int {
        let out = &mut *(ctx as *mut Vec<u8>);
        out.extend_from_slice(slice::from_raw_parts(buf, len));
        0
    }

    unsafe extern "C" fn read_vec(ctx: *mut c_void, pos: u64, buf: *mut u8, len: usize) -> c_int {
        let input = &*(ctx as *const Vec<u8>);
        let start = pos as usize;
        let end = start + len;
        if end > input.len() {
            return -1;
        }
        ptr::copy_nonoverlapping(input[start..end].as_ptr(), buf, len);
        0
    }

    #[test]
    fn ffi_round_trip_search_json() {
        unsafe {
            let writer = paimon_ftindex_writer_open(ptr::null(), ptr::null(), 0);
            assert!(!writer.is_null());

            let text = CString::new("Apache Paimon full text").unwrap();
            assert_eq!(
                paimon_ftindex_writer_add_document(writer, 7, text.as_ptr()),
                0
            );

            let mut bytes = Vec::new();
            let output = PaimonFtindexOutputFile {
                ctx: &mut bytes as *mut Vec<u8> as *mut c_void,
                write_fn: Some(write_vec),
                flush_fn: None,
            };
            assert_eq!(paimon_ftindex_writer_write_index(writer, output), 0);
            paimon_ftindex_writer_free(writer);

            let input = PaimonFtindexInputFile {
                ctx: &bytes as *const Vec<u8> as *mut c_void,
                read_at_fn: Some(read_vec),
            };
            let reader = paimon_ftindex_reader_open(input);
            assert!(!reader.is_null());

            let query = CString::new(
                paimon_ftindex_core::FullTextQuery::match_query("paimon", "text")
                    .to_json()
                    .unwrap(),
            )
            .unwrap();
            let mut row_ids = [0i64; 4];
            let mut scores = [0f32; 4];
            let mut result_len = 0usize;
            assert_eq!(
                paimon_ftindex_reader_search_json(
                    reader,
                    query.as_ptr(),
                    4,
                    row_ids.as_mut_ptr(),
                    scores.as_mut_ptr(),
                    row_ids.len(),
                    &mut result_len,
                ),
                0
            );
            assert_eq!(result_len, 1);
            assert_eq!(row_ids[0], 7);
            assert!(scores[0] > 0.0);
            paimon_ftindex_reader_free(reader);
        }
    }
}