<commit_before>
use std::io::{BufRead,ErrorKind,Result,Write};

/// This trait extends any type that implements BufRead with a
/// stream_until_token() function.
pub trait BufReadPlus: BufRead {
    /// Stream all bytes to out until the `token` delimiter is reached.
    ///
    /// This function will continue to read (and stream) bytes from the
    /// underlying stream until the token or EOF is found. Once found, all
    /// bytes up to the token (if found) will have been streamed to `out` and
    /// this input stream will advance past the token (the token will be
    /// discarded).
    ///
    /// This function will return `Ok(n)` where `n` is the number of bytes
    /// which were read, counting the token if it was found.  If the token was
    /// not found, it will still return `Ok(n)`
    ///
    /// # Errors
    ///
    /// This function will ignore all instances of `ErrorKind::Interrupted` and
    /// will otherwise return any errors returned by `fill_buf`.
    fn stream_until_token<W: Write>(&mut self, token: &[u8], out: &mut W) -> Result<usize> {
        stream_until_token(self, token, out)
    }
}

// Implement BufReadPlus for everything that implements BufRead
impl<T: BufRead> BufReadPlus for T {}

fn stream_until_token<R: BufRead + ?Sized, W: Write>(r: &mut R, token: &[u8], mut out: &mut W)
                                                     -> Result<usize> {
    let mut read = 0;

    // Partial represents the size of a token prefix found at the end of a buffer,
    // usually 0.  If not zero, the beginning of the next buffer is checked for the
    // suffix to find matches that straddle two buffers
    let mut partial: usize = 0;

    loop {
        let (found,used) = {
            let available = match r.fill_buf() {
                Ok(n) => n,
                Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
                Err(e) => return Err(e)
            };
            // If last buffer ended in a token prefix, check if this one starts with the
            // matching suffix
            if partial > 0 && available[..token.len() - partial] == token[partial..] {
                (true, token.len() - partial)
            }
            else {
                if partial > 0 {
                    // Last buffer ended in a token prefix, but it didn't pan out, so
                    // we need to push it along
                    try!( out.write_all(&token[..partial]) );
                }
                // Search for the token
                match available
                    .windows(token.len())
                    .enumerate()
                    .filter(|&(_,w)| { w == token })
                    .map(|(i,_)| { i })
                    .next()
                {
                    Some(i) => {
                        try!( out.write_all(&available[..i]) );
                        (true, i + token.len())
                    },
                    None => {
                        // Check for partial matches at the end of the buffer
                        partial = if available.len() > token.len() - 1 {
                            match (1..(token.len()-1))
                                .rev()
                                .filter(|&width| {
                                    token[..width] == available[available.len() - width..]
                                })
                                .next()
                            {
                                Some(width) => width,
                                None => 0
                            }
                        } else { 0 };
                        // Push all except the partial token at the end (if any)
                        try!( out.write_all(&available[..available.len()-partial]) );
                        (false, available.len()) // But mark it all consumed
                    }
                }
            }
        };
        r.consume(used);
        read += used;
        if found || used == 0 {
            return Ok(read);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::{Cursor,BufReader};
    use super::BufReadPlus;<commit_msg>Fix stream_until_token() partial token matches with buffer shorter than token<commit_after>
use std::io::{BufRead,ErrorKind,Result,Write};

/// This trait extends any type that implements BufRead with a
/// stream_until_token() function.
pub trait BufReadPlus: BufRead {
    /// Stream all bytes to out until the `token` delimiter is reached.
    ///
    /// This function will continue to read (and stream) bytes from the
    /// underlying stream until the token or EOF is found. Once found, all
    /// bytes up to the token (if found) will have been streamed to `out` and
    /// this input stream will advance past the token (the token will be
    /// discarded).
    ///
    /// This function will return `Ok(n)` where `n` is the number of bytes
    /// which were read, counting the token if it was found.  If the token was
    /// not found, it will still return `Ok(n)`
    ///
    /// # Errors
    ///
    /// This function will ignore all instances of `ErrorKind::Interrupted` and
    /// will otherwise return any errors returned by `fill_buf`.
    fn stream_until_token<W: Write>(&mut self, token: &[u8], out: &mut W) -> Result<usize> {
        stream_until_token(self, token, out)
    }
}

// Implement BufReadPlus for everything that implements BufRead
impl<T: BufRead> BufReadPlus for T {}

fn stream_until_token<R: BufRead + ?Sized, W: Write>(r: &mut R, token: &[u8], mut out: &mut W)
                                                     -> Result<usize> {
    let mut read = 0;

    // Partial represents the size of a token prefix found at the end of a buffer,
    // usually 0.  If not zero, the beginning of the next buffer is checked for the
    // suffix to find matches that straddle two buffers
    let mut partial: usize = 0;

    loop {
        let (found,used) = {
            let available = match r.fill_buf() {
                Ok(n) => n,
                Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
                Err(e) => return Err(e)
            };
            // If last buffer ended in a token prefix, check if this one starts with the
            // matching suffix
            if partial > 0 && available[..token.len() - partial] == token[partial..] {
                (true, token.len() - partial)
            }
            else {
                if partial > 0 {
                    // Last buffer ended in a token prefix, but it didn't pan out, so
                    // we need to push it along
                    try!( out.write_all(&token[..partial]) );
                }
                // Search for the token
                match available
                    .windows(token.len())
                    .enumerate()
                    .filter(|&(_,w)| { w == token })
                    .map(|(i,_)| { i })
                    .next()
                {
                    Some(i) => {
                        try!( out.write_all(&available[..i]) );
                        (true, i + token.len())
                    },
                    None => {
                        // Check for partial matches at the end of the buffer
                        let mut window = token.len() - 1;
                        if available.len() < window { window = available.len(); }

                        partial = match (1..window+1)
                            .rev()
                            .filter(|&width| {
                                token[..width] == available[available.len() - width..]
                            })
                            .next()
                        {
                            Some(width) => width,
                            None => 0
                        };

                        // Push all except the partial token at the end (if any)
                        try!( out.write_all(&available[..available.len()-partial]) );
                        (false, available.len()) // But mark it all consumed
                    }
                }
            }
        };
        r.consume(used);
        read += used;
        if found || used == 0 {
            return Ok(read);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::{Cursor,BufReader};
    use super::BufReadPlus;
{'fraction_non_alphanumeric': 0.11184447377895115, 'fraction_numerical': 0.003000120004800192, 'mean_word_length': 1.2714636140637776, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6830253', 'n_tokens_mistral': 2166, 'n_tokens_neox': 2064, 'n_words': 758}