From 2f8e85eac153fdea9fd73a1233431d7b388faa5a Mon Sep 17 00:00:00 2001
From: Chris Mikkelson <cmikk@fsi.io>
Date: Tue, 6 Aug 2024 22:53:00 -0600
Subject: [PATCH] writer: Move index update to data block write

Fixes missing index write on last block.
---
 src/writer/mod.rs | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/writer/mod.rs b/src/writer/mod.rs
index 6d96bd1..85090cd 100644
--- a/src/writer/mod.rs
+++ b/src/writer/mod.rs
@@ -15,6 +15,7 @@ pub struct Writer<W: std::io::Write> {
     meta: Metadata,
     block: BlockBuilder,
     index: BlockBuilder,
+    offset: usize,
     written: usize,
     last_key: Vec<u8>,
 }
@@ -105,6 +106,7 @@ impl<W: std::io::Write> Writer<W> {
             meta: Default::default(),
             block: BlockBuilder::default(),
             index: BlockBuilder::default(),
+            offset: 0,
             written: 0,
             last_key: Vec::new(),
         }
@@ -113,10 +115,7 @@ impl<W: std::io::Write> Writer<W> {
     pub fn add(&mut self, e: Entry) -> Result<()> {
         let est = e.key.len() + e.value.len() + 15;
         if self.block.len() + est >= self.blocksize {
-            let mut off_buf: [u8; 10] = Default::default();
-            let offlen = self.written.encode_var(&mut off_buf);
             bytesep(&mut self.last_key, e.key.as_slice());
-            self.index.add(self.last_key.as_slice(), &off_buf[..offlen]);
             self.write_block()?;
         }
         self.meta.add_entry(e.key.len(), e.value.len());
@@ -145,6 +144,13 @@ impl<W: std::io::Write> Writer<W> {
     }
 
     fn write_block(&mut self) -> Result<()> {
+        // update index entry
+        let mut off_buf: [u8; 10] = Default::default();
+        let offlen = self.offset.encode_var(&mut off_buf);
+        self.offset = self.written;
+        self.index.add(self.last_key.as_slice(), &off_buf[..offlen]);
+
+        // write block
         let b = self.block.as_slice();
         let cb = self.comp.compress(b)?;
         self.written += self.out.write_varint(cb.len())?;
-- 
2.50.1