Compiler Explorer

Source code

#![feature(pattern)]

use core::mem::{transmute, MaybeUninit};
use core::str::pattern::Pattern;

pub fn is_prefix_of(c: char, haystack: &str) -> bool {
    encode_char_4bytes(c, &mut [MaybeUninit::uninit(); 4]).is_prefix_of(haystack)
}

#[inline]
fn encode_utf8(ch: char, dst: &mut [u8]) -> &mut str {
    // SAFETY: `char` is not a surrogate, so this is valid UTF-8.
    unsafe { from_utf8_unchecked_mut(encode_utf8_raw(ch as u32, dst)) }
}

#[inline]
unsafe fn from_utf8_unchecked_mut(v: &mut [u8]) -> &mut str {
    // SAFETY: the caller must guarantee that the bytes `v`
    // are valid UTF-8, thus the cast to `*mut str` is safe.
    // Also, the pointer dereference is safe because that pointer
    // comes from a reference which is guaranteed to be valid for writes.
    unsafe { &mut *(v as *mut [u8] as *mut str) }
}

#[inline]
fn encode_utf8_raw(code: u32, dst: &mut [u8]) -> &mut [u8] {
    let len = len_utf8(code);
    if dst.len() < len {
        panic!(
            "encode_utf8: need {} bytes to encode U+{:X}, but the buffer has {}",
            len,
            code,
            dst.len(),
        );
    }
    unsafe { encode_utf8_raw_unchecked(code, transmute(dst)) }
}

#[inline]
unsafe fn encode_utf8_raw_unchecked(code: u32, dst: &mut [MaybeUninit<u8>]) -> &mut [u8] {
    let len = len_utf8(code);
    match len {
        1 => {
            dst.get_unchecked_mut(0).write(code as u8);
        }
        2 => {
            dst.get_unchecked_mut(0)
                .write((code >> 6 & 0x1F) as u8 | TAG_TWO_B);
            dst.get_unchecked_mut(1)
                .write((code & 0x3F) as u8 | TAG_CONT);
        }
        3 => {
            dst.get_unchecked_mut(0)
                .write((code >> 12 & 0x0F) as u8 | TAG_THREE_B);
            dst.get_unchecked_mut(1)
                .write((code >> 6 & 0x3F) as u8 | TAG_CONT);
            dst.get_unchecked_mut(2)
                .write((code & 0x3F) as u8 | TAG_CONT);
        }
        4 => {
            dst.get_unchecked_mut(0)
                .write((code >> 18 & 0x07) as u8 | TAG_FOUR_B);
            dst.get_unchecked_mut(1)
                .write((code >> 12 & 0x3F) as u8 | TAG_CONT);
            dst.get_unchecked_mut(2)
                .write((code >> 6 & 0x3F) as u8 | TAG_CONT);
            dst.get_unchecked_mut(3)
                .write((code & 0x3F) as u8 | TAG_CONT);
        }
        _ => unreachable!(),
    }
    transmute(dst.get_unchecked_mut(..len))
}

#[inline]
fn encode_char_4bytes(ch: char, dst: &mut [MaybeUninit<u8>]) -> &mut str {
    let encoded_char = unsafe { encode_utf8_raw_unchecked(ch as u32, dst) };
    unsafe { core::str::from_utf8_unchecked_mut(encoded_char) }
}

const TAG_CONT: u8 = 0b1000_0000;
const TAG_TWO_B: u8 = 0b1100_0000;
const TAG_THREE_B: u8 = 0b1110_0000;
const TAG_FOUR_B: u8 = 0b1111_0000;
const MAX_ONE_B: u32 = 0x80;
const MAX_TWO_B: u32 = 0x800;
const MAX_THREE_B: u32 = 0x10000;

#[inline]
const fn len_utf8(code: u32) -> usize {
    if code < MAX_ONE_B {
        1
    } else if code < MAX_TWO_B {
        2
    } else if code < MAX_THREE_B {
        3
    } else {
        4
    }
}