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(
117    reader: &mut dyn ReadSeek,
118    offset: u64,
119    buffer: &mut [u8],
120) -> Result<usize> {
121    let position = reader.stream_position()?;
122    let read_result = match reader.seek(SeekFrom::Start(offset)) {
123        Ok(_) => read_exact_or_eof(reader, buffer),
124        Err(error) => Err(error),
125    };
126    let restore_result = reader.seek(SeekFrom::Start(position));
127    match (read_result, restore_result) {
128        (Ok(count), Ok(_)) => Ok(count),
129        (Err(error), Ok(_)) => Err(error),
130        (_, Err(error)) => Err(error),
131    }
132}
133
134/// Reads from `reader` until `buffer` is full or EOF is reached.
135///
136/// # Parameters
137/// - `reader`: Source reader.
138/// - `buffer`: Destination buffer to fill.
139///
140/// # Returns
141/// The number of bytes written into `buffer`.
142///
143/// # Errors
144/// Returns the first non-interrupted read error reported by `reader`.
145fn read_exact_or_eof(reader: &mut dyn ReadSeek, buffer: &mut [u8]) -> Result<usize> {
146    let mut total = 0;
147    while total < buffer.len() {
148        match reader.read(&mut buffer[total..]) {
149            Ok(0) => break,
150            Ok(count) => total += count,
151            Err(error) => {
152                if error.kind() == ErrorKind::Interrupted {
153                    continue;
154                }
155                return Err(error);
156            }
157        }
158    }
159    Ok(total)
160}