From c27eccc1e3175628aa1be0de9caed313be1aee42 Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@cs.stanford.edu>
Date: Mon, 31 Jul 2023 11:06:57 -0700
Subject: [PATCH] work

---
 rust/src/raw.rs | 67 +++++++++++++++++++++++++------------------------
 1 file changed, 34 insertions(+), 33 deletions(-)

diff --git a/rust/src/raw.rs b/rust/src/raw.rs
index 7cb7238e9a..b09f94e73a 100644
--- a/rust/src/raw.rs
+++ b/rust/src/raw.rs
@@ -686,7 +686,7 @@ pub struct Document {
     pub pos: u64,
 
     /// The document, as an array of 80-byte lines.
-    pub lines: Vec<[u8; DOC_LINE_LEN as usize]>,
+    pub lines: Vec<[u8; Document::LINE_LEN as usize]>,
 }
 
 impl Document {
@@ -696,25 +696,22 @@ impl Document {
 
     /// Maximum number of lines we will accept in a document.  This is simply
     /// the maximum number that will fit in a 32-bit space.
-    pub const MAX_LINES: u32 = i32::MAX as u32 / DOC_LINE_LEN;
+    pub const MAX_LINES: u32 = i32::MAX as u32 / Self::LINE_LEN;
 
     fn read<R: Read + Seek>(r: &mut R, endian: Endian) -> Result<Document, Error> {
         let offset = r.stream_position()?;
         let n: u32 = endian.parse(read_bytes(r)?);
         match n {
-            0..=DOC_MAX_LINES => {
-                let pos = r.stream_position()?;
-                let mut lines = Vec::with_capacity(n as usize);
-                for _ in 0..n {
-                    let line: [u8; 80] = read_bytes(r)?;
-                    lines.push(line);
-                }
-                Ok(Document { pos, lines })
-            }
+            0..=Self::MAX_LINES => Ok(Document {
+                pos: r.stream_position()?,
+                lines: (0..n)
+                    .map(|_| read_bytes(r))
+                    .collect::<Result<Vec<_>, _>>()?,
+            }),
             _ => Err(Error::BadDocumentLength {
                 offset,
                 n,
-                max: DOC_MAX_LINES,
+                max: Self::MAX_LINES,
             }),
         }
     }
@@ -897,21 +894,35 @@ pub struct ZBlock {
     pub compressed_size: u32,
 }
 
-impl ZTrailer {
+impl ZBlock {
     fn read<R: Read + Seek>(
         r: &mut R,
         endian: Endian,
+    ) -> Result<ZBlock, Error> {
+        Ok(ZBlock {
+            uncompressed_ofs: endian.parse(read_bytes(r)?),
+            compressed_ofs: endian.parse(read_bytes(r)?),
+            uncompressed_size: endian.parse(read_bytes(r)?),
+            compressed_size: endian.parse(read_bytes(r)?),
+        })
+    }
+}
+
+impl ZTrailer {
+    fn read<R: Read + Seek>(
+        reader: &mut R,
+        endian: Endian,
         ztrailer_ofs: u64,
         ztrailer_len: u64,
     ) -> Result<Option<ZTrailer>, Error> {
-        let start_offset = r.stream_position()?;
-        if r.seek(SeekFrom::Start(ztrailer_ofs)).is_err() {
+        let start_offset = reader.stream_position()?;
+        if reader.seek(SeekFrom::Start(ztrailer_ofs)).is_err() {
             return Ok(None);
         }
-        let int_bias = endian.parse(read_bytes(r)?);
-        let zero = endian.parse(read_bytes(r)?);
-        let block_size = endian.parse(read_bytes(r)?);
-        let n_blocks: u32 = endian.parse(read_bytes(r)?);
+        let int_bias = endian.parse(read_bytes(reader)?);
+        let zero = endian.parse(read_bytes(reader)?);
+        let block_size = endian.parse(read_bytes(reader)?);
+        let n_blocks: u32 = endian.parse(read_bytes(reader)?);
         let expected_n_blocks = (ztrailer_len - 24) / 24;
         if n_blocks as u64 != expected_n_blocks {
             return Err(Error::BadZlibTrailerNBlocks {
@@ -921,20 +932,10 @@ impl ZTrailer {
                 ztrailer_len,
             });
         }
-        let mut blocks = Vec::with_capacity(n_blocks as usize);
-        for _ in 0..n_blocks {
-            let uncompressed_ofs = endian.parse(read_bytes(r)?);
-            let compressed_ofs = endian.parse(read_bytes(r)?);
-            let uncompressed_size = endian.parse(read_bytes(r)?);
-            let compressed_size = endian.parse(read_bytes(r)?);
-            blocks.push(ZBlock {
-                uncompressed_ofs,
-                compressed_ofs,
-                uncompressed_size,
-                compressed_size,
-            });
-        }
-        r.seek(SeekFrom::Start(start_offset))?;
+        let blocks = (0..n_blocks)
+                    .map(|_| ZBlock::read(reader, endian))
+                    .collect::<Result<Vec<_>, _>>()?;
+        reader.seek(SeekFrom::Start(start_offset))?;
         Ok(Some(ZTrailer {
             offset: ztrailer_ofs,
             int_bias,
-- 
2.30.2