Quelle  read_cache.rs   Sprache: unbekannt

use alloc::boxed::Box;
use alloc::vec::Vec;
use core::cell::RefCell;
use core::convert::TryInto;
use core::mem;
use core::ops::Range;
#[cfg(feature = "std")]
use std::io::{Read, Seek, SeekFrom};

#[cfg(not(feature = "std"))]
use alloc::collections::btree_map::{BTreeMap as Map, Entry};
#[cfg(feature = "std")]
use std::collections::hash_map::{Entry, HashMap as Map};

use crate::read::ReadRef;

/// An implementation of [`ReadRef`] for data in a stream that implements
/// `Read + Seek`.
///
/// Contains a cache of read-only blocks of data, allowing references to
/// them to be returned. Entries in the cache are never removed.
/// Entries are keyed on the offset and size of the read.
/// Currently overlapping reads are considered separate reads.
///
/// This is primarily intended for environments where memory mapped files
/// are not available or not suitable, such as WebAssembly.
///
/// Note that malformed files can cause the cache to grow much larger than
/// the file size.
#[derive(Debug)]
pub struct ReadCache<R: ReadCacheOps> {
    cache: RefCell<ReadCacheInternal<R>>,
}

#[derive(Debug)]
struct ReadCacheInternal<R: ReadCacheOps> {
    read: R,
    bufs: Map<(u64, u64), Box<[u8]>>,
    strings: Map<(u64, u8), Box<[u8]>>,
    len: Option<u64>,
}

impl<R: ReadCacheOps> ReadCacheInternal<R> {
    /// Ensures this range is contained in the len of the file
    fn range_in_bounds(&mut self, range: &Range<u64>) -> Result<(), ()> {
        if range.start <= range.end && range.end <= self.len()? {
            Ok(())
        } else {
            Err(())
        }
    }

    /// The length of the underlying read, memoized
    fn len(&mut self) -> Result<u64, ()> {
        match self.len {
            Some(len) => Ok(len),
            None => {
                let len = self.read.len()?;
                self.len = Some(len);
                Ok(len)
            }
        }
    }
}

impl<R: ReadCacheOps> ReadCache<R> {
    /// Create an empty `ReadCache` for the given stream.
    pub fn new(read: R) -> Self {
        ReadCache {
            cache: RefCell::new(ReadCacheInternal {
                read,
                bufs: Map::new(),
                strings: Map::new(),
                len: None,
            }),
        }
    }

    /// Return an implementation of `ReadRef` that restricts reads
    /// to the given range of the stream.
    pub fn range(&self, offset: u64, size: u64) -> ReadCacheRange<'_, R> {
        ReadCacheRange {
            r: self,
            offset,
            size,
        }
    }

    /// Free buffers used by the cache.
    pub fn clear(&mut self) {
        self.cache.borrow_mut().bufs.clear();
    }

    /// Unwrap this `ReadCache<R>`, returning the underlying reader.
    pub fn into_inner(self) -> R {
        self.cache.into_inner().read
    }
}

impl<'a, R: ReadCacheOps> ReadRef<'a> for &'a ReadCache<R> {
    fn len(self) -> Result<u64, ()> {
        self.cache.borrow_mut().len()
    }

    fn read_bytes_at(self, offset: u64, size: u64) -> Result<&'a [u8], ()> {
        if size == 0 {
            return Ok(&[]);
        }
        let cache = &mut *self.cache.borrow_mut();
        cache.range_in_bounds(&(offset..(offset.saturating_add(size))))?;
        let buf = match cache.bufs.entry((offset, size)) {
            Entry::Occupied(entry) => entry.into_mut(),
            Entry::Vacant(entry) => {
                let size = size.try_into().map_err(|_| ())?;
                cache.read.seek(offset)?;
                let mut bytes = Vec::new();
                bytes.try_reserve_exact(size).map_err(|_| ())?;
                bytes.resize(size, 0);
                let mut bytes = bytes.into_boxed_slice();
                cache.read.read_exact(&mut bytes)?;
                entry.insert(bytes)
            }
        };
        // Extend the lifetime to that of self.
        // This is OK because we never mutate or remove entries.
        Ok(unsafe { mem::transmute::<&[u8], &[u8]>(buf) })
    }

    fn read_bytes_at_until(self, range: Range<u64>, delimiter: u8) -> Result<&'a [u8], ()> {
        let cache = &mut *self.cache.borrow_mut();
        cache.range_in_bounds(&range)?;
        let buf = match cache.strings.entry((range.start, delimiter)) {
            Entry::Occupied(entry) => entry.into_mut(),
            Entry::Vacant(entry) => {
                cache.read.seek(range.start)?;

                let max_check: usize = (range.end - range.start).try_into().map_err(|_| ())?;
                // Strings should be relatively small.
                // TODO: make this configurable?
                let max_check = max_check.min(4096);

                let mut bytes = Vec::new();
                let mut checked = 0;
                loop {
                    bytes.resize((checked + 256).min(max_check), 0);
                    let read = cache.read.read(&mut bytes[checked..])?;
                    if read == 0 {
                        return Err(());
                    }
                    if let Some(len) = memchr::memchr(delimiter, &bytes[checked..][..read]) {
                        bytes.truncate(checked + len);
                        break entry.insert(bytes.into_boxed_slice());
                    }
                    checked += read;
                    if checked >= max_check {
                        return Err(());
                    }
                }
            }
        };
        // Extend the lifetime to that of self.
        // This is OK because we never mutate or remove entries.
        Ok(unsafe { mem::transmute::<&[u8], &[u8]>(buf) })
    }
}

/// An implementation of [`ReadRef`] for a range of data in a stream that
/// implements `Read + Seek`.
///
/// Shares an underlying [`ReadCache`] with a lifetime of `'a`.
#[derive(Debug)]
pub struct ReadCacheRange<'a, R: ReadCacheOps> {
    r: &'a ReadCache<R>,
    offset: u64,
    size: u64,
}

impl<'a, R: ReadCacheOps> Clone for ReadCacheRange<'a, R> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<'a, R: ReadCacheOps> Copy for ReadCacheRange<'a, R> {}

impl<'a, R: ReadCacheOps> ReadRef<'a> for ReadCacheRange<'a, R> {
    fn len(self) -> Result<u64, ()> {
        Ok(self.size)
    }

    fn read_bytes_at(self, offset: u64, size: u64) -> Result<&'a [u8], ()> {
        if size == 0 {
            return Ok(&[]);
        }
        let end = offset.checked_add(size).ok_or(())?;
        if end > self.size {
            return Err(());
        }
        let r_offset = self.offset.checked_add(offset).ok_or(())?;
        self.r.read_bytes_at(r_offset, size)
    }

    fn read_bytes_at_until(self, range: Range<u64>, delimiter: u8) -> Result<&'a [u8], ()> {
        let r_start = self.offset.checked_add(range.start).ok_or(())?;
        let r_end = self.offset.checked_add(range.end).ok_or(())?;
        let bytes = self.r.read_bytes_at_until(r_start..r_end, delimiter)?;
        let size = bytes.len().try_into().map_err(|_| ())?;
        let end = range.start.checked_add(size).ok_or(())?;
        if end > self.size {
            return Err(());
        }
        Ok(bytes)
    }
}

/// Operations required to implement [`ReadCache`].
///
/// This is a subset of the `Read` and `Seek` traits.
/// A blanket implementation is provided for all types that implement
/// `Read + Seek`.
#[allow(clippy::len_without_is_empty)]
pub trait ReadCacheOps {
    /// Return the length of the stream.
    ///
    /// Equivalent to `std::io::Seek::seek(SeekFrom::End(0))`.
    fn len(&mut self) -> Result<u64, ()>;

    /// Seek to the given position in the stream.
    ///
    /// Equivalent to `std::io::Seek::seek` with `SeekFrom::Start(pos)`.
    fn seek(&mut self, pos: u64) -> Result<u64, ()>;

    /// Read up to `buf.len()` bytes into `buf`.
    ///
    /// Equivalent to `std::io::Read::read`.
    fn read(&mut self, buf: &mut [u8]) -> Result<usize, ()>;

    /// Read exactly `buf.len()` bytes into `buf`.
    ///
    /// Equivalent to `std::io::Read::read_exact`.
    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), ()>;
}

#[cfg(feature = "std")]
impl<T: Read + Seek> ReadCacheOps for T {
    fn len(&mut self) -> Result<u64, ()> {
        self.seek(SeekFrom::End(0)).map_err(|_| ())
    }

    fn seek(&mut self, pos: u64) -> Result<u64, ()> {
        self.seek(SeekFrom::Start(pos)).map_err(|_| ())
    }

    fn read(&mut self, buf: &mut [u8]) -> Result<usize, ()> {
        Read::read(self, buf).map_err(|_| ())
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), ()> {
        Read::read_exact(self, buf).map_err(|_| ())
    }
}