diff --git a/mkpfs/pfs.py b/mkpfs/pfs.py index 6ed986f..33fa8de 100755 --- a/mkpfs/pfs.py +++ b/mkpfs/pfs.py @@ -1005,6 +1005,7 @@ 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. @@ -1012,17 +1013,22 @@ def _should_store_pfsc_block_compressed( 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: diff --git a/tests/mkpfs/test_pfs.py b/tests/mkpfs/test_pfs.py index e094fca..c9933a7 100644 --- a/tests/mkpfs/test_pfs.py +++ b/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(" 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(" 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(" 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(" 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(" 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."""