hurray_io/file/reader.rs
1use hurray_core::TensorDescriptor;
2use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt, SeekFrom};
3
4use crate::file::sans_io::{FileReaderCore, FileStep};
5// `pub use`, not `use`: 0.1.0 published these at `file::reader::*` as well as at
6// `file::*`, because this module is public and defined them. They belong to the core now,
7// and the published paths keep resolving. `hurray_io::file::*` is canonical.
8pub use crate::file::types::{FileComposite, FileItem, FileTensor, DEFAULT_MAX_COMPOSITE_DEPTH};
9use crate::file::types::{IndexEntry, KvValue};
10use crate::{Error, Result};
11
12/// Reads tensors from a seekable Hurray file.
13///
14/// Open with [`FileReader::open`], then look up tensors by name with
15/// [`read_tensor`][FileReader::read_tensor] or inspect metadata with
16/// [`tensor_names`][FileReader::tensor_names] and [`kv`][FileReader::kv].
17///
18/// This is a driver, not a parser: [`FileReaderCore`] holds the container format and
19/// decides which bytes it needs, and this type only seeks to the range the core asks for
20/// and reports how many bytes arrived (ADR-041 § 7).
21///
22/// # Examples
23///
24/// ```no_run
25/// # #[tokio::main]
26/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
27/// use hurray_io::file::FileReader;
28///
29/// let file = tokio::fs::File::open("model.hrry").await?;
30/// let mut reader = FileReader::open(file).await?;
31/// println!("tensors: {:?}", reader.tensor_names().collect::<Vec<_>>());
32/// let tensor = reader.read_tensor("embeddings").await?;
33/// println!("buffer[0]: {} bytes", tensor.buffers[0].len());
34/// # Ok(())
35/// # }
36/// ```
37#[derive(Debug)]
38pub struct FileReader<R> {
39 inner: R,
40 core: FileReaderCore,
41 /// Footer index: one entry per tensor, in file-write order (or sorted if
42 /// `SORTED_INDEX` was set by the writer).
43 pub index: Vec<IndexEntry>,
44 /// File-level KV metadata (empty if the file has no KV section).
45 pub kv: Vec<(String, KvValue)>,
46}
47
48impl<R: AsyncRead + AsyncSeek + Unpin> FileReader<R> {
49 /// Opens a Hurray file: reads the header, trailer, index, and KV section.
50 ///
51 /// # Errors
52 ///
53 /// - [`Error::InvalidFileMagic`] — first 8 bytes are not `HRRYFILE`
54 /// - [`Error::UnsupportedContainerVersion`] — major version > 1
55 /// - [`Error::ReservedFileFlagBits`] — header `file_flags` has unknown bits set
56 /// - [`Error::InvalidTrailerMagic`] — trailer `trailer_magic` is not `HRRY`
57 /// - [`Error::IndexOverrunsTrailer`] — index section overlaps the trailer
58 /// - [`Error::IndexCrc32cMismatch`] — CRC-32C verification failed
59 /// - [`Error::DuplicateTensorName`] — index contains duplicate names
60 /// - [`Error::Io`] / [`Error::UnexpectedEof`] — I/O failures
61 pub async fn open(mut inner: R) -> Result<Self> {
62 let size = inner.seek(SeekFrom::End(0)).await?;
63 let mut core = FileReaderCore::new(size);
64
65 loop {
66 match core.step()? {
67 FileStep::Range { offset, len } => fill(&mut inner, &mut core, offset, len).await?,
68 FileStep::Ready => break,
69 // Until the file is open the core answers nothing else.
70 _ => return Err(Error::CoreFailed),
71 }
72 }
73
74 // Copies, because 0.1.0 published `index` and `kv` as fields. The core keeps its
75 // own: it needs them for every name lookup and for the composite walk.
76 Ok(Self {
77 index: core.index().to_vec(),
78 kv: core.kv().to_vec(),
79 inner,
80 core,
81 })
82 }
83
84 /// Sets the maximum composite nesting depth accepted by
85 /// [`read_composite`][FileReader::read_composite]. Default:
86 /// [`DEFAULT_MAX_COMPOSITE_DEPTH`].
87 pub fn with_max_composite_depth(mut self, max_composite_depth: usize) -> Self {
88 self.core = self.core.with_max_composite_depth(max_composite_depth);
89 self
90 }
91
92 /// Returns the names of all tensors in the file, in index order.
93 pub fn tensor_names(&self) -> impl Iterator<Item = &str> {
94 self.index.iter().map(|e| e.name.as_str())
95 }
96
97 /// Returns the file-level KV metadata.
98 pub fn kv(&self) -> &[(String, KvValue)] {
99 &self.kv
100 }
101
102 /// Decodes and returns the descriptor for `name` without reading buffer data.
103 ///
104 /// Useful for inspecting dtype, shape, or layout before deciding whether to
105 /// load the full tensor.
106 pub async fn read_descriptor(&mut self, name: &str) -> Result<TensorDescriptor> {
107 self.core.begin_descriptor(name)?;
108 match self.drive().await? {
109 FileStep::Descriptor => self.core.take_descriptor().ok_or(Error::CoreFailed),
110 _ => Err(Error::CoreFailed),
111 }
112 }
113
114 /// Reads and returns a complete tensor (descriptor + all buffers).
115 pub async fn read_tensor(&mut self, name: &str) -> Result<FileTensor> {
116 self.core.begin_tensor(name)?;
117 match self.drive().await? {
118 FileStep::Tensor => self.core.take_tensor().ok_or(Error::CoreFailed),
119 _ => Err(Error::CoreFailed),
120 }
121 }
122
123 /// Reads a composite tensor by its head name, reassembling the head with its members.
124 ///
125 /// Members are recovered by **file-offset adjacency**: the head's declared
126 /// `member_count` tensors written immediately after it, recursively for nested
127 /// composites. Recovery keys on the descriptor offset, not index position, so it is
128 /// correct even when the file used the `SORTED_INDEX` option (which reorders the index
129 /// array but not the file layout). The reassembled group is validated with
130 /// [`CompositeValidator`][hurray_core::CompositeValidator].
131 ///
132 /// # Examples
133 ///
134 /// ```no_run
135 /// # #[tokio::main]
136 /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
137 /// use hurray_io::file::FileReader;
138 ///
139 /// let file = tokio::fs::File::open("model.hrry").await?;
140 /// let mut reader = FileReader::open(file).await?;
141 /// let composite = reader.read_composite("weight").await?;
142 /// println!("{} member(s)", composite.members.len());
143 /// # Ok(())
144 /// # }
145 /// ```
146 ///
147 /// # Errors
148 ///
149 /// - [`Error::TensorNotFound`] — no tensor named `head_name`
150 /// - [`Error::NotAComposite`] — `head_name` exists but is not a composite head
151 /// - [`Error::TornComposite`] — fewer members follow the head than it declares
152 /// - [`Error::CompositeNestingTooDeep`] — nesting exceeded `max_composite_depth`
153 /// - [`Error::Core`] — composite validation failed
154 pub async fn read_composite(&mut self, head_name: &str) -> Result<FileComposite> {
155 self.core.begin_composite(head_name)?;
156 match self.drive().await? {
157 FileStep::Composite => self.core.take_composite().ok_or(Error::CoreFailed),
158 _ => Err(Error::CoreFailed),
159 }
160 }
161
162 /// Returns the underlying reader.
163 pub fn into_inner(self) -> R {
164 self.inner
165 }
166
167 /// Answers every range the core asks for, and returns the step that ends the job.
168 async fn drive(&mut self) -> Result<FileStep> {
169 loop {
170 match self.core.step()? {
171 FileStep::Range { offset, len } => {
172 fill(&mut self.inner, &mut self.core, offset, len).await?
173 }
174 other => return Ok(other),
175 }
176 }
177 }
178}
179
180/// Answers one [`FileStep::Range`]: seek, read into the core's destination, report the
181/// count.
182///
183/// A free function rather than a method so the seek and the destination can come from
184/// disjoint field borrows, which is also what `open` needs before the reader exists.
185///
186/// The transfer and the report are one step with nothing suspended between them, the shape
187/// ADR-041 § 6a requires of every driver.
188async fn fill<R: AsyncRead + AsyncSeek + Unpin>(
189 inner: &mut R,
190 core: &mut FileReaderCore,
191 offset: u64,
192 len: usize,
193) -> Result<()> {
194 inner.seek(SeekFrom::Start(offset)).await?;
195 let n = loop {
196 // `read`, not `read_exact`: a bulk helper cannot report how many bytes it moved
197 // when it fails, and that count is the read position (§ 6). Here the position is
198 // absolute anyway, so even a lost count would be recoverable — but a short read
199 // must still be reported as what it is.
200 match inner.read(&mut core.input_dst()[..len]).await {
201 Ok(n) => break n,
202 // `read_exact`, which this replaces, absorbed this internally.
203 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
204 // A read that errors moved no bytes, so the count is known to be zero and the
205 // core's position survives.
206 Err(e) => return Err(Error::Io(e)),
207 }
208 };
209 // `n == 0` means the file is shorter than its own index claims; the core calls it.
210 core.report_filled(n)
211}