Optimize utf16 decoding

library/alloc/src/lib.rs

@@ -274,6 +274,8 @@ pub mod string;
 pub mod sync;
 #[cfg(all(not(no_global_oom_handling), not(no_rc), not(no_sync)))]
 pub mod task;
+#[cfg(not(no_global_oom_handling))]
+mod utf16_decode;
 pub mod vec;
 #[cfg(all(not(no_rc), not(no_sync), not(no_global_oom_handling)))]
 pub mod wtf8;

library/alloc/src/string.rs

@@ -725,7 +725,12 @@ impl String {
     #[cfg(not(no_global_oom_handling))]
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn from_utf16(v: &[u16]) -> Result<String, FromUtf16Error> {
-        Self::from_utf16_units(v.iter().cloned(), v.len())
+        let mut out = Vec::with_capacity(v.len());
+        match crate::utf16_decode::decode_utf16_into(v, &mut out, |_| Err(())) {
+            // SAFETY: Lone surrogates return `Err`; successful output is UTF-8.
+            Ok(()) => Ok(unsafe { String::from_utf8_unchecked(out) }),
+            Err(()) => Err(FromUtf16Error { kind: FromUtf16ErrorKind::LoneSurrogate }),
+        }
     }
 
     /// Decodes an iterator of UTF-16 code units into a `String`, returning
@@ -776,9 +781,12 @@ impl String {
     #[inline]
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn from_utf16_lossy(v: &[u16]) -> String {
-        char::decode_utf16(v.iter().cloned())
-            .map(|r| r.unwrap_or(char::REPLACEMENT_CHARACTER))
-            .collect()
+        let mut out = Vec::with_capacity(v.len());
+        let _ = crate::utf16_decode::decode_utf16_into(v, &mut out, |_| {
+            Ok(char::REPLACEMENT_CHARACTER as u16)
+        });
+        // SAFETY: Lone surrogates are replaced with `char::REPLACEMENT_CHARACTER`.
+        unsafe { String::from_utf8_unchecked(out) }
     }
 
     /// Decode a UTF-16LE–encoded vector `v` into a `String`,

library/alloc/src/utf16_decode.rs

@@ -0,0 +1,202 @@
+use core::mem::MaybeUninit;
+
+use crate::vec::Vec;
+
+const BLOCK: usize = 16;
+// Three bytes per unit, plus one for the last store or a pair crossing the block.
+const ENCODE_BLOCK_LEN: usize = 3 * BLOCK + 1;
+
+/// Appends generalized UTF-8; `on_unpaired` handles lone surrogates.
+/// Callers must preserve their encoding invariant. On error, `out` is unspecified.
+pub(crate) fn decode_utf16_into(
+    v: &[u16],
+    out: &mut Vec<u8>,
+    mut on_unpaired: impl FnMut(u16) -> Result<u16, ()>,
+) -> Result<(), ()> {
+    let mut i = 0;
+    while let Some(units) = v.get(i..i + BLOCK) {
+        let units: &[u16; BLOCK] = units.try_into().unwrap();
+        let mut max = 0;
+        let mut surrogates = false;
+        for &u in units {
+            max = max.max(u);
+            surrogates |= (u & 0xF800) == 0xD800;
+        }
+        let decoded_len = if max < 0x80 {
+            let rest = &v[i..];
+            out.reserve(rest.len());
+
+            let ascii_len = ascii_prefix(rest, out.spare_capacity_mut());
+            // SAFETY: `ascii_prefix` initialized `ascii_len` bytes in the reserved space.
+            unsafe { out.set_len(out.len() + ascii_len) };
+            ascii_len
+        } else if surrogates {
+            out.reserve(ENCODE_BLOCK_LEN);
+            // SAFETY: We reserved `3 * BLOCK + 1` bytes.
+            let (len, written) = unsafe {
+                decode_scalar(v, i, i + BLOCK, out.spare_capacity_mut(), &mut on_unpaired)
+            }?;
+            // SAFETY: `decode_scalar` initialized `written` bytes in the reserved space.
+            unsafe { out.set_len(out.len() + written) };
+            len
+        } else {
+            out.reserve(v.len() - i + ENCODE_BLOCK_LEN);
+            let len = encode_block(units, out.spare_capacity_mut());
+            // SAFETY: `encode_block` initialized `len` bytes in the reserved space.
+            unsafe {
+                out.set_len(out.len() + len);
+            };
+            BLOCK
+        };
+        i += decoded_len;
+    }
+    decode_small(&v[i..], out, &mut on_unpaired)?;
+    Ok(())
+}
+
+fn decode_small(
+    v: &[u16],
+    out: &mut Vec<u8>,
+    on_unpaired: &mut impl FnMut(u16) -> Result<u16, ()>,
+) -> Result<(), ()> {
+    let mut i = 0;
+    while i < v.len() {
+        let u = v[i];
+        i += 1;
+        let c = if u < 0x80 {
+            out.push(u as u8);
+            continue;
+        } else if (u & 0xF800) != 0xD800 {
+            u as u32
+        } else if u < 0xDC00 && i < v.len() && (v[i] & 0xFC00) == 0xDC00 {
+            let c = 0x10000 + (((u as u32) & 0x3FF) << 10) + ((v[i] as u32) & 0x3FF);
+            i += 1;
+            c
+        } else {
+            on_unpaired(u)? as u32
+        };
+        push_code_point(out, c);
+    }
+    Ok(())
+}
+
+#[inline]
+fn push_code_point(out: &mut Vec<u8>, code: u32) {
+    let n = 1 + (code >= 0x80) as usize + (code >= 0x800) as usize + (code >= 0x10000) as usize;
+    let len = out.len();
+    out.reserve(n);
+
+    // SAFETY: We reserved `n` bytes, and the encoder initializes all of them.
+    unsafe {
+        core::char::encode_utf8_raw_unchecked(code, out.as_mut_ptr().add(len));
+        out.set_len(len + n);
+    }
+}
+
+fn ascii_prefix(src: &[u16], out: &mut [MaybeUninit<u8>]) -> usize {
+    const CHUNK_SIZE: usize = 32;
+    let mut i = 0;
+    while let Some(units) = src.get(i..i + CHUNK_SIZE) {
+        let mut max: u16 = 0;
+        let units: &[u16; CHUNK_SIZE] = units.try_into().unwrap();
+        for j in 0..CHUNK_SIZE {
+            max = max.max(units[j]);
+        }
+        if max >= 0x80 {
+            break;
+        }
+        let dst: &mut [MaybeUninit<u8>; CHUNK_SIZE] =
+            (&mut out[i..i + CHUNK_SIZE]).try_into().unwrap();
+        for j in 0..CHUNK_SIZE {
+            dst[j].write(units[j] as u8);
+        }
+        i += CHUNK_SIZE;
+    }
+    while i < src.len() && src[i] < 0x80 {
+        out[i].write(src[i] as u8);
+        i += 1;
+    }
+    i
+}
+
+/// Decodes `v[start..end]`, returning `(units_consumed, bytes_written)`.
+/// A pair crossing `end` consumes one extra unit.
+///
+/// # Safety
+///
+/// `out.len()` must be at least `3 * (end - start) + 1`.
+unsafe fn decode_scalar(
+    v: &[u16],
+    start: usize,
+    end: usize,
+    out: &mut [MaybeUninit<u8>],
+    on_unpaired: &mut impl FnMut(u16) -> Result<u16, ()>,
+) -> Result<(usize, usize), ()> {
+    let mut i = start;
+    let mut written = 0;
+    while i < end {
+        let u = v[i];
+        i += 1;
+        let (word, n) = if u < 0x80 {
+            (u as u32, 1)
+        } else if (u & 0xF800) != 0xD800 {
+            encode_bmp(u)
+        } else if u < 0xDC00 && i < v.len() && (v[i] & 0xFC00) == 0xDC00 {
+            let c = 0x10000 + (((u as u32) & 0x3FF) << 10) + ((v[i] as u32) & 0x3FF);
+            let cont = |shift: u32| 0x80 | (c >> shift & 0x3F);
+            i += 1;
+            ((0xF0 | c >> 18) | cont(12) << 8 | cont(6) << 16 | cont(0) << 24, 4)
+        } else {
+            let Ok(unit) = on_unpaired(u) else {
+                return Err(());
+            };
+            encode_bmp(unit)
+        };
+
+        // SAFETY: The caller's output bound covers the final store and a pair crossing `end`.
+        unsafe {
+            let dst = out.get_unchecked_mut(written..).as_mut_ptr();
+            match n {
+                1 => dst.cast::<u8>().write(word as u8),
+                2 => dst.cast::<u16>().write_unaligned((word as u16).to_le()),
+                3 | 4 => dst.cast::<u32>().write_unaligned(word.to_le()),
+                _ => unreachable!(),
+            }
+        }
+        written += n;
+    }
+    Ok((i - start, written))
+}
+
+/// Packs UTF-8 bytes least-significant byte first, and returns their length.
+/// Surrogates are encoded as-is.
+#[inline]
+fn encode_bmp(u: u16) -> (u32, usize) {
+    let u = u as u32;
+    if u < 0x80 {
+        (u, 1)
+    } else if u < 0x800 {
+        ((0xC0 | u >> 6) | (0x80 | u & 0x3F) << 8, 2)
+    } else {
+        ((0xE0 | u >> 12) | (0x80 | (u >> 6) & 0x3F) << 8 | (0x80 | u & 0x3F) << 16, 3)
+    }
+}
+
+/// Encodes a non-surrogate block and returns the byte count.
+fn encode_block(units: &[u16; BLOCK], out: &mut [MaybeUninit<u8>]) -> usize {
+    // Byte lengths keep the encoding loop vectorizable.
+    let mut words = [0u32; BLOCK];
+    let mut lens = [0u8; BLOCK];
+    // Encoding separately from the stores lets LLVM vectorize the first loop.
+    for k in 0..BLOCK {
+        let (word, n) = encode_bmp(units[k]);
+        (words[k], lens[k]) = (word, n as u8);
+    }
+    let mut len = 0;
+    for k in 0..BLOCK {
+        // Four-byte stores overlap; the last may extend past the returned count.
+        out[len..len + 4].write_copy_of_slice(&words[k].to_le_bytes());
+        len += lens[k] as usize;
+    }
+    len
+}

library/alloc/src/wtf8/mod.rs

@@ -143,26 +143,13 @@ impl Wtf8Buf {
     /// will always return the original code units.
     pub fn from_wide(v: &[u16]) -> Wtf8Buf {
         let mut string = Wtf8Buf::with_capacity(v.len());
-        for item in char::decode_utf16(v.iter().cloned()) {
-            match item {
-                Ok(ch) => string.push_char(ch),
-                Err(surrogate) => {
-                    let surrogate = surrogate.unpaired_surrogate();
-                    // SAFETY: Surrogates are known to be in the code point range.
-                    let code_point = unsafe { CodePoint::from_u32_unchecked(surrogate as u32) };
-                    // The string will now contain an unpaired surrogate.
-                    string.is_known_utf8 = false;
-                    // SAFETY: `decode_utf16` reports only unpaired surrogates here,
-                    // so this code point cannot form a surrogate pair with the
-                    // preceding and succeeding contents. The existing buffer is
-                    // well-formed WTF-8, and appending this encoded surrogate
-                    // preserves that invariant.
-                    unsafe {
-                        string.push_code_point_unchecked(code_point);
-                    }
-                }
-            }
-        }
+        let mut is_known_utf8 = true;
+        let _ = crate::utf16_decode::decode_utf16_into(v, &mut string.bytes, |surrogate| {
+            // The string will now contain an unpaired surrogate.
+            is_known_utf8 = false;
+            Ok(surrogate)
+        });
+        string.is_known_utf8 = is_known_utf8;
         string
     }
 

library/alloc/src/wtf8/tests.rs

@@ -133,6 +133,73 @@ fn wtf8buf_from_wide() {
     assert!(!buf.is_known_utf8);
 }
 
+#[test]
+fn wtf8buf_from_wide_block_boundaries() {
+    fn check(input: &[u16]) {
+        let output: Vec<u16> = Wtf8Buf::from_wide(input).encode_wide().collect();
+        assert_eq!(output, input);
+    }
+
+    for len in 0..=80 {
+        let mut state = 0x2545_f491_u32;
+        let input: Vec<u16> = (0..len)
+            .map(|_| {
+                state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
+                state as u16
+            })
+            .collect();
+        check(&input);
+    }
+
+    for boundary in [16, 32, 48, 64] {
+        let mut input = vec![0x0061; boundary + 1];
+        input[boundary - 1] = 0xd83d;
+        input[boundary] = 0xde00;
+        check(&input);
+
+        input[boundary] = 0x0062;
+        check(&input);
+
+        input[boundary - 1] = 0x0061;
+        input[boundary] = 0xdc00;
+        check(&input);
+    }
+}
+
+#[test]
+fn wtf8buf_from_wide_dense_block_boundaries() {
+    fn check(input: &[u16]) {
+        let buf = Wtf8Buf::from_wide(input);
+        assert_eq!(buf.encode_wide().collect::<Vec<_>>(), input);
+        let valid_utf8 = char::decode_utf16(input.iter().copied()).all(|unit| unit.is_ok());
+        assert_eq!(buf.is_known_utf8, valid_utf8);
+    }
+
+    for len in [15, 16, 17, 31, 32, 33, 47, 48, 49, 63, 64, 65, 80] {
+        for unit in [0x0080, 0x07ff, 0x0800, 0xd7ff, 0xe000, 0xffff] {
+            check(&vec![unit; len]);
+        }
+    }
+
+    for boundary in [16, 32, 48, 64] {
+        for tail in [1, 2, 16, 17] {
+            for (lead, trail) in [(0xd800, 0xdc00), (0xdbff, 0xdfff)] {
+                let mut input = vec![0xffff; boundary + tail];
+                input[boundary - 1] = lead;
+                input[boundary] = trail;
+                check(&input);
+
+                input[boundary] = 0xffff;
+                check(&input);
+
+                input[boundary - 1] = 0xffff;
+                input[boundary] = trail;
+                check(&input);
+            }
+        }
+    }
+}
+
 #[test]
 fn wtf8buf_push_str() {
     let mut string = Wtf8Buf::new();

library/alloctests/benches/lib.rs

@@ -7,8 +7,10 @@
 #![feature(slice_partition_dedup)]
 #![feature(strict_provenance_lints)]
 #![feature(test)]
+#![feature(wtf8_internals)]
 #![deny(implicit_provenance_casts)]
 
+extern crate alloc;
 extern crate test;
 
 mod binary_heap;