Skip to main content

qubit_io/ext/
read_seek_ext.rs

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