An equivalent of `bytes()` using internal iteration

custom_try_fold.rs

#![feature(try_trait_v2)]
use std::{io::{Write, BufRead}, ops::{Try, ControlFlow}};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let test_file_path = std::env::args_os().nth(1).expect("No input file specified!");

    let test_file = std::fs::File::open(test_file_path)?;
    let stdout = std::io::stdout();
    // Lock the stdout once so we don't have to acquire the lock for every line written, which is slow
    let stdout = stdout.lock();
    // And wrap it in a BufWriter, because we're going to make lots of small writes
    // and without this it would be very slow
    let mut writer = std::io::BufWriter::new(stdout);

    let reader = std::io::BufReader::new(test_file);

    let result = custom_try_fold(reader, 0, |acc: u64, i| ControlFlow::Continue(acc + i as u64));
    match result.branch() {
        ControlFlow::Continue(sum) => writeln!(writer, "{sum}")?,
        ControlFlow::Break(_) => todo!(),
    };

    Ok(())
}

fn custom_try_fold<B, F, I>(mut reader: I, init: B, mut f: F) -> impl Try<Output = B>
where
    I: BufRead,
    F: FnMut(B, u8) -> ControlFlow<B, B>,
{
    let mut accum = init;
    loop {
        match reader.fill_buf() {
            Ok(chunk) => {
                // check for EOF
                if chunk.len() == 0 {
                    return Ok(accum);
                }
                let mut iterator = chunk.iter().copied();
                let result = iterator.try_fold(accum, &mut f);
                let consumed = chunk.len() - iterator.len();
                reader.consume(consumed);
                match result {
                    ControlFlow::Continue(a) => accum = a,
                    ControlFlow::Break(a) => {
                        accum = a;
                        return Ok(accum);
                    },
                }
            },
            Err(e) => return Err(e),
        };
    }
}