From cd9a4f869f7689606495c416030d2f653ca6d982 Mon Sep 17 00:00:00 2001
From: "Joseph C. Osborn" <joseph.osborn@pomona.edu>
Date: Wed, 29 Oct 2025 12:37:33 -0700
Subject: [PATCH] fix encoding

---
 src/bin/reencode.rs |  3 ++-
 src/rply.rs         | 17 ++++++++++++-----
 src/statestream.rs  |  4 ++++
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/bin/reencode.rs b/src/bin/reencode.rs
index bbef427..6a5af96 100644
--- a/src/bin/reencode.rs
+++ b/src/bin/reencode.rs
@@ -4,7 +4,7 @@ fn main() {
     let args: Vec<_> = std::env::args().collect();
     let file =
         std::fs::File::open(args.get(1).unwrap_or(&"examples/bobl.replay".to_string())).unwrap();
-    let outfile = std::fs::File::open(
+    let outfile = std::fs::File::create(
         args.get(2)
             .unwrap_or(&"examples/bobl_smallblocks.replay".to_string()),
     )
@@ -41,4 +41,5 @@ fn main() {
     out.finish().unwrap();
     assert_eq!(out.frame_number, rply.frame_number);
     assert_eq!(out.header.frame_count(), rply.header.frame_count());
+    assert_eq!(out.header.frame_count(), Some(out.frame_number));
 }
diff --git a/src/rply.rs b/src/rply.rs
index 23ae64f..0dbeb2c 100644
--- a/src/rply.rs
+++ b/src/rply.rs
@@ -154,6 +154,8 @@ pub enum ReplayError {
     Encoding(InvalidDeterminant),
     #[error("I/O Error")]
     IO(#[from] std::io::Error),
+    #[error("Too many frames to {0} fit framecount header")]
+    TooManyFrames(std::num::TryFromIntError),
     #[error("Coreless frame read for version 0 not possible")]
     NoCoreRead(),
     #[error("Checkpoint too big {0}")]
@@ -421,7 +423,7 @@ pub fn decode<R: std::io::BufRead>(rply: &mut R) -> Result<ReplayDecoder<'_, R>>
 pub struct ReplayEncoder<'a, W: std::io::Write + std::io::Seek> {
     rply: &'a mut W,
     pub header: Header,
-    pub initial_state: Vec<u8>,
+    // pub initial_state: Vec<u8>,
     pub frame_number: u64,
     last_pos: u64,
     ss_state: statestream::Ctx,
@@ -447,7 +449,6 @@ impl<'w, W: std::io::Write + std::io::Seek> ReplayEncoder<'w, W> {
         let mut replay = ReplayEncoder {
             rply,
             header,
-            initial_state: vec![],
             frame_number: 0,
             last_pos: 0,
             ss_state,
@@ -475,6 +476,10 @@ impl<'w, W: std::io::Write + std::io::Seek> ReplayEncoder<'w, W> {
             .write_u32::<LittleEndian>(self.header.initial_state_size())?;
         self.rply
             .write_u64::<LittleEndian>(self.header.identifier())?;
+        self.rply.write_u32::<LittleEndian>(
+            u32::try_from(self.header.frame_count().unwrap())
+                .map_err(ReplayError::TooManyFrames)?,
+        )?;
         self.rply
             .write_u32::<LittleEndian>(self.header.block_size())?;
         self.rply
@@ -567,15 +572,17 @@ impl<'w, W: std::io::Write + std::io::Seek> ReplayEncoder<'w, W> {
         Ok(())
     }
     fn encode_initial_checkpoint(&mut self, checkpoint: &[u8]) -> Result<()> {
-        let initial = std::mem::take(&mut self.initial_state);
+        // let initial = std::mem::take(&mut self.initial_state);
         let old_pos = self.rply.stream_position()?;
         self.rply
             .seek(std::io::SeekFrom::Start(HEADERV2_LEN_BYTES as u64))?;
         self.encode_checkpoint(checkpoint, 0)?;
+        let encoded_size = self.rply.stream_position()? - HEADERV2_LEN_BYTES as u64;
         self.header.set_initial_state_size(
-            u32::try_from(initial.len()).map_err(ReplayError::CheckpointTooBig)?,
+            u32::try_from(encoded_size).map_err(ReplayError::CheckpointTooBig)?,
         );
-        self.initial_state = initial;
+        // self.initial_state = initial;
+        // dbg!("initial state size", self.header.initial_state_size());
         // Have to rewrite header to account for initial state size
         self.write_header()?;
         self.last_pos = self.rply.stream_position()?;
diff --git a/src/statestream.rs b/src/statestream.rs
index 91df561..7c90d11 100644
--- a/src/statestream.rs
+++ b/src/statestream.rs
@@ -256,6 +256,9 @@ impl<'w, 'c, W: std::io::Write> Encoder<'w, 'c, W> {
         let superblock_size = self.ctx.superblock_size as usize;
         let superblock_size_bytes = block_size * superblock_size;
         let superblock_count = ((checkpoint.len() - 1) / superblock_size_bytes) + 1;
+        self.ctx
+            .last_superseq
+            .resize(superblock_count.max(self.ctx.last_superseq.len()), 0);
         let mut superblock_contents = vec![0_u32; superblock_size];
         for (superblock_i, superblock_bytes) in checkpoint.chunks(superblock_size_bytes).enumerate()
         {
@@ -306,6 +309,7 @@ impl<'w, 'c, W: std::io::Write> Encoder<'w, 'c, W> {
                 }
             }
         }
+        self.ctx.last_superseq.truncate(superblock_count);
         bytes_out += rmp_size(r::write_uint(
             self.writer,
             u64::from(u8::from(SSToken::SuperblockSeq)),