Skip to main content

urna_format/
chunk.rs

1//! Canonical chunk types and chunk_id derivation.
2//!
3//! `chunk_id` is a deterministic SHA-256 over the canonical text plus the
4//! original location and the chunker version. The exact preimage is fixed
5//! by spec - anything else would let two runs of the same chunker produce
6//! different IDs for the same content, which breaks reproducibility.
7
8use crate::error::UrnaError;
9use sha2::{Digest, Sha256};
10
11/// One chunk's worth of input that the writer needs to produce a v1 file.
12///
13/// `embedding` must be a normalized f32 vector of the embedding dimension
14/// declared in the manifest. The writer validates dimensions and rejects
15/// NaN/Inf - callers do not get to pass garbage in.
16#[derive(Clone, Debug, PartialEq)]
17pub struct ChunkInput {
18    pub canonical_text: String,
19    pub source_uri: String,
20    pub byte_start: u64,
21    pub byte_end: u64,
22    pub embedding: Vec<f32>,
23}
24
25/// Compute the canonical chunk_id for the given inputs.
26///
27/// Preimage layout (UTF-8, no separators ambiguity - all length-prefixed):
28///
29/// ```text
30///   "urna:chunk_id:v1\n"
31///   u32 LE   len(canonical_text)
32///   bytes    canonical_text
33///   u32 LE   len(source_uri)
34///   bytes    source_uri
35///   u64 LE   byte_start
36///   u64 LE   byte_end
37///   u32 LE   len(chunker_version)
38///   bytes    chunker_version
39/// ```
40///
41/// Output: `sha256:<64 hex chars>`.
42pub fn chunk_id(
43    canonical_text: &str,
44    source_uri: &str,
45    byte_start: u64,
46    byte_end: u64,
47    chunker_version: &str,
48) -> String {
49    const DOMAIN: &[u8] = b"urna:chunk_id:v1\n";
50    let mut h = Sha256::new();
51    h.update(DOMAIN);
52    write_lp(&mut h, canonical_text.as_bytes());
53    write_lp(&mut h, source_uri.as_bytes());
54    h.update(byte_start.to_le_bytes());
55    h.update(byte_end.to_le_bytes());
56    write_lp(&mut h, chunker_version.as_bytes());
57    let digest = h.finalize();
58    format!("sha256:{}", hex::encode(digest))
59}
60
61fn write_lp(h: &mut Sha256, bytes: &[u8]) {
62    let len = bytes.len() as u32;
63    h.update(len.to_le_bytes());
64    h.update(bytes);
65}
66
67/// Validate a `ChunkInput` against the embedding dimension declared in the
68/// manifest. Returns the validated reference for ergonomics.
69pub fn validate_chunk(c: &ChunkInput, embedding_dim: usize) -> crate::Result<()> {
70    if c.embedding.len() != embedding_dim {
71        return Err(UrnaError::DimensionMismatch {
72            expected: embedding_dim,
73            got: c.embedding.len(),
74        });
75    }
76    for v in &c.embedding {
77        if v.is_nan() || v.is_infinite() {
78            return Err(UrnaError::InvalidEmbeddingValue);
79        }
80    }
81    if c.byte_end < c.byte_start {
82        return Err(UrnaError::InvalidInput(format!(
83            "byte_end ({}) < byte_start ({})",
84            c.byte_end, c.byte_start
85        )));
86    }
87    Ok(())
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn chunk_id_is_deterministic() {
96        let a = chunk_id("hello", "doc.txt", 0, 5, "v1");
97        let b = chunk_id("hello", "doc.txt", 0, 5, "v1");
98        assert_eq!(a, b);
99        assert!(a.starts_with("sha256:"));
100        assert_eq!(a.len(), "sha256:".len() + 64);
101    }
102
103    #[test]
104    fn chunk_id_changes_with_inputs() {
105        let base = chunk_id("hello", "doc.txt", 0, 5, "v1");
106        assert_ne!(base, chunk_id("HELLO", "doc.txt", 0, 5, "v1"));
107        assert_ne!(base, chunk_id("hello", "other.txt", 0, 5, "v1"));
108        assert_ne!(base, chunk_id("hello", "doc.txt", 1, 5, "v1"));
109        assert_ne!(base, chunk_id("hello", "doc.txt", 0, 6, "v1"));
110        assert_ne!(base, chunk_id("hello", "doc.txt", 0, 5, "v2"));
111    }
112
113    #[test]
114    fn validate_chunk_dim_mismatch() {
115        let c = ChunkInput {
116            canonical_text: "x".into(),
117            source_uri: "y".into(),
118            byte_start: 0,
119            byte_end: 1,
120            embedding: vec![1.0, 0.0],
121        };
122        assert!(validate_chunk(&c, 4).is_err());
123    }
124
125    #[test]
126    fn validate_chunk_nan() {
127        let c = ChunkInput {
128            canonical_text: "x".into(),
129            source_uri: "y".into(),
130            byte_start: 0,
131            byte_end: 1,
132            embedding: vec![f32::NAN, 0.0, 0.0, 0.0],
133        };
134        assert!(matches!(
135            validate_chunk(&c, 4),
136            Err(UrnaError::InvalidEmbeddingValue)
137        ));
138    }
139
140    #[test]
141    fn validate_chunk_bad_span() {
142        let c = ChunkInput {
143            canonical_text: "x".into(),
144            source_uri: "y".into(),
145            byte_start: 10,
146            byte_end: 5,
147            embedding: vec![1.0, 0.0, 0.0, 0.0],
148        };
149        assert!(matches!(
150            validate_chunk(&c, 4),
151            Err(UrnaError::InvalidInput(_))
152        ));
153    }
154}