superfluid_engine/
tokenizer_abi.rs

1//! A [`Tokenizer`] across a shared-library boundary.
2
3use std::ffi::{c_char, c_void, CStr};
4use std::path::Path;
5
6use crate::Tokenizer;
7
8pub const ABI_VERSION: u32 = 1;
9
10pub mod symbol {
11    pub const ABI_VERSION: &[u8] = b"superfluid_tokenizer_abi_version\0";
12    pub const OPEN: &[u8] = b"superfluid_tokenizer_open\0";
13    pub const CLOSE: &[u8] = b"superfluid_tokenizer_close\0";
14    pub const ENCODE: &[u8] = b"superfluid_tokenizer_encode\0";
15    pub const TOKEN_BYTES: &[u8] = b"superfluid_tokenizer_token_bytes\0";
16    pub const VOCAB_SIZE: &[u8] = b"superfluid_tokenizer_vocab_size\0";
17    pub const SPECIAL_TOKENS: &[u8] = b"superfluid_tokenizer_special_tokens\0";
18    pub const BOS: &[u8] = b"superfluid_tokenizer_bos\0";
19    pub const EOS: &[u8] = b"superfluid_tokenizer_eos\0";
20    pub const CHAT_TEMPLATE: &[u8] = b"superfluid_tokenizer_chat_template\0";
21}
22
23pub type AbiVersionFn = unsafe extern "C" fn() -> u32;
24pub type OpenFn = unsafe extern "C" fn(path: *const c_char, err: *mut c_char, err_len: usize) -> *mut c_void;
25pub type CloseFn = unsafe extern "C" fn(tok: *mut c_void);
26pub type EncodeFn = unsafe extern "C" fn(tok: *const c_void, text: *const u8, len: usize, out: *mut u32, cap: usize) -> usize;
27pub type TokenBytesFn = unsafe extern "C" fn(tok: *const c_void, id: u32, out: *mut u8, cap: usize) -> usize;
28pub type VocabSizeFn = unsafe extern "C" fn(tok: *const c_void) -> u32;
29pub type SpecialTokensFn = unsafe extern "C" fn(tok: *const c_void, out: *mut u8, cap: usize) -> usize;
30pub type BosFn = unsafe extern "C" fn(tok: *const c_void) -> i64;
31pub type EosFn = unsafe extern "C" fn(tok: *const c_void) -> u32;
32pub type ChatTemplateFn = unsafe extern "C" fn(tok: *const c_void, name: *const c_char, out: *mut u8, cap: usize) -> usize;
33
34pub fn pack_special(tokens: &[(String, u32)]) -> Vec<u8> {
35    let mut out = Vec::new();
36    for (text, id) in tokens {
37        out.extend_from_slice(&id.to_le_bytes());
38        out.extend_from_slice(&(text.len() as u32).to_le_bytes());
39        out.extend_from_slice(text.as_bytes());
40    }
41    out
42}
43
44pub fn unpack_special(mut bytes: &[u8]) -> Vec<(String, u32)> {
45    let mut out = Vec::new();
46    while bytes.len() >= 8 {
47        let id = u32::from_le_bytes(bytes[0..4].try_into().expect("4 bytes"));
48        let len = u32::from_le_bytes(bytes[4..8].try_into().expect("4 bytes")) as usize;
49        let Some(text) = bytes.get(8..8 + len) else { break };
50        out.push((String::from_utf8_lossy(text).into_owned(), id));
51        bytes = &bytes[8 + len..];
52    }
53    out
54}
55
56#[doc(hidden)]
57pub mod export {
58    use super::*;
59    use std::panic::{catch_unwind, AssertUnwindSafe};
60
61    type Boxed = Box<dyn Tokenizer>;
62
63    fn fill(bytes: &[u8], out: *mut u8, cap: usize) -> usize {
64        let n = bytes.len().min(cap);
65        if n > 0 && !out.is_null() {
66            // SAFETY: the caller's buffer holds `cap` bytes; `n <= cap`.
67            unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), out, n) };
68        }
69        bytes.len()
70    }
71
72    fn write_err(err: *mut c_char, err_len: usize, msg: &str) {
73        if err.is_null() || err_len == 0 {
74            return;
75        }
76        let n = msg.len().min(err_len - 1);
77        // SAFETY: `err` holds `err_len` bytes; `n + 1 <= err_len`.
78        unsafe {
79            std::ptr::copy_nonoverlapping(msg.as_ptr(), err.cast::<u8>(), n);
80            *err.add(n) = 0;
81        }
82    }
83
84    /// # Safety
85    /// `tok` is a live handle [`open`] returned, or NULL.
86    unsafe fn tok<'a>(tok: *const c_void) -> Option<&'a dyn Tokenizer> {
87        // SAFETY: per this function's contract; NULL is refused.
88        (!tok.is_null()).then(|| unsafe { &**(tok as *const Boxed) })
89    }
90
91    /// # Safety
92    /// `path` is NUL-terminated; `err` holds `err_len` bytes (or is NULL).
93    pub unsafe fn open<T, E>(
94        path: *const c_char,
95        err: *mut c_char,
96        err_len: usize,
97        load: impl Fn(&Path) -> Result<T, E>,
98    ) -> *mut c_void
99    where
100        T: Tokenizer + 'static,
101        E: std::fmt::Display,
102    {
103        if path.is_null() {
104            write_err(err, err_len, "no path");
105            return std::ptr::null_mut();
106        }
107        // SAFETY: per this function's contract.
108        let path = unsafe { CStr::from_ptr(path) }.to_string_lossy().into_owned();
109        match catch_unwind(AssertUnwindSafe(|| load(Path::new(&path)))) {
110            Ok(Ok(t)) => Box::into_raw(Box::new(Box::new(t) as Boxed)) as *mut c_void,
111            Ok(Err(e)) => {
112                write_err(err, err_len, &e.to_string());
113                std::ptr::null_mut()
114            }
115            Err(_) => {
116                write_err(err, err_len, "the tokenizer panicked while loading");
117                std::ptr::null_mut()
118            }
119        }
120    }
121
122    /// # Safety
123    /// `t` is a live handle [`open`] returned, released once.
124    pub unsafe fn close(t: *mut c_void) {
125        if !t.is_null() {
126            // SAFETY: per this function's contract.
127            let boxed = unsafe { Box::from_raw(t as *mut Boxed) };
128            let _ = catch_unwind(AssertUnwindSafe(|| drop(boxed)));
129        }
130    }
131
132    /// # Safety
133    /// `t` is live; `text` holds `len` bytes; `out` holds `cap` ids.
134    pub unsafe fn encode(t: *const c_void, text: *const u8, len: usize, out: *mut u32, cap: usize) -> usize {
135        catch_unwind(AssertUnwindSafe(|| {
136            let bytes = if len == 0 {
137                &[][..]
138            } else {
139                // SAFETY: per this function's contract.
140                unsafe { std::slice::from_raw_parts(text, len) }
141            };
142            // SAFETY: per this function's contract.
143            let Some(tk) = (unsafe { tok(t) }) else { return 0 };
144            let ids = tk.encode(&String::from_utf8_lossy(bytes));
145            let n = ids.len().min(cap);
146            if n > 0 && !out.is_null() {
147                // SAFETY: `out` holds `cap` ids; `n <= cap`.
148                unsafe { std::ptr::copy_nonoverlapping(ids.as_ptr(), out, n) };
149            }
150            ids.len()
151        }))
152        .unwrap_or(0)
153    }
154
155    /// # Safety
156    /// `t` is live; `out` holds `cap` bytes.
157    pub unsafe fn token_bytes(t: *const c_void, id: u32, out: *mut u8, cap: usize) -> usize {
158        // SAFETY: per this function's contract.
159        catch_unwind(AssertUnwindSafe(|| unsafe { tok(t) }.map_or(0, |tk| fill(&tk.token_bytes(id), out, cap))))
160            .unwrap_or(0)
161    }
162
163    /// # Safety
164    /// `t` is live.
165    pub unsafe fn vocab_size(t: *const c_void) -> u32 {
166        // SAFETY: per this function's contract.
167        catch_unwind(AssertUnwindSafe(|| unsafe { tok(t) }.map_or(0, |tk| tk.vocab_size()))).unwrap_or(0)
168    }
169
170    /// # Safety
171    /// `t` is live; `out` holds `cap` bytes.
172    pub unsafe fn special_tokens(t: *const c_void, out: *mut u8, cap: usize) -> usize {
173        // SAFETY: per this function's contract.
174        catch_unwind(AssertUnwindSafe(|| {
175            // SAFETY: per this function's contract.
176            unsafe { tok(t) }.map_or(0, |tk| fill(&pack_special(&tk.special_tokens()), out, cap))
177        }))
178        .unwrap_or(0)
179    }
180
181    /// # Safety
182    /// `t` is live.
183    pub unsafe fn bos(t: *const c_void) -> i64 {
184        // SAFETY: per this function's contract.
185        catch_unwind(AssertUnwindSafe(|| {
186            // SAFETY: per this function's contract.
187            unsafe { tok(t) }.and_then(|tk| tk.bos_token()).map_or(-1, i64::from)
188        }))
189        .unwrap_or(-1)
190    }
191
192    /// # Safety
193    /// `t` is live.
194    pub unsafe fn eos(t: *const c_void) -> u32 {
195        // SAFETY: per this function's contract.
196        catch_unwind(AssertUnwindSafe(|| unsafe { tok(t) }.map_or(0, |tk| tk.eos_token()))).unwrap_or(0)
197    }
198
199    /// # Safety
200    /// `t` is live; `name` is NUL-terminated or NULL; `out` holds `cap` bytes.
201    pub unsafe fn chat_template(t: *const c_void, name: *const c_char, out: *mut u8, cap: usize) -> usize {
202        catch_unwind(AssertUnwindSafe(|| {
203            // SAFETY: per this function's contract.
204            let Some(tk) = (unsafe { tok(t) }) else { return usize::MAX };
205            let template = if name.is_null() {
206                Some(tk.chat_template_jinja())
207            } else {
208                // SAFETY: per this function's contract.
209                tk.chat_template_named(&unsafe { CStr::from_ptr(name) }.to_string_lossy())
210            };
211            match template {
212                Some(s) => fill(s.as_bytes(), out, cap),
213                None => usize::MAX,
214            }
215        }))
216        .unwrap_or(usize::MAX)
217    }
218}
219
220#[macro_export]
221macro_rules! export_tokenizer {
222    ($load:expr) => {
223        #[no_mangle]
224        pub extern "C" fn superfluid_tokenizer_abi_version() -> u32 {
225            $crate::tokenizer_abi::ABI_VERSION
226        }
227
228        /// # Safety
229        /// See `superfluid_engine::tokenizer_abi::OpenFn`.
230        #[no_mangle]
231        pub unsafe extern "C" fn superfluid_tokenizer_open(
232            path: *const ::std::ffi::c_char,
233            err: *mut ::std::ffi::c_char,
234            err_len: usize,
235        ) -> *mut ::std::ffi::c_void {
236            // SAFETY: forwarded under the same contract.
237            unsafe { $crate::tokenizer_abi::export::open(path, err, err_len, $load) }
238        }
239
240        /// # Safety
241        /// See `superfluid_engine::tokenizer_abi::CloseFn`.
242        #[no_mangle]
243        pub unsafe extern "C" fn superfluid_tokenizer_close(tok: *mut ::std::ffi::c_void) {
244            // SAFETY: forwarded under the same contract.
245            unsafe { $crate::tokenizer_abi::export::close(tok) }
246        }
247
248        /// # Safety
249        /// See `superfluid_engine::tokenizer_abi::EncodeFn`.
250        #[no_mangle]
251        pub unsafe extern "C" fn superfluid_tokenizer_encode(
252            tok: *const ::std::ffi::c_void,
253            text: *const u8,
254            len: usize,
255            out: *mut u32,
256            cap: usize,
257        ) -> usize {
258            // SAFETY: forwarded under the same contract.
259            unsafe { $crate::tokenizer_abi::export::encode(tok, text, len, out, cap) }
260        }
261
262        /// # Safety
263        /// See `superfluid_engine::tokenizer_abi::TokenBytesFn`.
264        #[no_mangle]
265        pub unsafe extern "C" fn superfluid_tokenizer_token_bytes(
266            tok: *const ::std::ffi::c_void,
267            id: u32,
268            out: *mut u8,
269            cap: usize,
270        ) -> usize {
271            // SAFETY: forwarded under the same contract.
272            unsafe { $crate::tokenizer_abi::export::token_bytes(tok, id, out, cap) }
273        }
274
275        /// # Safety
276        /// See `superfluid_engine::tokenizer_abi::VocabSizeFn`.
277        #[no_mangle]
278        pub unsafe extern "C" fn superfluid_tokenizer_vocab_size(tok: *const ::std::ffi::c_void) -> u32 {
279            // SAFETY: forwarded under the same contract.
280            unsafe { $crate::tokenizer_abi::export::vocab_size(tok) }
281        }
282
283        /// # Safety
284        /// See `superfluid_engine::tokenizer_abi::SpecialTokensFn`.
285        #[no_mangle]
286        pub unsafe extern "C" fn superfluid_tokenizer_special_tokens(
287            tok: *const ::std::ffi::c_void,
288            out: *mut u8,
289            cap: usize,
290        ) -> usize {
291            // SAFETY: forwarded under the same contract.
292            unsafe { $crate::tokenizer_abi::export::special_tokens(tok, out, cap) }
293        }
294
295        /// # Safety
296        /// See `superfluid_engine::tokenizer_abi::BosFn`.
297        #[no_mangle]
298        pub unsafe extern "C" fn superfluid_tokenizer_bos(tok: *const ::std::ffi::c_void) -> i64 {
299            // SAFETY: forwarded under the same contract.
300            unsafe { $crate::tokenizer_abi::export::bos(tok) }
301        }
302
303        /// # Safety
304        /// See `superfluid_engine::tokenizer_abi::EosFn`.
305        #[no_mangle]
306        pub unsafe extern "C" fn superfluid_tokenizer_eos(tok: *const ::std::ffi::c_void) -> u32 {
307            // SAFETY: forwarded under the same contract.
308            unsafe { $crate::tokenizer_abi::export::eos(tok) }
309        }
310
311        /// # Safety
312        /// See `superfluid_engine::tokenizer_abi::ChatTemplateFn`.
313        #[no_mangle]
314        pub unsafe extern "C" fn superfluid_tokenizer_chat_template(
315            tok: *const ::std::ffi::c_void,
316            name: *const ::std::ffi::c_char,
317            out: *mut u8,
318            cap: usize,
319        ) -> usize {
320            // SAFETY: forwarded under the same contract.
321            unsafe { $crate::tokenizer_abi::export::chat_template(tok, name, out, cap) }
322        }
323    };
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329
330    struct PanicsOnDrop;
331
332    impl Tokenizer for PanicsOnDrop {
333        fn encode(&self, _text: &str) -> Vec<u32> {
334            Vec::new()
335        }
336        fn token_bytes(&self, _token: u32) -> Vec<u8> {
337            Vec::new()
338        }
339        fn vocab_size(&self) -> u32 {
340            1
341        }
342        fn special_tokens(&self) -> Vec<(String, u32)> {
343            Vec::new()
344        }
345        fn bos_token(&self) -> Option<u32> {
346            None
347        }
348        fn eos_token(&self) -> u32 {
349            0
350        }
351        fn chat_template_jinja(&self) -> String {
352            String::new()
353        }
354    }
355
356    impl Drop for PanicsOnDrop {
357        fn drop(&mut self) {
358            panic!("a tokenizer that panics on drop");
359        }
360    }
361
362    #[test]
363    fn closing_a_tokenizer_whose_drop_panics_does_not_unwind_into_the_caller() {
364        // SAFETY: a NUL-terminated path and no error buffer; the handle is
365        // closed once.
366        unsafe {
367            let h = export::open(c"unused".as_ptr(), std::ptr::null_mut(), 0, |_| Ok::<_, String>(PanicsOnDrop));
368            assert!(!h.is_null());
369            assert_eq!(export::vocab_size(h), 1);
370            export::close(h);
371        }
372    }
373
374    #[test]
375    fn a_null_handle_answers_each_call_with_its_failure_value() {
376        let null = std::ptr::null();
377        // SAFETY: every call is handed NULL, which each refuses.
378        unsafe {
379            assert_eq!(export::encode(null, b"hi".as_ptr(), 2, std::ptr::null_mut(), 0), 0);
380            assert_eq!(export::token_bytes(null, 1, std::ptr::null_mut(), 0), 0);
381            assert_eq!(export::vocab_size(null), 0);
382            assert_eq!(export::special_tokens(null, std::ptr::null_mut(), 0), 0);
383            assert_eq!(export::bos(null), -1);
384            assert_eq!(export::eos(null), 0);
385            assert_eq!(export::chat_template(null, std::ptr::null(), std::ptr::null_mut(), 0), usize::MAX);
386            export::close(std::ptr::null_mut());
387        }
388    }
389
390    #[test]
391    fn special_tokens_pack_and_unpack() {
392        let t = vec![("<|im_start|>".to_string(), 151644), ("".to_string(), 7), ("é".to_string(), 3)];
393        assert_eq!(unpack_special(&pack_special(&t)), t);
394        let packed = pack_special(&t);
395        assert_eq!(unpack_special(&packed[..packed.len() - 1]).len(), 2, "a truncated tail is dropped");
396    }
397}
superfluid_engine/tokenizer_abi.rs

superfluid_engine/
tokenizer_abi.rs
