[EXPERIMENT] Avoid corruption from the last compressed PFSC block

mkpfs/pfs.py

@@ -1005,24 +1005,30 @@ def _should_store_pfsc_block_compressed(
     logical_block_size: int,
     gain_pct: float,
     threshold_gain: int,
+    is_last_block: bool,
 ) -> bool:
     """Return whether a PFSC block can be stored in compressed form.
 
     PFSC decoding distinguishes raw and compressed blocks only by comparing the
     stored block span with the logical block size. A compressed block must
     therefore be strictly smaller than the logical block size, otherwise the
     decoder interprets it as raw bytes and the payload becomes self-inconsistent.
+    This policy also keeps the final logical block raw as a workaround for the
+    target runtime's last-block handling.
 
     Args:
         compressed_block_size: Encoded zlib block length in bytes.
         logical_block_size: PFSC logical block size in bytes.
         gain_pct: Percent gain achieved by compressing the padded logical block.
         threshold_gain: Minimum gain required to keep the compressed bytes.
+        is_last_block: Whether the current logical block is the file's final block.
 
     Returns:
         ``True`` when the block should be stored compressed, ``False`` when it
         must remain raw.
     """
+    if is_last_block:
+        return False
     return compressed_block_size < logical_block_size and gain_pct >= threshold_gain
 
 
@@ -1065,7 +1071,7 @@ def encode_pfsc_payload(
     all_compressed_size: int = 0
     compressed_blocks: int = 0
 
-    for block in logical_blocks:
+    for block_index, block in enumerate(logical_blocks):
         padded_block: bytes = block.ljust(logical_block_size, b"\x00")
         compressed_block: bytes = zlib.compress(padded_block, level=zlib_level)
         all_compressed_size += len(compressed_block)
@@ -1075,6 +1081,7 @@ def encode_pfsc_payload(
             logical_block_size=logical_block_size,
             gain_pct=gain_pct,
             threshold_gain=threshold_gain,
+            is_last_block=block_index == (block_count - 1),
         )
         chosen_block: bytes = compressed_block if store_compressed else padded_block
         if store_compressed:
@@ -1309,7 +1316,7 @@ def _analyze_pfsc_file_storage(
 
     if effective_block_workers == 1:
         with abs_path.open("rb") as source_file:
-            for _idx in range(block_count):
+            for block_index in range(block_count):
                 chunk: bytes = source_file.read(logical_block_size)
                 padded_chunk: bytes = chunk.ljust(logical_block_size, b"\x00")
                 compressed_chunk: bytes = zlib.compress(padded_chunk, level=zlib_level)
@@ -1320,6 +1327,7 @@ def _analyze_pfsc_file_storage(
                     logical_block_size=logical_block_size,
                     gain_pct=gain_pct,
                     threshold_gain=threshold_gain,
+                    is_last_block=block_index == (block_count - 1),
                 ):
                     chosen_payload_size += len(compressed_chunk)
                     compressed_blocks += 1
@@ -1341,7 +1349,7 @@ def _analyze_pfsc_file_storage(
             results_iter = pool.imap(_compress_pfsc_block_lengths_worker, worker_args_iter, chunksize=1)
             raw_block_len: int
             compressed_block_len: int
-            for raw_block_len, compressed_block_len in results_iter:
+            for block_index, (raw_block_len, compressed_block_len) in enumerate(results_iter):
                 all_compressed_size += compressed_block_len
                 padded_block_len: int = logical_block_size
                 gain_pct: float = ((padded_block_len - compressed_block_len) / padded_block_len) * 100.0
@@ -1350,6 +1358,7 @@ def _analyze_pfsc_file_storage(
                     logical_block_size=logical_block_size,
                     gain_pct=gain_pct,
                     threshold_gain=threshold_gain,
+                    is_last_block=block_index == (block_count - 1),
                 ):
                     chosen_payload_size += compressed_block_len
                     compressed_blocks += 1
@@ -1418,7 +1427,7 @@ def _encode_pfsc_into_handle(
     out.seek(base_offset + header_size)
     if effective_block_workers == 1:
         with source_path.open("rb") as source_file:
-            for _idx in range(block_count):
+            for block_index in range(block_count):
                 chunk: bytes = source_file.read(logical_block_size)
                 padded_chunk: bytes = chunk.ljust(logical_block_size, b"\x00")
                 compressed_chunk: bytes = zlib.compress(padded_chunk, level=zlib_level)
@@ -1429,6 +1438,7 @@ def _encode_pfsc_into_handle(
                     logical_block_size=logical_block_size,
                     gain_pct=gain_pct,
                     threshold_gain=threshold_gain,
+                    is_last_block=block_index == (block_count - 1),
                 )
                 selected_chunk: bytes = compressed_chunk if store_compressed else padded_chunk
                 if store_compressed:
@@ -1451,7 +1461,7 @@ def _encode_pfsc_into_handle(
             results_iter = pool.imap(_compress_pfsc_block_payload_worker, worker_args_iter, chunksize=1)
             raw_chunk: bytes
             compressed_chunk: bytes
-            for raw_chunk, compressed_chunk in results_iter:
+            for block_index, (raw_chunk, compressed_chunk) in enumerate(results_iter):
                 padded_chunk: bytes = raw_chunk.ljust(logical_block_size, b"\x00")
                 all_compressed_size += len(compressed_chunk)
                 gain_pct: float = ((len(padded_chunk) - len(compressed_chunk)) / len(padded_chunk)) * 100.0
@@ -1460,6 +1470,7 @@ def _encode_pfsc_into_handle(
                     logical_block_size=logical_block_size,
                     gain_pct=gain_pct,
                     threshold_gain=threshold_gain,
+                    is_last_block=block_index == (block_count - 1),
                 )
                 selected_chunk: bytes = compressed_chunk if store_compressed else padded_chunk
                 if store_compressed:

tests/mkpfs/test_pfs.py

@@ -1010,7 +1010,7 @@ def test_pfsc_encode_decode_round_trip(self) -> None:
         )
         assert encoded != raw
         assert gain_pct > 0.0
-        assert hypothetical_size >= len(encoded)
+        assert hypothetical_size > 0
         assert encoded[:4] == b"PFSC"
         decoded: bytes = pfs_mod.decode_pfsc_payload(payload=encoded, expected_logical_size=len(raw))
         assert decoded == raw
@@ -1026,6 +1026,62 @@ def test_pfsc_encode_decode_round_trip(self) -> None:
         assert data_start >= c.PFSC_INITIAL_DATA_OFFSET
         assert data_length == 3 * c.PFSC_LOGICAL_BLOCK_SIZE
 
+    def test_pfsc_encode_keeps_last_logical_block_raw(self) -> None:
+        """PFSC encoding should keep the final logical block raw even when compressible."""
+        raw: bytes = (b"A" * c.PFSC_LOGICAL_BLOCK_SIZE) + (b"B" * c.PFSC_LOGICAL_BLOCK_SIZE) + (b"\x00" * 1234)
+        encoded: bytes
+        gain_pct: float
+        _hypothetical_size: int
+        encoded, gain_pct, _hypothetical_size = pfs_mod.encode_pfsc_payload(
+            raw=raw,
+            threshold_gain=0,
+            zlib_level=9,
+            logical_block_size=c.PFSC_LOGICAL_BLOCK_SIZE,
+        )
+
+        assert encoded != raw
+        assert gain_pct > 0.0
+        _magic: int
+        _unk4: int
+        _unk8: int
+        logical_block_size: int
+        _logical_block_size_2: int
+        block_offsets_offset: int
+        _data_offset: int
+        _logical_size: int
+        (
+            _magic,
+            _unk4,
+            _unk8,
+            logical_block_size,
+            _logical_block_size_2,
+            block_offsets_offset,
+            _data_offset,
+            _logical_size,
+        ) = struct.unpack_from("<iiiiqqQq", encoded, 0)
+        offsets: tuple[int, int, int, int]
+        offsets = struct.unpack_from("<4Q", encoded, block_offsets_offset)
+        assert offsets[1] - offsets[0] < logical_block_size
+        assert offsets[2] - offsets[1] < logical_block_size
+        assert offsets[3] - offsets[2] == logical_block_size
+        assert pfs_mod.decode_pfsc_payload(payload=encoded, expected_logical_size=len(raw)) == raw
+
+    def test_pfsc_encode_single_block_file_falls_back_to_raw(self) -> None:
+        """A one-block file should fall back to raw storage when the last block cannot compress."""
+        raw: bytes = b"\x00" * 1234
+        encoded: bytes
+        gain_pct: float
+        _hypothetical_size: int
+        encoded, gain_pct, _hypothetical_size = pfs_mod.encode_pfsc_payload(
+            raw=raw,
+            threshold_gain=0,
+            zlib_level=9,
+            logical_block_size=c.PFSC_LOGICAL_BLOCK_SIZE,
+        )
+
+        assert encoded == raw
+        assert gain_pct <= 0.0
+
     def test_pfsc_encode_reports_incremental_progress_bytes(self) -> None:
         """PFSC encoding should report raw bytes as each logical block is processed."""
         raw: bytes = (b"A" * c.PFSC_LOGICAL_BLOCK_SIZE) + (b"B" * 1234)
@@ -1048,6 +1104,7 @@ def test_pfsc_encode_keeps_equal_sized_compressed_block_raw(self) -> None:
             (b"A" * c.PFSC_LOGICAL_BLOCK_SIZE)
             + (b"B" * c.PFSC_LOGICAL_BLOCK_SIZE)
             + (b"C" * c.PFSC_LOGICAL_BLOCK_SIZE)
+            + (b"\x00" * 1234)
         )
         real_compress: Callable[..., bytes] = pfs_mod.zlib.compress
         compress_call_count: int = 0
@@ -1089,11 +1146,12 @@ def fake_compress(data: bytes, *, level: int) -> bytes:
             _data_offset,
             _logical_size,
         ) = struct.unpack_from("<iiiiqqQq", encoded, 0)
-        offsets: tuple[int, int, int, int]
-        offsets = struct.unpack_from("<4Q", encoded, block_offsets_offset)
+        offsets: tuple[int, int, int, int, int]
+        offsets = struct.unpack_from("<5Q", encoded, block_offsets_offset)
         assert offsets[1] - offsets[0] == logical_block_size
         assert offsets[2] - offsets[1] < logical_block_size
         assert offsets[3] - offsets[2] < logical_block_size
+        assert offsets[4] - offsets[3] == logical_block_size
         assert pfs_mod.decode_pfsc_payload(payload=encoded, expected_logical_size=len(raw)) == raw
 
     def test_pfsc_spool_keeps_equal_sized_compressed_block_raw(self) -> None:
@@ -1105,6 +1163,7 @@ def test_pfsc_spool_keeps_equal_sized_compressed_block_raw(self) -> None:
             (b"A" * c.PFSC_LOGICAL_BLOCK_SIZE)
             + (b"B" * c.PFSC_LOGICAL_BLOCK_SIZE)
             + (b"C" * c.PFSC_LOGICAL_BLOCK_SIZE)
+            + (b"\x00" * 1234)
         )
         source_path.write_bytes(raw)
         real_compress: Callable[..., bytes] = pfs_mod.zlib.compress
@@ -1132,9 +1191,114 @@ def fake_compress(data: bytes, *, level: int) -> bytes:
                 logical_block_size=c.PFSC_LOGICAL_BLOCK_SIZE,
                 block_worker_count=1,
             )
+        encoded: bytes = spool_path.read_bytes()
+        _magic: int
+        _unk4: int
+        _unk8: int
+        logical_block_size: int
+        _logical_block_size_2: int
+        block_offsets_offset: int
+        _data_offset: int
+        _logical_size: int
+        (
+            _magic,
+            _unk4,
+            _unk8,
+            logical_block_size,
+            _logical_block_size_2,
+            block_offsets_offset,
+            _data_offset,
+            _logical_size,
+        ) = struct.unpack_from("<iiiiqqQq", encoded, 0)
+        offsets: tuple[int, int, int, int, int]
+        offsets = struct.unpack_from("<5Q", encoded, block_offsets_offset)
+        assert is_compressed
+        assert stored_size == spool_path.stat().st_size
+        assert offsets[1] - offsets[0] == logical_block_size
+        assert offsets[2] - offsets[1] < logical_block_size
+        assert offsets[3] - offsets[2] < logical_block_size
+        assert offsets[4] - offsets[3] == logical_block_size
+        assert pfs_mod.decode_pfsc_payload(payload=encoded, expected_logical_size=len(raw)) == raw
+
+    def test_pfsc_spool_keeps_last_logical_block_raw(self) -> None:
+        """Streaming PFSC encoding should keep the final logical block raw."""
+        tmp_path: Path = self.make_temp_path()
+        source_path: Path = tmp_path / "source.bin"
+        spool_path: Path = tmp_path / "out.pfsc"
+        raw: bytes = (b"A" * c.PFSC_LOGICAL_BLOCK_SIZE) + (b"B" * c.PFSC_LOGICAL_BLOCK_SIZE) + (b"\x00" * 1234)
+        source_path.write_bytes(raw)
+
+        stored_size: int
+        is_compressed: bool
+        gain_pct: float
+        _hypothetical_size: int
+        stored_size, is_compressed, gain_pct, _hypothetical_size = pfs_mod._encode_pfsc_file_to_spool(
+            abs_path=source_path,
+            spool_path=spool_path,
+            threshold_gain=0,
+            min_file_gain=0,
+            zlib_level=9,
+            logical_block_size=c.PFSC_LOGICAL_BLOCK_SIZE,
+            block_worker_count=1,
+        )
+
+        encoded: bytes = spool_path.read_bytes()
+        _magic: int
+        _unk4: int
+        _unk8: int
+        logical_block_size: int
+        _logical_block_size_2: int
+        block_offsets_offset: int
+        _data_offset: int
+        _logical_size: int
+        (
+            _magic,
+            _unk4,
+            _unk8,
+            logical_block_size,
+            _logical_block_size_2,
+            block_offsets_offset,
+            _data_offset,
+            _logical_size,
+        ) = struct.unpack_from("<iiiiqqQq", encoded, 0)
+        offsets: tuple[int, int, int, int]
+        offsets = struct.unpack_from("<4Q", encoded, block_offsets_offset)
         assert is_compressed
+        assert gain_pct > 0.0
         assert stored_size == spool_path.stat().st_size
-        assert pfs_mod.decode_pfsc_payload(payload=spool_path.read_bytes(), expected_logical_size=len(raw)) == raw
+        assert offsets[1] - offsets[0] < logical_block_size
+        assert offsets[2] - offsets[1] < logical_block_size
+        assert offsets[3] - offsets[2] == logical_block_size
+        assert pfs_mod.decode_pfsc_payload(payload=encoded, expected_logical_size=len(raw)) == raw
+
+    def test_pfsc_analyze_matches_last_block_spool_storage(self) -> None:
+        """PFSC analysis should match spool output when the final block stays raw."""
+        tmp_path: Path = self.make_temp_path()
+        source_path: Path = tmp_path / "source.bin"
+        raw: bytes = (b"A" * c.PFSC_LOGICAL_BLOCK_SIZE) + (b"B" * c.PFSC_LOGICAL_BLOCK_SIZE) + (b"\x00" * 1234)
+        source_path.write_bytes(raw)
+        spool_path: Path = tmp_path / "out.pfsc"
+
+        analyzed: tuple[int, bool, float, int] = pfs_mod._analyze_pfsc_file_storage(
+            abs_path=source_path,
+            threshold_gain=0,
+            min_file_gain=0,
+            zlib_level=9,
+            logical_block_size=c.PFSC_LOGICAL_BLOCK_SIZE,
+            block_worker_count=2,
+        )
+        stored: tuple[int, bool, float, int] = pfs_mod._encode_pfsc_file_to_spool(
+            abs_path=source_path,
+            spool_path=spool_path,
+            threshold_gain=0,
+            min_file_gain=0,
+            zlib_level=9,
+            logical_block_size=c.PFSC_LOGICAL_BLOCK_SIZE,
+            block_worker_count=2,
+        )
+
+        assert analyzed == stored
+        assert stored[0] == spool_path.stat().st_size
 
     def test_compute_file_storage_worker_batches_progress_updates(self) -> None:
         """Compression workers should batch byte deltas before reporting them upstream."""
@@ -1322,8 +1486,30 @@ def test_parallel_pfsc_spool_matches_single_worker_output(self) -> None:
             block_worker_count=2,
         )
 
+        sequential_encoded: bytes = sequential_spool.read_bytes()
+        _magic: int
+        _unk4: int
+        _unk8: int
+        logical_block_size: int
+        _logical_block_size_2: int
+        block_offsets_offset: int
+        _data_offset: int
+        _logical_size: int
+        (
+            _magic,
+            _unk4,
+            _unk8,
+            logical_block_size,
+            _logical_block_size_2,
+            block_offsets_offset,
+            _data_offset,
+            _logical_size,
+        ) = struct.unpack_from("<iiiiqqQq", sequential_encoded, 0)
+        offsets: tuple[int, int, int, int]
+        offsets = struct.unpack_from("<4Q", sequential_encoded, block_offsets_offset)
+        assert offsets[3] - offsets[2] == logical_block_size
         assert sequential_result == parallel_result
-        assert sequential_spool.read_bytes() == parallel_spool.read_bytes()
+        assert sequential_encoded == parallel_spool.read_bytes()
 
     def test_run_image_check_reports_logical_and_stored_bytes_for_compressed_files(self) -> None:
         """Verify report output should label logical and stored byte counts correctly."""
@@ -1371,7 +1557,7 @@ def test_run_image_check_reports_logical_and_stored_bytes_for_compressed_files(s
         assert errors == []
         report_text: str = output_buffer.getvalue()
         assert "Logical file bytes:    196,608" in report_text
-        assert "Stored file bytes:     65,791" in report_text
+        assert "Stored file bytes:     131,242" in report_text
 
     def test_compression_phase_emits_intermediate_progress_for_single_worker(self) -> None:
         """Single-worker compression should emit intermediate progress updates."""