Skip to main content

axhash_ffi/
lib.rs

1#![allow(clippy::not_unsafe_ptr_arg_deref)]
2
3use axhash_core::hasher::AxHasher;
4use axhash_core::hasher::api::{axhash, axhash_seeded};
5use axhash_core::{RuntimeBackend, runtime_backend, runtime_has_simd};
6use core::ffi::c_char;
7use core::hash::Hasher;
8use core::ptr::NonNull;
9
10extern crate alloc;
11use alloc::boxed::Box;
12
13pub enum AxHashState {}
14
15#[repr(C)]
16pub struct AxHashIovec {
17    pub ptr: *const u8,
18    pub len: usize,
19}
20
21#[non_exhaustive]
22#[repr(C)]
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24pub enum AxHashRuntimeBackend {
25    Scalar = 0,
26    Aarch64Neon = 1,
27    X86_64Avx2 = 2,
28}
29
30const MAX_BATCH: usize = 1 << 16;
31
32#[inline(always)]
33fn as_state_ptr(hasher: AxHasher) -> *mut AxHashState {
34    Box::into_raw(Box::new(hasher)).cast()
35}
36
37#[inline(always)]
38unsafe fn state_mut<'a>(state: *mut AxHashState) -> Option<&'a mut AxHasher> {
39    NonNull::new(state.cast()).map(|p| unsafe { &mut *p.as_ptr() })
40}
41
42#[inline(always)]
43unsafe fn slice_from_raw<'a>(ptr: *const u8, len: usize) -> &'a [u8] {
44    unsafe { core::slice::from_raw_parts(ptr, len) }
45}
46
47#[inline(always)]
48fn is_invalid_input(ptr: *const u8, len: usize) -> bool {
49    len != 0 && ptr.is_null()
50}
51
52impl From<RuntimeBackend> for AxHashRuntimeBackend {
53    #[inline(always)]
54    fn from(b: RuntimeBackend) -> Self {
55        match b {
56            RuntimeBackend::Scalar => Self::Scalar,
57            RuntimeBackend::Aarch64Neon => Self::Aarch64Neon,
58            RuntimeBackend::X86_64Avx2 => Self::X86_64Avx2,
59            _ => Self::Scalar,
60        }
61    }
62}
63
64#[cold]
65fn fail_u64() -> u64 {
66    0
67}
68
69#[unsafe(no_mangle)]
70pub extern "C" fn axhash_ffi_version() -> *const c_char {
71    static VERSION: &[u8] = concat!(env!("CARGO_PKG_VERSION"), "\0").as_bytes();
72    VERSION.as_ptr().cast()
73}
74
75#[unsafe(no_mangle)]
76pub extern "C" fn axhash_bytes(bytes: *const u8, len: usize) -> u64 {
77    if is_invalid_input(bytes, len) {
78        return fail_u64();
79    }
80
81    if len == 0 {
82        return axhash(&[]);
83    }
84
85    unsafe { axhash(slice_from_raw(bytes, len)) }
86}
87
88#[unsafe(no_mangle)]
89pub extern "C" fn axhash_bytes_seeded(bytes: *const u8, len: usize, seed: u64) -> u64 {
90    if is_invalid_input(bytes, len) {
91        return fail_u64();
92    }
93
94    if len == 0 {
95        return axhash_seeded(&[], seed);
96    }
97
98    unsafe { axhash_seeded(slice_from_raw(bytes, len), seed) }
99}
100
101#[unsafe(no_mangle)]
102pub extern "C" fn axhash_batch_seeded(
103    iovecs: *const AxHashIovec,
104    count: usize,
105    seed: u64,
106    out_hashes: *mut u64,
107) {
108    if iovecs.is_null() || out_hashes.is_null() || count == 0 {
109        return;
110    }
111
112    let count = count.min(MAX_BATCH);
113
114    let jobs = unsafe { core::slice::from_raw_parts(iovecs, count) };
115    let outs = unsafe { core::slice::from_raw_parts_mut(out_hashes, count) };
116
117    for (job, out) in jobs.iter().zip(outs.iter_mut()) {
118        *out = if is_invalid_input(job.ptr, job.len) {
119            0
120        } else if job.len == 0 {
121            axhash_seeded(&[], seed)
122        } else {
123            unsafe { axhash_seeded(slice_from_raw(job.ptr, job.len), seed) }
124        };
125    }
126}
127
128#[unsafe(no_mangle)]
129pub extern "C" fn axhash_hasher_new() -> *mut AxHashState {
130    as_state_ptr(AxHasher::new())
131}
132
133#[unsafe(no_mangle)]
134pub extern "C" fn axhash_hasher_new_seeded(seed: u64) -> *mut AxHashState {
135    as_state_ptr(AxHasher::new_with_seed(seed))
136}
137
138#[unsafe(no_mangle)]
139pub extern "C" fn axhash_hasher_reset(state: *mut AxHashState, seed: u64) -> bool {
140    let hasher = unsafe { state_mut(state) };
141    let Some(hasher) = hasher else {
142        return false;
143    };
144
145    *hasher = AxHasher::new_with_seed(seed);
146    true
147}
148
149#[unsafe(no_mangle)]
150pub extern "C" fn axhash_hasher_write(
151    state: *mut AxHashState,
152    bytes: *const u8,
153    len: usize,
154) -> bool {
155    if is_invalid_input(bytes, len) {
156        return false;
157    }
158
159    if len == 0 {
160        return true;
161    }
162
163    let hasher = unsafe { state_mut(state) };
164
165    let Some(hasher) = hasher else {
166        return false;
167    };
168
169    unsafe {
170        hasher.write(slice_from_raw(bytes, len));
171    }
172
173    true
174}
175
176#[unsafe(no_mangle)]
177pub extern "C" fn axhash_hasher_finish(state: *mut AxHashState) -> u64 {
178    let hasher = unsafe { state_mut(state) };
179    let Some(hasher) = hasher else {
180        return 0;
181    };
182
183    hasher.finish()
184}
185
186#[unsafe(no_mangle)]
187pub extern "C" fn axhash_hasher_free(state: *mut AxHashState) {
188    if state.is_null() {
189        return;
190    }
191
192    unsafe {
193        drop(Box::from_raw(state.cast::<AxHasher>()));
194    }
195}
196
197#[unsafe(no_mangle)]
198pub extern "C" fn axhash_runtime_backend() -> AxHashRuntimeBackend {
199    runtime_backend().into()
200}
201
202#[unsafe(no_mangle)]
203pub extern "C" fn axhash_runtime_has_simd() -> bool {
204    runtime_has_simd()
205}