From f43300b5ff3a708195d842be082e7be558d21aae Mon Sep 17 00:00:00 2001
From: Ben Pfaff <blp@cs.stanford.edu>
Date: Mon, 24 Jul 2023 22:58:30 -0700
Subject: [PATCH] work

---
 rust/src/lib.rs | 67 ++++++++++++++++++++++++++++++-------------------
 1 file changed, 41 insertions(+), 26 deletions(-)

diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index becdb20ad5..b51f1cf209 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -157,43 +157,32 @@ impl TryFrom<[u8; 4]> for Magic {
 
 pub struct Reader<R: Read> {
     r: BufReader<R>,
-    state: u32,
+    state: ReaderState,
     endianness: Option<Endian>,
 }
 
+enum ReaderState {
+    Start,
+    Headers(Endian),
+    End,
+}
+
 impl<R: Read + Seek> Reader<R> {
     pub fn new(r: R) -> Result<Reader<R>, Error> {
         Ok(Reader {
             r: BufReader::new(r),
-            state: 0,
+            state: ReaderState::Start,
             endianness: None,
         })
     }
-
-    pub fn read(&mut self) -> Result<Option<Record>, Error> {
-        let retval = self.do_read();
-        match retval {
-            Ok(None) => {
-                self.state = u32::MAX;
-                Ok(None)
-            }
-            Ok(Some((record, next_state))) => {
-                self.state = next_state;
-                Ok(Some(record))
-            }
-            Err(error) => Err(error),
-        }
-    }
-
-    pub fn do_read(&mut self) -> Result<Option<(Record, u32)>, Error> {
+    fn _next(&mut self) -> Result<Option<(Record, ReaderState)>, Error> {
         match self.state {
-            0 => {
+            ReaderState::Start => {
                 let header = read_header(&mut self.r)?;
-                self.endianness = Some(header.endianness);
-                Ok(Some((Record::Header(header), 1)))
+                let endianness = header.endianness;
+                Ok(Some((Record::Header(header), ReaderState::Headers(endianness))))
             }
-            1 => {
-                let e = self.endianness.unwrap();
+            ReaderState::Headers(e) => {
                 let rec_type: u32 = e.parse(read_bytes(&mut self.r)?);
                 let record = match rec_type {
                     2 => Record::Variable(read_variable_record(&mut self.r, e)?),
@@ -201,7 +190,10 @@ impl<R: Read + Seek> Reader<R> {
                     4 => Record::VarIndexes(read_var_indexes_record(&mut self.r, e)?),
                     6 => Record::Document(read_document_record(&mut self.r, e)?),
                     7 => Record::Extension(read_extension_record(&mut self.r, e)?),
-                    999 => Record::EndOfHeaders,
+                    999 => {
+                        let _: [u8; 4] = read_bytes(&mut self.r)?;
+                        return Ok(Some((Record::EndOfHeaders, ReaderState::End)))
+                    },
                     _ => {
                         return Err(Error::BadRecordType {
                             offset: self.r.stream_position()?,
@@ -209,7 +201,30 @@ impl<R: Read + Seek> Reader<R> {
                         })
                     }
                 };
-                Ok(Some((record, 1)))
+                Ok(Some((record, ReaderState::Headers(e))))
+            }
+            ReaderState::End => Ok(None),
+        }
+    }
+}
+
+impl<R: Read + Seek> Iterator for Reader<R> {
+    type Item = Result<Record, Error>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let retval = self._next();
+        match retval {
+            Ok(None) => {
+                self.state = ReaderState::End;
+                None
+            }
+            Ok(Some((record, next_state))) => {
+                self.state = next_state;
+                Some(Ok(record))
+            }
+            Err(error) => {
+                self.state = ReaderState::End;
+                Some(Err(error))
             }
         }
     }
-- 
2.30.2