feat: register_escape_fn API
This commit is contained in:
parent
0e99d66840
commit
bc4bb47082
|
@ -14,20 +14,20 @@ const VECTOR_BYTES: usize = std::mem::size_of::<__m256i>();
|
||||||
const VECTOR_ALIGN: usize = VECTOR_BYTES - 1;
|
const VECTOR_ALIGN: usize = VECTOR_BYTES - 1;
|
||||||
|
|
||||||
#[target_feature(enable = "avx2")]
|
#[target_feature(enable = "avx2")]
|
||||||
pub unsafe fn escape(buffer: &mut Buffer, bytes: &[u8]) {
|
pub unsafe fn escape(feed: &str, buffer: &mut Buffer) {
|
||||||
let len = bytes.len();
|
let len = feed.len();
|
||||||
|
|
||||||
if len < 8 {
|
if len < 8 {
|
||||||
let start_ptr = bytes.as_ptr();
|
let start_ptr = feed.as_ptr();
|
||||||
let end_ptr = start_ptr.add(len);
|
let end_ptr = start_ptr.add(len);
|
||||||
naive::escape(buffer, start_ptr, start_ptr, end_ptr);
|
naive::escape(buffer, start_ptr, start_ptr, end_ptr);
|
||||||
return;
|
return;
|
||||||
} else if len < VECTOR_BYTES {
|
} else if len < VECTOR_BYTES {
|
||||||
sse2::escape(buffer, bytes);
|
sse2::escape(feed, buffer);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut start_ptr = bytes.as_ptr();
|
let mut start_ptr = feed.as_ptr();
|
||||||
let end_ptr = start_ptr.add(len);
|
let end_ptr = start_ptr.add(len);
|
||||||
|
|
||||||
let v_independent1 = _mm256_set1_epi8(5);
|
let v_independent1 = _mm256_set1_epi8(5);
|
||||||
|
|
|
@ -39,12 +39,12 @@ fn contains_key(x: usize) -> bool {
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
pub unsafe fn escape(buffer: &mut Buffer, bytes: &[u8]) {
|
pub unsafe fn escape(feed: &str, buffer: &mut Buffer) {
|
||||||
let len = bytes.len();
|
let len = feed.len();
|
||||||
let mut start_ptr = bytes.as_ptr();
|
let mut start_ptr = feed.as_ptr();
|
||||||
let end_ptr = start_ptr.add(len);
|
let end_ptr = start_ptr.add(len);
|
||||||
|
|
||||||
if bytes.len() < USIZE_BYTES {
|
if feed.len() < USIZE_BYTES {
|
||||||
naive::escape(buffer, start_ptr, start_ptr, end_ptr);
|
naive::escape(buffer, start_ptr, start_ptr, end_ptr);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
|
@ -7,8 +7,13 @@ mod fallback;
|
||||||
mod naive;
|
mod naive;
|
||||||
mod sse2;
|
mod sse2;
|
||||||
|
|
||||||
|
use std::mem;
|
||||||
|
use std::sync::atomic::{AtomicPtr, Ordering};
|
||||||
|
|
||||||
use super::buffer::Buffer;
|
use super::buffer::Buffer;
|
||||||
|
|
||||||
|
type FnRaw = *mut ();
|
||||||
|
|
||||||
static ESCAPE_LUT: [u8; 256] = [
|
static ESCAPE_LUT: [u8; 256] = [
|
||||||
9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9,
|
9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9,
|
||||||
9, 9, 9, 9, 9, 9, 0, 9, 9, 9, 1, 2, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9,
|
9, 9, 9, 9, 9, 9, 0, 9, 9, 9, 1, 2, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9,
|
||||||
|
@ -25,38 +30,37 @@ static ESCAPE_LUT: [u8; 256] = [
|
||||||
const ESCAPED: [&str; 5] = [""", "&", "'", "<", ">"];
|
const ESCAPED: [&str; 5] = [""", "&", "'", "<", ">"];
|
||||||
const ESCAPED_LEN: usize = 5;
|
const ESCAPED_LEN: usize = 5;
|
||||||
|
|
||||||
|
static FN: AtomicPtr<()> = AtomicPtr::new(escape as FnRaw);
|
||||||
|
|
||||||
#[cfg(target_feature = "avx2")]
|
#[cfg(target_feature = "avx2")]
|
||||||
pub(crate) fn escape_to_buf(feed: &str, buf: &mut Buffer) {
|
pub fn escape(feed: &str, buf: &mut Buffer) {
|
||||||
unsafe { avx2::escape(buf, feed.as_bytes()) }
|
unsafe { avx2::escape(buf, feed.as_bytes()) }
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// default escape function
|
||||||
#[cfg(not(target_feature = "avx2"))]
|
#[cfg(not(target_feature = "avx2"))]
|
||||||
pub(crate) fn escape_to_buf(feed: &str, buf: &mut Buffer) {
|
pub fn escape(feed: &str, buf: &mut Buffer) {
|
||||||
use std::mem;
|
|
||||||
use std::sync::atomic::{AtomicPtr, Ordering};
|
|
||||||
|
|
||||||
type FnRaw = *mut ();
|
|
||||||
|
|
||||||
static FN: AtomicPtr<()> = AtomicPtr::new(detect as FnRaw);
|
|
||||||
|
|
||||||
fn detect(buffer: &mut Buffer, bytes: &[u8]) {
|
|
||||||
let fun = if is_x86_feature_detected!("avx2") {
|
let fun = if is_x86_feature_detected!("avx2") {
|
||||||
avx2::escape as FnRaw
|
avx2::escape
|
||||||
} else if is_x86_feature_detected!("sse2") {
|
} else if is_x86_feature_detected!("sse2") {
|
||||||
sse2::escape as FnRaw
|
sse2::escape
|
||||||
} else {
|
} else {
|
||||||
fallback::escape as FnRaw
|
fallback::escape
|
||||||
};
|
};
|
||||||
|
|
||||||
FN.store(fun as FnRaw, Ordering::Relaxed);
|
FN.store(fun as FnRaw, Ordering::Relaxed);
|
||||||
unsafe {
|
unsafe { fun(feed, buf) };
|
||||||
mem::transmute::<FnRaw, fn(&mut Buffer, &[u8])>(fun)(buffer, bytes);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn register_escape_fn(fun: fn(&str, &mut Buffer)) {
|
||||||
|
FN.store(fun as FnRaw, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
pub(crate) fn escape_to_buf(feed: &str, buf: &mut Buffer) {
|
||||||
unsafe {
|
unsafe {
|
||||||
let fun = FN.load(Ordering::Relaxed);
|
let fun = FN.load(Ordering::Relaxed);
|
||||||
mem::transmute::<FnRaw, fn(&mut Buffer, &[u8])>(fun)(buf, feed.as_bytes());
|
mem::transmute::<FnRaw, fn(&str, &mut Buffer)>(fun)(feed, buf);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -157,8 +161,8 @@ mod tests {
|
||||||
buf3.clear();
|
buf3.clear();
|
||||||
|
|
||||||
unsafe {
|
unsafe {
|
||||||
escape_to_buf(&*s, &mut buf1);
|
escape_to_buf(s, &mut buf1);
|
||||||
fallback::escape(&mut buf2, s.as_bytes());
|
fallback::escape(s, &mut buf2);
|
||||||
naive::escape(&mut buf3, s.as_ptr(), s.as_ptr(), s.as_ptr().add(s.len()));
|
naive::escape(&mut buf3, s.as_ptr(), s.as_ptr(), s.as_ptr().add(s.len()));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -15,16 +15,16 @@ const VECTOR_ALIGN: usize = VECTOR_BYTES - 1;
|
||||||
|
|
||||||
#[target_feature(enable = "sse2")]
|
#[target_feature(enable = "sse2")]
|
||||||
#[inline]
|
#[inline]
|
||||||
pub unsafe fn escape(buffer: &mut Buffer, bytes: &[u8]) {
|
pub unsafe fn escape(feed: &str, buffer: &mut Buffer) {
|
||||||
let len = bytes.len();
|
let len = feed.len();
|
||||||
if len < VECTOR_BYTES {
|
if len < VECTOR_BYTES {
|
||||||
let start_ptr = bytes.as_ptr();
|
let start_ptr = feed.as_ptr();
|
||||||
let end_ptr = start_ptr.add(len);
|
let end_ptr = start_ptr.add(len);
|
||||||
naive::escape(buffer, start_ptr, start_ptr, end_ptr);
|
naive::escape(buffer, start_ptr, start_ptr, end_ptr);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
let mut start_ptr = bytes.as_ptr();
|
let mut start_ptr = feed.as_ptr();
|
||||||
let end_ptr = start_ptr.add(len);
|
let end_ptr = start_ptr.add(len);
|
||||||
|
|
||||||
let v_independent1 = _mm_set1_epi8(5);
|
let v_independent1 = _mm_set1_epi8(5);
|
||||||
|
|
Loading…
Reference in New Issue