luby/crates/luaffi/src/string.rs

use crate::{
    __internal::{disp, display, export},
    FromFfi, IntoFfi,
};
use bstr::{BStr, BString};
use luaffi_impl::{cdef, metatype};
use std::{fmt::Display, mem::ManuallyDrop, ptr, slice};

pub(crate) const IS_UTF8_FN: &str = "luaffi_is_utf8";
pub(crate) const DROP_BUFFER_FN: &str = "luaffi_drop_buffer";

#[unsafe(export_name = "luaffi_is_utf8")]
unsafe extern "C" fn __is_utf8(ptr: *const u8, len: usize) -> bool {
    debug_assert!(!ptr.is_null());
    simdutf8::basic::from_utf8(unsafe { slice::from_raw_parts(ptr, len) }).is_ok()
}

#[unsafe(export_name = "luaffi_drop_buffer")]
unsafe extern "C" fn __drop_buffer(buf: *mut lua_buffer) {
    debug_assert!(!buf.is_null());
    debug_assert!(!unsafe { (*buf).__ptr.is_null() });
    drop(unsafe { Vec::from_raw_parts((*buf).__ptr, (*buf).__len, (*buf).__cap) })
}

export![__is_utf8, __drop_buffer];

#[derive(Debug, Clone, Copy)]
#[cdef]
pub struct lua_buf {
    __ptr: *const u8,
    __len: usize,
}

#[metatype]
impl lua_buf {
    // this takes a slice and decomposes it into its raw parts. caller should ensure the result is
    // used only as long as the original buffer is still alive.
    pub(crate) fn new(s: &[u8]) -> Self {
        Self {
            __ptr: s.as_ptr(),
            __len: s.len(),
        }
    }

    pub(crate) fn null() -> Self {
        Self {
            __ptr: ptr::null(),
            __len: 0,
        }
    }
}

#[derive(Debug, Clone, Copy)]
#[cdef]
pub struct lua_buffer {
    __ptr: *mut u8,
    __len: usize,
    __cap: usize,
}

#[metatype]
impl lua_buffer {
    // this takes ownership of the Vec and decomposes it into its raw parts. the result must be
    // dropped by `__drop_buffer` (see [`DROP_BUFFER_FN`]).
    pub(crate) fn new(s: impl Into<Vec<u8>>) -> Self {
        let s = s.into();
        Self {
            __cap: s.capacity(),
            __len: s.len(),
            __ptr: ManuallyDrop::new(s).as_mut_ptr(),
        }
    }

    pub(crate) fn null() -> Self {
        Self {
            __ptr: ptr::null_mut(),
            __len: 0,
            __cap: 0,
        }
    }
}

unsafe impl<'s> FromFfi for &'s [u8] {
    type From = Option<&'s lua_buf>;

    fn require_keepalive() -> bool {
        true
    }

    fn prelude(arg: &str) -> impl Display {
        // this converts string arguments to a `lua_buf` with a pointer to the string and its length
        disp(move |f| {
            write!(
                f,
                r#"assert(type({arg}) == "string", "string expected in argument '{arg}', got " .. type({arg})); "#
            )?;
            // SAFETY: the lua_buf is only valid for as long as the string is alive. we've ensured
            // that it is alive for at least the duration of the ffi call via `require_keepalive()`.
            write!(f, "{arg} = __new(__ct.lua_buf, {arg}, #{arg}); end; ")
        })
    }

    fn convert(from: Self::From) -> Self {
        // SAFETY: we already checked that the string is nonnull from the lua side
        debug_assert!(from.is_some());
        let from = unsafe { from.unwrap_unchecked() };
        debug_assert!(!from.__ptr.is_null());
        unsafe { slice::from_raw_parts(from.__ptr, from.__len) }
    }
}

unsafe impl<'s> FromFfi for &'s str {
    type From = Option<&'s lua_buf>;

    fn require_keepalive() -> bool {
        true
    }

    fn prelude(arg: &str) -> impl Display {
        // this converts string arguments to a `lua_buf` with a pointer to the string and its length
        // and ensures that the string is valid utf8
        disp(move |f| {
            write!(
                f,
                r#"assert(type({arg}) == "string", "string expected in argument '{arg}', got " .. type({arg})); "#
            )?;
            write!(
                f,
                r#"assert(__C.{IS_UTF8_FN}({arg}, #{arg}), "argument '{arg}' must be a valid utf-8 string"); "#
            )?;
            write!(f, "{arg} = __new(__ct.lua_buf, {arg}, #{arg}); ")
        })
    }

    fn convert(from: Self::From) -> Self {
        // SAFETY: we already checked that the string is nonnull and valid utf8 from the lua side
        debug_assert!(from.is_some());
        let from = unsafe { from.unwrap_unchecked() };
        debug_assert!(!from.__ptr.is_null());
        let from = unsafe { slice::from_raw_parts(from.__ptr, from.__len) };
        debug_assert!(
            std::str::from_utf8(from).is_ok(),
            "<&str>::convert() called on an invalid utf8 string when it was checked to be valid"
        );
        unsafe { std::str::from_utf8_unchecked(from) }
    }
}

unsafe impl IntoFfi for &'static [u8] {
    type Into = lua_buf;

    fn convert(self) -> Self::Into {
        // SAFETY: the slice is 'static so the resulting lua_buf is always valid
        lua_buf::new(self)
    }

    fn postlude(ret: &str) -> impl Display {
        display!("{ret} = __intern({ret}.__ptr, {ret}.__len)")
    }
}

unsafe impl IntoFfi for Vec<u8> {
    type Into = lua_buffer;

    fn convert(self) -> Self::Into {
        lua_buffer::new(self)
    }

    fn postlude(ret: &str) -> impl Display {
        display!(
            "do local __{ret} = {ret}; {ret} = __intern({ret}.__ptr, {ret}.__len); __C.{DROP_BUFFER_FN}(__{ret}); end; "
        )
    }
}

macro_rules! impl_from_via {
    ($ty:ty, $via:ty) => {
        unsafe impl<'s> FromFfi for $ty {
            type From = <$via as FromFfi>::From;

            fn require_keepalive() -> bool {
                <$via as FromFfi>::require_keepalive()
            }

            fn prelude(arg: &str) -> impl Display {
                <$via as FromFfi>::prelude(arg)
            }

            fn convert(from: Self::From) -> Self {
                <$via as FromFfi>::convert(from).into()
            }
        }
    };
}

impl_from_via!(&'s BStr, &'s [u8]);

macro_rules! impl_into_via {
    ($ty:ty, $via:ty) => {
        unsafe impl IntoFfi for $ty {
            type Into = <$via as IntoFfi>::Into;

            fn convert(self) -> Self::Into {
                <$via as IntoFfi>::convert(self.into())
            }

            fn postlude(ret: &str) -> impl Display {
                <$via as IntoFfi>::postlude(ret)
            }
        }
    };
}

impl_into_via!(&'static BStr, &'static [u8]);
impl_into_via!(&'static str, &'static BStr);
impl_into_via!(BString, Vec<u8>);
impl_into_via!(String, BString);

macro_rules! impl_optional_from {
    ($ty:ty) => {
        unsafe impl<'s> FromFfi for Option<$ty> {
            type From = <$ty as FromFfi>::From;

            fn require_keepalive() -> bool {
                <$ty as FromFfi>::require_keepalive()
            }

            fn prelude(arg: &str) -> impl Display {
                // just pass a null pointer if argument is nil
                display!(
                    "if {arg} ~= nil then {}end; ",
                    <$ty as FromFfi>::prelude(arg)
                )
            }

            fn convert(from: Self::From) -> Self {
                from.map(|s| <$ty as FromFfi>::convert(Some(s)))
            }
        }
    };
}

impl_optional_from!(&'s [u8]);
impl_optional_from!(&'s BStr);
impl_optional_from!(&'s str);

macro_rules! impl_optional_into {
    ($ty:ty, $null:expr) => {
        unsafe impl IntoFfi for Option<$ty> {
            type Into = <$ty as IntoFfi>::Into;

            fn convert(self) -> Self::Into {
                self.map_or($null, <$ty as IntoFfi>::convert)
            }

            fn postlude(ret: &str) -> impl Display {
                display!(
                    "if {ret}.__ptr == nil then {ret} = nil; else {}end; ",
                    <$ty as IntoFfi>::postlude(ret)
                )
            }
        }
    };
}

impl_optional_into!(&'static [u8], lua_buf::null());
impl_optional_into!(&'static BStr, lua_buf::null());
impl_optional_into!(&'static str, lua_buf::null());
impl_optional_into!(Vec<u8>, lua_buffer::null());
impl_optional_into!(BString, lua_buffer::null());
impl_optional_into!(String, lua_buffer::null());