Creating a file streaming abstraction

user_292731 0
Viewed 80

This is probably a super basic aspect of Rust, but I'm struggling with seeing through the levels of reference checks to make this make sense.

Imagine my current program sums all of the binary-encoded integers in a file:

pub fn read_int(file: &mut File) -> io::Result<i64> {
  let mut int_buffer = [0; 8];
  file.read_exact(&mut int_buffer)?;
  Ok(i64::from_be_bytes(int_buffer))
}

fn read_loop() -> io::Result<f64> {
    let mut file = File::open("data.dat")?;
    let mut found_end = false;
    let mut dat_sum = 0;
    let mut dat_count = 0;
    while !found_end {
        match read_int(&mut file) {
            Ok(d) => {
                dat_sum   += d;
                dat_count += 1;
            },
            Err(_) => found_end = true
        };
    }
    Ok(dat_sum as f64 / dat_count as f64)
}

I'd like to abstract over this to instead use a specialized i64 iterator:

fn read_loop() -> io::Result<f64> {
    let mut int_iter = IntIterator::open("data.dat")?;
    let mut dat_sum = 0;
    let mut dat_count = 0;
    for d in iiter {
        dat_sum   += d;
        dat_count += 1;
    }
    Ok(dat_sum as f64 / dat_count as f64)
}

I tried implementing this, but:

pub struct TickIterator<'r> {
    file: &'r mut File,
    found_end: bool
}

impl<'r> TickIterator<'r> {
    pub fn new(fp: String) -> io::Result<TickIterator<'r>> {
      let mut file = File::open(fp)?;
      Ok(TickIterator {
          file: &mut file,
          found_end: false
      })
    }
}

impl<'a> Iterator for TickIterator<'a> {
    type Item = i64;

    fn next(&mut self) -> Option<Self::Item> {
        if self.found_end {
            None
        } else {
            match read_int(self.file) {
                Err(_) => {
                    self.found_end = true;
                    None
                },
                Ok(d) => Some(d)
            }
        }
    }
}

It looks like I have a lot of problems to work out; the first error seems to be that I can't return the reference to file inside the struct because it is owned by the function new itself.

Is this the right approach, or am I going about things wrong here?

borrow-checker rust
1 Answers

Here is one approach for how you could do this. The trick here is to just ignore the lifetime entirely by leaving the type up to the generic. We don't really need to know what we are reading the numbers from since all we need is std::io::Read to be implemented for the type for this to work. With this approach, you can File, TcpStream, BufReader, and many other types from different crates which also use Read. &mut T is a type just like T and if we look in the docs, Read is implemented for all &mut T where T: Read.

use std::io::{self, ErrorKind, Read};

pub struct ValueIter<R> {
    reader: R,
    found_end: bool,
}

impl<R: Read> ValueIter<R> {
    pub fn for_reader(reader: R) -> Self {
        ValueIter { 
            reader,
            found_end: false,
        }
    }
}

// `Read` is the trait that gives us `read_exact` 
impl<R: Read> Iterator for ValueIter<R> {
    type Item = io::Result<i64>;
    
    fn next(&mut self) -> Option<Self::Item> {
        if self.found_end {
            return None
        }
        
        let mut buffer = [0u8; 8];
        match self.reader.read_exact(&mut buffer) {
            // Read the next number
            Ok(_) => Some(Ok(i64::from_be_bytes(buffer))),
            // Handle end of file reached
            Err(e) if e.kind() == ErrorKind::UnexpectedEof => {
                self.found_end = true;
                None
            }
            // Return any other errors as part of the iterator
            Err(e) => Some(Err(e)),
        }
    }
}

And then it could be used like such.

pub fn read_loop<R: Read>(reader: &mut R) -> io::Result<i64> {
    let mut sum = 0;
    let mut count = 0;
    
    for value in ValueIter::for_reader(reader) {
        sum += value?;
        count += 1;
    }
    
    Ok(sum / count)
}

The neat part is that since we made it generic over Read, we can use more than just File. For example, you will improve your performance by wrapping it in a BufReader so it can get the file from the OS in larger chunks instead of needing to ask the OS every time a new number is read.

use std::fs::File;
use std::io::BufReader;
let mut file = BufReader::new(File::open("data.dat").expect("Unable to open file"));

println!("{:?}", read_loop(&mut file).unwrap());
user_59876690
Related