urna_format/sections/
spans.rs

1//! `chunks_original_spans` section (`SECTION_CHUNKS_ORIGINAL_SPANS = 0x03`).
2//! `(source_uri, byte_start, byte_end)` per chunk - the offset into the
3//! original source document the chunk text came from. Required for
4//! citation resolution.
5//!
6//! the raw encoding repeats the full `source_uri` for every chunk. the
7//! `intpack` repack (encoding id 4, kind 1) dedups the uris into a pool
8//! and bitpacks `(uri_index, byte_start, byte_end - byte_start)`, then
9//! reconstructs the exact raw payload on decode, so `content_hash` is
10//! byte-identical and this canonical section is never version-bumped.
11
12use std::collections::HashMap;
13
14use super::REPACK_KIND_SPANS;
15use super::codec::{Cursor, read_prefix, write_lp_str, write_prefix};
16use crate::encoding::{pack_u64s, unpack_u64s};
17use crate::error::UrnaError;
18use crate::layout::SECTION_CHUNKS_ORIGINAL_SPANS;
19
20#[derive(Clone, Debug, PartialEq)]
21pub struct OriginalSpan {
22    pub source_uri: String,
23    pub byte_start: u64,
24    pub byte_end: u64,
25}
26
27pub fn encode_chunks_original_spans(spans: &[OriginalSpan]) -> crate::Result<Vec<u8>> {
28    let mut buf = Vec::new();
29    write_prefix(&mut buf, spans.len() as u64);
30    for s in spans {
31        write_lp_str(&mut buf, &s.source_uri)?;
32        buf.extend_from_slice(&s.byte_start.to_le_bytes());
33        buf.extend_from_slice(&s.byte_end.to_le_bytes());
34    }
35    Ok(buf)
36}
37
38pub fn decode_chunks_original_spans(
39    data: &[u8],
40    expected_count: usize,
41) -> crate::Result<Vec<OriginalSpan>> {
42    let mut c = Cursor::new(data, SECTION_CHUNKS_ORIGINAL_SPANS);
43    let count = read_prefix(&mut c)? as usize;
44    if count != expected_count {
45        return Err(UrnaError::SectionCountMismatch {
46            section_id: SECTION_CHUNKS_ORIGINAL_SPANS,
47            expected: expected_count,
48            got: count,
49        });
50    }
51    let mut out = Vec::with_capacity(count);
52    for _ in 0..count {
53        let source_uri = c.read_lp_str()?;
54        let byte_start = c.read_u64()?;
55        let byte_end = c.read_u64()?;
56        out.push(OriginalSpan {
57            source_uri,
58            byte_start,
59            byte_end,
60        });
61    }
62    c.finish()?;
63    Ok(out)
64}
65
66fn read_blob<'a>(c: &mut Cursor<'a>) -> crate::Result<&'a [u8]> {
67    let len = c.read_u32()? as usize;
68    c.read_bytes(len)
69}
70
71/// encode the spans section as an `intpack` repack: a kind byte, the
72/// count, a deduped uri pool, then three bitpacked columns (uri index,
73/// byte_start, and byte_end - byte_start). lengths are stored as a
74/// wrapping difference so reconstruction is byte-exact for any input.
75pub fn encode_chunks_original_spans_intpack(spans: &[OriginalSpan]) -> Vec<u8> {
76    let mut pool: Vec<&str> = Vec::new();
77    let mut pos_of: HashMap<&str, usize> = HashMap::new();
78    let mut idx: Vec<u64> = Vec::with_capacity(spans.len());
79    let mut starts: Vec<u64> = Vec::with_capacity(spans.len());
80    let mut lens: Vec<u64> = Vec::with_capacity(spans.len());
81    for s in spans {
82        let uri = s.source_uri.as_str();
83        // first-appearance index, O(1) via the map. the numbering is
84        // identical to a linear pool scan (`pool.len()` at first sight), so
85        // the serialized pool + indices stay BYTE-IDENTICAL; this only avoids
86        // the O(n*distinct) scan that turns quadratic on many-uri corpora.
87        let pos = *pos_of.entry(uri).or_insert_with(|| {
88            let p = pool.len();
89            pool.push(uri);
90            p
91        });
92        idx.push(pos as u64);
93        starts.push(s.byte_start);
94        lens.push(s.byte_end.wrapping_sub(s.byte_start));
95    }
96    let mut out = Vec::new();
97    out.push(REPACK_KIND_SPANS);
98    out.extend_from_slice(&(spans.len() as u32).to_le_bytes());
99    out.extend_from_slice(&(pool.len() as u32).to_le_bytes());
100    for uri in &pool {
101        out.extend_from_slice(&(uri.len() as u32).to_le_bytes());
102        out.extend_from_slice(uri.as_bytes());
103    }
104    for col in [&idx, &starts, &lens] {
105        let blob = pack_u64s(col);
106        out.extend_from_slice(&(blob.len() as u32).to_le_bytes());
107        out.extend_from_slice(&blob);
108    }
109    out
110}
111
112/// reconstruct the canonical (raw-encoding) spans payload from the body
113/// of an `intpack` repack (the bytes after the kind byte). byte-identical
114/// to [`encode_chunks_original_spans`] so `content_hash` is preserved.
115pub fn decode_chunks_original_spans_intpack(rest: &[u8]) -> crate::Result<Vec<u8>> {
116    let mut c = Cursor::new(rest, SECTION_CHUNKS_ORIGINAL_SPANS);
117    let count = c.read_u32()? as usize;
118    let n_uris = c.read_u32()? as usize;
119    // bound the claim against the bytes before allocating: every pooled uri
120    // costs at least its 4-byte length prefix, so a hostile `n_uris` of 1.3G
121    // (a real cargo-fuzz finding: a 90-byte payload asked for 31 GB) is a
122    // typed error instead of an allocation abort.
123    if n_uris > (rest.len() - c.pos) / 4 {
124        return Err(c.malformed("spans intpack: uri pool count exceeds payload"));
125    }
126    let mut pool: Vec<String> = Vec::with_capacity(n_uris);
127    for _ in 0..n_uris {
128        pool.push(c.read_lp_str()?);
129    }
130    let idx = unpack_u64s(read_blob(&mut c)?)?;
131    let starts = unpack_u64s(read_blob(&mut c)?)?;
132    let lens = unpack_u64s(read_blob(&mut c)?)?;
133    if idx.len() != count || starts.len() != count || lens.len() != count {
134        return Err(c.malformed("spans intpack: column length mismatch"));
135    }
136    c.finish()?;
137    let mut buf = Vec::with_capacity(12 + count * 24);
138    write_prefix(&mut buf, count as u64);
139    for i in 0..count {
140        let uri = pool
141            .get(idx[i] as usize)
142            .ok_or_else(|| UrnaError::MalformedSectionPayload {
143                section_id: SECTION_CHUNKS_ORIGINAL_SPANS,
144                reason: "spans intpack: uri index out of range".into(),
145            })?;
146        write_lp_str(&mut buf, uri)?;
147        buf.extend_from_slice(&starts[i].to_le_bytes());
148        buf.extend_from_slice(&starts[i].wrapping_add(lens[i]).to_le_bytes());
149    }
150    Ok(buf)
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn roundtrip() {
159        let spans = vec![
160            OriginalSpan {
161                source_uri: "doc.txt".into(),
162                byte_start: 0,
163                byte_end: 10,
164            },
165            OriginalSpan {
166                source_uri: "doc.txt".into(),
167                byte_start: 10,
168                byte_end: 25,
169            },
170        ];
171        let bytes = encode_chunks_original_spans(&spans).unwrap();
172        let back = decode_chunks_original_spans(&bytes, 2).unwrap();
173        assert_eq!(spans, back);
174    }
175
176    fn corpus(n: u64) -> Vec<OriginalSpan> {
177        (0..n)
178            .map(|i| OriginalSpan {
179                source_uri: "corpus-next/v1.txt".into(), // shared -> deduped
180                byte_start: i * 500,
181                byte_end: i * 500 + 480,
182            })
183            .collect()
184    }
185
186    #[test]
187    fn intpack_decodes_byte_identical_to_raw() {
188        let spans = corpus(200);
189        let raw = encode_chunks_original_spans(&spans).unwrap();
190        let packed = encode_chunks_original_spans_intpack(&spans);
191        assert_eq!(packed[0], REPACK_KIND_SPANS);
192        assert!(packed.len() < raw.len(), "dedup+intpack must shrink spans");
193        let reconstructed = decode_chunks_original_spans_intpack(&packed[1..]).unwrap();
194        assert_eq!(reconstructed, raw, "spans repack must rebuild raw bytes");
195        assert_eq!(
196            decode_chunks_original_spans(&reconstructed, spans.len()).unwrap(),
197            spans
198        );
199    }
200
201    #[test]
202    fn intpack_handles_multiple_uris_and_empty() {
203        let spans = vec![
204            OriginalSpan {
205                source_uri: "a.txt".into(),
206                byte_start: 5,
207                byte_end: 9,
208            },
209            OriginalSpan {
210                source_uri: "b.txt".into(),
211                byte_start: 0,
212                byte_end: 100,
213            },
214            OriginalSpan {
215                source_uri: "a.txt".into(),
216                byte_start: 9,
217                byte_end: 12,
218            },
219        ];
220        let raw = encode_chunks_original_spans(&spans).unwrap();
221        let packed = encode_chunks_original_spans_intpack(&spans);
222        assert_eq!(
223            decode_chunks_original_spans_intpack(&packed[1..]).unwrap(),
224            raw
225        );
226
227        let empty = encode_chunks_original_spans_intpack(&[]);
228        assert_eq!(
229            decode_chunks_original_spans_intpack(&empty[1..]).unwrap(),
230            encode_chunks_original_spans(&[]).unwrap()
231        );
232    }
233
234    #[test]
235    fn intpack_truncated_errors_never_panic() {
236        let packed = encode_chunks_original_spans_intpack(&corpus(10));
237        for cut in 1..packed.len() {
238            let _ = decode_chunks_original_spans_intpack(&packed[1..cut]);
239        }
240    }
241}
urna_format/sections/spans.rs

urna_format/sections/
spans.rs
