Skip to main content

rapidgzip_core/index/
gzidx.rs

1//! indexed_gzip `GZIDX` import and export.
2
3use super::{
4    Checkpoint, CheckpointKind, DeflateIndex, IndexError, IndexKind, IndexReadOptions,
5    StoredWindow, WINDOW_SIZE, read_exact_bytes, read_u8, read_u32_le, read_u64_le, write_u32_le,
6    write_u64_le,
7};
8use std::io::{Read, Write};
9
10const MAGIC: &[u8; 5] = b"GZIDX";
11const MAX_VERSION: u8 = 1;
12
13/// Encodes an absolute raw-DEFLATE bit offset as zran's byte/bits pair.
14#[must_use]
15pub fn encode_bit_offset(compressed_offset_in_bits: u64) -> (u64, u8) {
16    let remainder = (compressed_offset_in_bits % 8) as u8;
17    if remainder == 0 {
18        (compressed_offset_in_bits / 8, 0)
19    } else {
20        (compressed_offset_in_bits / 8 + 1, 8 - remainder)
21    }
22}
23
24/// Decodes zran's byte/bits pair into an absolute raw-DEFLATE bit offset.
25pub fn decode_bit_offset(byte_offset: u64, bits_field: u8) -> Result<u64, IndexError> {
26    if bits_field >= 8 {
27        return Err(IndexError::InvalidCheckpoint(
28            "denormal compressed offset: bits field is 8 or more",
29        ));
30    }
31    let bit_offset = byte_offset
32        .checked_mul(8)
33        .ok_or(IndexError::InvalidCheckpoint(
34            "compressed byte offset overflows a bit count",
35        ))?;
36    if bits_field == 0 {
37        return Ok(bit_offset);
38    }
39    if bit_offset == 0 {
40        return Err(IndexError::InvalidCheckpoint(
41            "denormal compressed offset: bits field before the source start",
42        ));
43    }
44    Ok(bit_offset - u64::from(bits_field))
45}
46
47pub(crate) fn write_gzidx(index: &DeflateIndex, writer: &mut impl Write) -> Result<(), IndexError> {
48    index.validate()?;
49    if !matches!(index.kind, IndexKind::Gzip | IndexKind::Bgzf) {
50        return Err(IndexError::IncompatibleFormat {
51            operation: "GZIDX export",
52            kind: index.kind,
53        });
54    }
55    let compressed_size = index
56        .compressed_size_in_bytes
57        .ok_or(IndexError::MissingMetadata(
58            "compressed size for GZIDX export",
59        ))?;
60    let uncompressed_size = index
61        .uncompressed_size_in_bytes
62        .ok_or(IndexError::MissingMetadata(
63            "uncompressed size for GZIDX export",
64        ))?;
65    let spacing = index
66        .checkpoint_spacing_in_bytes
67        .ok_or(IndexError::MissingMetadata(
68            "checkpoint spacing for GZIDX export",
69        ))?;
70    let spacing = u32::try_from(spacing).map_err(|_| IndexError::ExcessiveLength {
71        what: "GZIDX checkpoint spacing",
72        value: spacing,
73    })?;
74    let count =
75        u32::try_from(index.checkpoints.len()).map_err(|_| IndexError::ExcessiveLength {
76            what: "checkpoint count",
77            value: u64::try_from(index.checkpoints.len()).unwrap_or(u64::MAX),
78        })?;
79    if index
80        .checkpoints
81        .iter()
82        .any(|point| matches!(point.kind, CheckpointKind::GzipMemberHeader))
83    {
84        return Err(IndexError::InvalidCheckpoint(
85            "GZIDX export requires raw-DEFLATE resume offsets",
86        ));
87    }
88
89    writer.write_all(MAGIC).map_err(IndexError::io)?;
90    writer
91        .write_all(&[MAX_VERSION, 0])
92        .map_err(IndexError::io)?;
93    write_u64_le(writer, compressed_size)?;
94    write_u64_le(writer, uncompressed_size)?;
95    write_u32_le(writer, spacing)?;
96    write_u32_le(writer, WINDOW_SIZE as u32)?;
97    write_u32_le(writer, count)?;
98
99    for checkpoint in &index.checkpoints {
100        let (byte_offset, bits_field) = encode_bit_offset(checkpoint.compressed_offset_in_bits);
101        write_u64_le(writer, byte_offset)?;
102        write_u64_le(writer, checkpoint.uncompressed_offset_in_bytes)?;
103        let has_window = index
104            .windows
105            .get(checkpoint.compressed_offset_in_bits)
106            .is_some();
107        writer
108            .write_all(&[bits_field, u8::from(has_window)])
109            .map_err(IndexError::io)?;
110    }
111    for checkpoint in &index.checkpoints {
112        if let Some(window) = index.windows.get(checkpoint.compressed_offset_in_bits) {
113            writer
114                .write_all(&window.decompressed()?)
115                .map_err(IndexError::io)?;
116        }
117    }
118    Ok(())
119}
120
121pub(crate) fn read_gzidx(
122    reader: &mut impl Read,
123    archive_size: Option<u64>,
124    options: IndexReadOptions,
125) -> Result<DeflateIndex, IndexError> {
126    let mut magic = [0_u8; 5];
127    read_exact_bytes(reader, &mut magic)?;
128    if &magic != MAGIC {
129        return Err(IndexError::BadMagic {
130            found: magic.to_vec(),
131        });
132    }
133    let version = read_u8(reader)?;
134    if version > MAX_VERSION {
135        return Err(IndexError::UnsupportedVersion(u64::from(version)));
136    }
137    let flags = read_u8(reader)?;
138    if flags != 0 {
139        return Err(IndexError::UnsupportedFlags {
140            flags: u64::from(flags),
141        });
142    }
143
144    let compressed_size = read_u64_le(reader)?;
145    let uncompressed_size = read_u64_le(reader)?;
146    let spacing = read_u32_le(reader)?;
147    let window_size = read_u32_le(reader)?;
148    if window_size as usize != WINDOW_SIZE {
149        return Err(IndexError::InvalidWindowSize(u64::from(window_size)));
150    }
151    if archive_size.is_some_and(|size| size != compressed_size) {
152        return Err(IndexError::ArchiveSizeMismatch {
153            index_size: compressed_size,
154            archive_size: archive_size.expect("checked as some"),
155        });
156    }
157
158    let count = read_u32_le(reader)? as usize;
159    if count > options.max_checkpoints {
160        return Err(IndexError::ExcessiveLength {
161            what: "checkpoint count",
162            value: count as u64,
163        });
164    }
165    let mut records = Vec::new();
166    records
167        .try_reserve(count)
168        .map_err(|_| IndexError::AllocationFailed {
169            what: "GZIDX checkpoint records",
170        })?;
171    let mut window_count = 0_u64;
172    for position in 0..count {
173        let byte_offset = read_u64_le(reader)?;
174        let uncompressed_offset_in_bytes = read_u64_le(reader)?;
175        let bits_field = read_u8(reader)?;
176        let has_window = if version == 0 {
177            position != 0
178        } else {
179            match read_u8(reader)? {
180                0 => false,
181                1 => true,
182                _ => {
183                    return Err(IndexError::InvalidCheckpoint(
184                        "GZIDX window flag is not zero or one",
185                    ));
186                }
187            }
188        };
189        if byte_offset > compressed_size {
190            return Err(IndexError::InvalidCheckpoint(
191                "checkpoint compressed offset is after the source end",
192            ));
193        }
194        if has_window {
195            window_count = window_count
196                .checked_add(1)
197                .ok_or(IndexError::ExcessiveLength {
198                    what: "aggregate window bytes",
199                    value: u64::MAX,
200                })?;
201        }
202        records.push((
203            Checkpoint {
204                compressed_offset_in_bits: decode_bit_offset(byte_offset, bits_field)?,
205                uncompressed_offset_in_bytes,
206                kind: CheckpointKind::DeflateBlock,
207                line_offset: None,
208            },
209            has_window,
210        ));
211    }
212    let window_bytes =
213        window_count
214            .checked_mul(WINDOW_SIZE as u64)
215            .ok_or(IndexError::ExcessiveLength {
216                what: "aggregate window bytes",
217                value: u64::MAX,
218            })?;
219    if WINDOW_SIZE > options.max_window_payload_bytes || window_bytes > options.max_window_bytes {
220        return Err(IndexError::ExcessiveLength {
221            what: "aggregate window bytes",
222            value: window_bytes,
223        });
224    }
225
226    let mut index = DeflateIndex::new();
227    index.compressed_size_in_bytes = Some(compressed_size);
228    index.uncompressed_size_in_bytes = Some(uncompressed_size);
229    index.checkpoint_spacing_in_bytes = (spacing != 0).then_some(u64::from(spacing));
230    for (checkpoint, has_window) in records {
231        let window = if has_window {
232            let mut payload = Vec::new();
233            payload
234                .try_reserve_exact(WINDOW_SIZE)
235                .map_err(|_| IndexError::AllocationFailed {
236                    what: "GZIDX window",
237                })?;
238            payload.resize(WINDOW_SIZE, 0);
239            read_exact_bytes(reader, &mut payload)?;
240            StoredWindow::from_raw(payload)?
241        } else {
242            StoredWindow::empty()
243        };
244        index.push(checkpoint, window)?;
245    }
246    index.validate()?;
247    Ok(index)
248}