Compiler Explorer

Source code

use std::iter::Peekable;
use std::borrow::Cow;

type Window = [u8; 3];
type Iter<'a> = Peekable<std::slice::Iter<'a, u8>>;

pub struct WindowedReplacer<'a> {
    window: Window,
    input_iter: Iter<'a>,
}

#[inline]
fn is_replaceable(byte: u8) -> bool {
    matches!(byte, b'a'..=b'z')
}

#[inline]
fn next_byte(iter: &mut Iter) -> Option<u8> {
    let byte = iter.next().copied()?;
    if is_replaceable(byte) {
        while iter.next_if(|b| is_replaceable(**b)).is_some() {}
        Some(b'#')
    } else {
        Some(byte)
    }
}

impl<'a> WindowedReplacer<'a> {
    pub fn new(input: &'a [u8]) -> Option<Self> {
        let mut window: Window = Default::default();
        let mut iter = input.iter().peekable();
        for byte in window.iter_mut().skip(1) {
            *byte = next_byte(&mut iter)?;
        }
        Some(WindowedReplacer {
            window,
            input_iter: iter,
        })
    }
}

impl<'a> Iterator for WindowedReplacer<'a> {
    type Item = Window;

#[inline]
    fn next(&mut self) -> Option<Self::Item> {
        for i in 0..2 {
            self.window[i] = self.window[i + 1];
        }
        let byte = next_byte(&mut self.input_iter)?;
        self.window[2] = byte;
        Some(self.window)
    }
}

pub fn tensor_populate_ngrams(tensor: &mut [f32], input: Cow<'_, [u8]>) {
    if let Some(replacer) = WindowedReplacer::new(&input) {                
        for ngram in replacer {
            let idx = ((ngram[0] as usize) << 16) | ((ngram[1] as usize) << 8) | (ngram[2] as usize);
            tensor[idx] += 1.0;
        }
    }
}