tikz-gallery-generator

Custum build of stapix for tikz.pablopie.xyz

Commit
35c597cce7b7d0fa36975ae442b93e3568c04551
Parent
cf11502b59177ed6a54c43af3a454fab8ef50c0f
Author
Pablo <pablo-pie@riseup.net>
Date

Added SIMD-optimized HTML escaping

Diffstats

3 files changed, 201 insertions, 21 deletions

Status Name Changes Insertions Deletions
Modified TODO.md 2 files changed 0 1
Added src/escape.rs 1 file changed 198 0
Modified src/main.rs 2 files changed 3 20
diff --git a/TODO.md b/TODO.md
@@ -1,6 +1,5 @@
 # TODO
 
-* use SIMD for HTML escaping
 * refactor the page visuals
   * Also make the page root configurable
 * add a flag to configure the size of the thread pool
diff --git /dev/null b/src/escape.rs
@@ -0,0 +1,198 @@
+//! HTML Escaping
+//!
+//! Stolen from pulldown-cmark-escape
+//! <https://github.com/pulldown-cmark/pulldown-cmark/>
+
+use std::fmt::{self, Display};
+
+const ESCAPE_TABLE: [Option<&str>; 256] = create_escape_table();
+const fn create_escape_table() -> [Option<&'static str>; 256] {
+  let mut table = [None; 256];
+  table[b'<'  as usize] = Some("&lt;");
+  table[b'>'  as usize] = Some("&gt;");
+  table[b'&'  as usize] = Some("&amp;");
+  table[b'"'  as usize] = Some("&quot;");
+  table[b'\'' as usize] = Some("&apos;");
+  table
+}
+
+/// A wrapper for HTML-escaped strings
+pub struct Escaped<'a>(pub &'a str);
+
+// stolen from pulldown-cmark-escape
+impl Display for Escaped<'_> {
+  #[cfg(target_arch = "x86_64")]
+  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    // the SIMD accelerated code uses the PSHUFB instruction, which is part
+    // of the SSSE3 instruction set
+    if is_x86_feature_detected!("ssse3") {
+      simd::fmt_escaped_html(self.0, f)
+    } else {
+      fmt_escaped_html_scalar(self.0, f)
+    }
+  }
+
+  #[cfg(not(target_arch = "x86_64"))]
+  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    fmt_escaped_html_scalar(self.0, f)
+  }
+}
+
+// stolen from pulldown-cmark-escape
+fn fmt_escaped_html_scalar(
+  s: &str,
+  f: &mut fmt::Formatter<'_>
+) -> fmt::Result {
+  let bytes = s.as_bytes();
+  let mut mark = 0;
+  let mut i = 0;
+
+  while i < s.len() {
+    let next_escaped = bytes[i..]
+      .iter()
+      .enumerate()
+      .find_map(|(offset, c)| Some((offset, ESCAPE_TABLE[*c as usize]?)));
+
+    if let Some((offset, escape_seq)) = next_escaped {
+      i += offset;
+      f.write_str(&s[mark..i])?;
+      f.write_str(escape_seq)?;
+
+      i += 1;
+      mark = i; // all escaped characters are ASCII
+    } else {
+      break;
+    }
+  }
+
+  f.write_str(&s[mark..])
+}
+
+// stolen from pulldown-cmark-escape
+#[cfg(target_arch = "x86_64")]
+mod simd {
+  use std::{arch::x86_64::*, mem, fmt};
+
+  const VECTOR_SIZE: usize = mem::size_of::<__m128i>();
+
+  #[inline]
+  pub fn fmt_escaped_html(s: &str, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    // the strategy here is to walk s in chunks of VECTOR_SIZE (16) bytes at
+    // a time:
+    //
+    // 1. for each chunk, we compute a a bitmask indicating whether the
+    //    corresponding byte is a HTML special byte
+    // 2. for each bit set in this mask, we print the escaped character
+    //    accordingly, as well as the surrounding characters that don't need
+    //    escaping
+    //
+    // when the number of HTML special bytes in the buffer is relatively low,
+    // this allows us to quickly go through the buffer without a lookup and
+    // for every single byte
+    if s.len() < VECTOR_SIZE {
+      return super::fmt_escaped_html_scalar(s, f);
+    }
+
+    let bytes = s.as_bytes();
+    let mut mark = 0;
+    let mut offset = 0;
+
+    unsafe {
+      let upperbound = bytes.len() - VECTOR_SIZE;
+      while offset < upperbound {
+        let mut mask = compute_mask(bytes, offset);
+
+        while mask != 0 {
+          let first_special = mask.trailing_zeros();
+          let i = offset + first_special as usize;
+          let c = *bytes.get_unchecked(i) as usize;
+
+          // here we know c = s[i] is a character that should be escaped,
+          // so it is safe to unwrap ESCAPE_TABLE[c]
+          let escape_seq = super::ESCAPE_TABLE[c].unwrap();
+          f.write_str(s.get_unchecked(mark..i))?;
+          f.write_str(escape_seq)?;
+
+          mark = i + 1; // all escaped characters are ASCII
+          mask ^= mask & -mask;
+        }
+
+        offset += VECTOR_SIZE;
+      }
+
+      // ======================================================================
+      // final iteration: we align the read with the end of the slice
+      // and shift off the bytes at start we have already scanned
+      let mut mask = compute_mask(bytes, upperbound);
+      mask >>= offset - upperbound;
+
+      while mask != 0 {
+        let first_special = mask.trailing_zeros();
+        let i = offset + first_special as usize;
+        let c = *bytes.get_unchecked(i) as usize;
+
+        // here we know c = s[i] is a character that should be escaped,
+        // so it is safe to unwrap ESCAPE_TABLE[c]
+        let escape_seq = super::ESCAPE_TABLE[c].unwrap();
+        f.write_str(s.get_unchecked(mark..i))?;
+        f.write_str(escape_seq)?;
+
+        mark = i + 1; // all escaped characters are ASCII
+        mask ^= mask & -mask;
+      }
+
+      f.write_str(s.get_unchecked(mark..))
+    }
+  }
+
+
+  #[inline]
+  #[target_feature(enable = "ssse3")]
+  /// Computes a byte mask at given offset in the byte buffer. Its first 16
+  /// (least significant) bits correspond to whether there is an HTML special
+  /// byte at the first VECTOR_SIZE bytes `bytes[offset..]`.
+  ///
+  /// It is only safe to call this function when `bytes.len() >= offset +
+  /// VECTOR_SIZE`.
+  unsafe fn compute_mask(bytes: &[u8], offset: usize) -> i32 {
+    debug_assert!(bytes.len() >= offset + VECTOR_SIZE);
+
+    const LOOKUP_TABLE: [u8; VECTOR_SIZE] = create_lookup();
+    const fn create_lookup() -> [u8; VECTOR_SIZE] {
+      let mut table = [0; VECTOR_SIZE];
+      table[(b'<'  & 0x0f) as usize] = b'<';
+      table[(b'>'  & 0x0f) as usize] = b'>';
+      table[(b'&'  & 0x0f) as usize] = b'&';
+      table[(b'"'  & 0x0f) as usize] = b'"';
+      table[(b'\'' & 0x0f) as usize] = b'\'';
+      table[0]                       = 0b01111111;
+      table
+    }
+
+    let lookup_table = _mm_loadu_si128(
+      LOOKUP_TABLE.as_ptr() as *const __m128i
+    );
+    let raw_ptr = bytes.as_ptr().add(offset) as *const __m128i;
+    let vector = _mm_loadu_si128(raw_ptr);
+
+    // mask the vector using the lookup table:
+    //
+    // 1. bytes whose lower nibbles are special HTML characters get mapped to
+    //    their lower nibbles
+    // 2. bytes whose lower nibbles are nonzero and *not* special HTML
+    //    characters get mapped to 0
+    // 3. bytes whose lower nibbles are 0 get mapped to 0b01111111
+    let masked = _mm_shuffle_epi8(lookup_table, vector);
+
+    // compare the original vector to the masked one:
+    //
+    // 1. bytes that shared a lower nibble with an HTML special byte match
+    //    *only* if they are that special byte
+    // 2. all other bytes will never match
+    let matches = _mm_cmpeq_epi8(masked, vector);
+
+    // translate matches to a bitmask: every 1 corresponds to a HTML
+    // special character and a 0 is a non-HTML byte
+    _mm_movemask_epi8(matches)
+  }
+}
diff --git a/src/main.rs b/src/main.rs
@@ -12,18 +12,18 @@ use std::{
 };
 use gallery_entry::{GalleryEntry, FileFormat, LicenseType};
 use threadpool::ThreadPool;
+
 use log::JobListLogger;
+use escape::Escaped;
 
 #[macro_use]
 mod log;
+mod escape;
 mod gallery_entry;
 
 /// A wrapper for displaying the path for the thumbnail of a given path
 pub struct ThumbPath<'a>(pub &'a GalleryEntry);
 
-/// A wrapper for HTML-escaped strings
-pub struct Escaped<'a>(pub &'a str);
-
 /// A wrapper to display lists of command line arguments
 struct ArgList<'a>(pub &'a [String]);
 
@@ -568,23 +568,6 @@ impl From> for PathBuf {
   }
 }
 
-impl Display for Escaped<'_> {
-  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
-    for c in self.0.chars() {
-      match c {
-        '<'  => write!(f, "&lt;")?,
-        '>'  => write!(f, "&gt;")?,
-        '&'  => write!(f, "&amp;")?,
-        '"'  => write!(f, "&quot;")?,
-        '\'' => write!(f, "&apos;")?,
-        c    => c.fmt(f)?,
-      }
-    }
-
-    Ok(())
-  }
-}
-
 impl Display for ArgList<'_> {
   fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
     let mut first = true;