Skip to main content

qubit_io/ext/
read_seek_ext.rs

1/*******************************************************************************
2 *
3 *    Copyright (c) 2026 Haixing Hu.
4 *
5 *    SPDX-License-Identifier: Apache-2.0
6 *
7 *    Licensed under the Apache License, Version 2.0.
8 *
9 ******************************************************************************/
10use std::io::{
11    ErrorKind,
12    Read,
13    Result,
14    Seek,
15    SeekFrom,
16};
17
18use crate::ReadSeek;
19
20/// Extension methods for values that implement both [`Read`] and [`Seek`].
21///
22/// `ReadSeekExt` provides position-preserving read helpers for common
23/// inspection use cases such as file signature checks, MIME detection, and
24/// random-offset probing.
25pub trait ReadSeekExt: Read + Seek {
26    /// Reads from the current position and restores the original position.
27    ///
28    /// This method has the same partial-EOF semantics as
29    /// [`crate::ReadExt::read_exact_or_eof`], but it leaves the stream positioned
30    /// where it was before the call when restoration succeeds.
31    ///
32    /// # Parameters
33    /// - `buffer`: Destination buffer to fill.
34    ///
35    /// # Returns
36    /// The number of bytes written into `buffer`.
37    ///
38    /// # Errors
39    /// Returns an error when reading the current position, reading bytes, or
40    /// restoring the original position fails. If both reading and restoration
41    /// fail, the restoration error is returned because the caller's stream
42    /// position contract was not preserved.
43    fn peek_exact_or_eof(&mut self, buffer: &mut [u8]) -> Result<usize>;
44
45    /// Reads from `offset` and restores the original position.
46    ///
47    /// This method seeks to `offset`, reads until `buffer` is full or EOF is
48    /// reached, and then restores the position that was current before the
49    /// call.
50    ///
51    /// # Parameters
52    /// - `offset`: Absolute byte offset from the start of the stream.
53    /// - `buffer`: Destination buffer to fill.
54    ///
55    /// # Returns
56    /// The number of bytes written into `buffer`.
57    ///
58    /// # Errors
59    /// Returns an error when reading the current position, seeking to `offset`,
60    /// reading bytes, or restoring the original position fails. If restoration
61    /// fails, the restoration error is returned.
62    fn read_exact_or_eof_at(&mut self, offset: u64, buffer: &mut [u8]) -> Result<usize>;
63}
64
65impl<T> ReadSeekExt for T
66where
67    T: Read + Seek + ?Sized,
68{
69    #[inline]
70    fn peek_exact_or_eof(&mut self, buffer: &mut [u8]) -> Result<usize> {
71        let mut reader = self;
72        peek_exact_or_eof_impl(&mut reader, buffer)
73    }
74
75    #[inline]
76    fn read_exact_or_eof_at(&mut self, offset: u64, buffer: &mut [u8]) -> Result<usize> {
77        let mut reader = self;
78        read_exact_or_eof_at_impl(&mut reader, offset, buffer)
79    }
80}
81
82/// Reads from the current stream position and restores that position.
83///
84/// # Parameters
85/// - `reader`: Seekable reader to inspect.
86/// - `buffer`: Destination buffer to fill.
87///
88/// # Returns
89/// The number of bytes written into `buffer`.
90///
91/// # Errors
92/// Returns an error when position lookup, reading, or position restoration fails.
93fn peek_exact_or_eof_impl(reader: &mut dyn ReadSeek, buffer: &mut [u8]) -> Result<usize> {
94    let position = reader.stream_position()?;
95    let read_result = read_exact_or_eof(reader, buffer);
96    let restore_result = reader.seek(SeekFrom::Start(position));
97    match (read_result, restore_result) {
98        (Ok(count), Ok(_)) => Ok(count),
99        (Err(error), Ok(_)) => Err(error),
100        (_, Err(error)) => Err(error),
101    }
102}
103
104/// Reads from `offset` and restores the original stream position.
105///
106/// # Parameters
107/// - `reader`: Seekable reader to inspect.
108/// - `offset`: Absolute byte offset from the start of the stream.
109/// - `buffer`: Destination buffer to fill.
110///
111/// # Returns
112/// The number of bytes written into `buffer`.
113///
114/// # Errors
115/// Returns an error when position lookup, seeking, reading, or position restoration fails.
116fn read_exact_or_eof_at_impl(reader: &mut dyn ReadSeek, offset: u64, buffer: &mut [u8]) -> Result<usize> {
117    let position = reader.stream_position()?;
118    let read_result = match reader.seek(SeekFrom::Start(offset)) {
119        Ok(_) => read_exact_or_eof(reader, buffer),
120        Err(error) => Err(error),
121    };
122    let restore_result = reader.seek(SeekFrom::Start(position));
123    match (read_result, restore_result) {
124        (Ok(count), Ok(_)) => Ok(count),
125        (Err(error), Ok(_)) => Err(error),
126        (_, Err(error)) => Err(error),
127    }
128}
129
130/// Reads from `reader` until `buffer` is full or EOF is reached.
131///
132/// # Parameters
133/// - `reader`: Source reader.
134/// - `buffer`: Destination buffer to fill.
135///
136/// # Returns
137/// The number of bytes written into `buffer`.
138///
139/// # Errors
140/// Returns the first non-interrupted read error reported by `reader`.
141fn read_exact_or_eof(reader: &mut dyn ReadSeek, buffer: &mut [u8]) -> Result<usize> {
142    let mut total = 0;
143    while total < buffer.len() {
144        match reader.read(&mut buffer[total..]) {
145            Ok(0) => break,
146            Ok(count) => total += count,
147            Err(error) => {
148                if error.kind() == ErrorKind::Interrupted {
149                    continue;
150                }
151                return Err(error);
152            }
153        }
154    }
155    Ok(total)
156}