vole_document/adapter/pdf/
channels.rs1use crate::error::{Error, ErrorClass, Result};
22use crate::limits::Limits;
23
24use super::lexer::lex;
25use super::span::SpanKind;
26
27pub const KIND_COUNT: usize = 12;
29
30pub const fn kind_id(k: SpanKind) -> u8 {
32 match k {
33 SpanKind::Whitespace => 0,
34 SpanKind::Comment => 1,
35 SpanKind::LiteralString => 2,
36 SpanKind::HexString => 3,
37 SpanKind::DictOpen => 4,
38 SpanKind::DictClose => 5,
39 SpanKind::ArrayOpen => 6,
40 SpanKind::ArrayClose => 7,
41 SpanKind::BraceOpen => 8,
42 SpanKind::BraceClose => 9,
43 SpanKind::Name => 10,
44 SpanKind::Regular => 11,
45 }
46}
47
48pub const fn kind_from_id(id: u8) -> Option<SpanKind> {
50 Some(match id {
51 0 => SpanKind::Whitespace,
52 1 => SpanKind::Comment,
53 2 => SpanKind::LiteralString,
54 3 => SpanKind::HexString,
55 4 => SpanKind::DictOpen,
56 5 => SpanKind::DictClose,
57 6 => SpanKind::ArrayOpen,
58 7 => SpanKind::ArrayClose,
59 8 => SpanKind::BraceOpen,
60 9 => SpanKind::BraceClose,
61 10 => SpanKind::Name,
62 11 => SpanKind::Regular,
63 _ => return None,
64 })
65}
66
67pub const fn role_id(k: SpanKind) -> u8 {
78 match k {
79 SpanKind::DictOpen
80 | SpanKind::DictClose
81 | SpanKind::ArrayOpen
82 | SpanKind::ArrayClose
83 | SpanKind::BraceOpen
84 | SpanKind::BraceClose
85 | SpanKind::Name => 0,
86 SpanKind::LiteralString
87 | SpanKind::HexString
88 | SpanKind::Comment
89 | SpanKind::Whitespace => 1,
90 SpanKind::Regular => 2,
91 }
92}
93
94#[derive(Debug, Clone, PartialEq, Eq, Default)]
100pub struct TokenChannelPlan {
101 pub kinds: Vec<u8>,
103 pub lengths: Vec<u32>,
105 pub payloads: Vec<Vec<u8>>,
108}
109
110impl TokenChannelPlan {
111 pub fn token_count(&self) -> usize {
113 self.kinds.len()
114 }
115
116 pub fn total_len(&self) -> u64 {
119 self.lengths
120 .iter()
121 .fold(0u64, |acc, &len| acc.saturating_add(u64::from(len)))
122 }
123}
124
125pub fn split(input: &[u8], limits: Limits) -> Result<Option<TokenChannelPlan>> {
132 split_with(input, limits, kind_id)
133}
134
135pub fn split_role(input: &[u8], limits: Limits) -> Result<Option<TokenChannelPlan>> {
146 split_with(input, limits, role_id)
147}
148
149fn split_with(
152 input: &[u8],
153 limits: Limits,
154 map: fn(SpanKind) -> u8,
155) -> Result<Option<TokenChannelPlan>> {
156 if input.len() as u64 > limits.max_output_bytes {
159 return Err(Error::resource_limit(format!(
160 "pdf channel payload {} exceeds max_output_bytes {}",
161 input.len(),
162 limits.max_output_bytes
163 )));
164 }
165
166 let lexed = match lex(input, limits) {
170 Ok(result) => result,
171 Err(e) if e.class() == ErrorClass::ResourceLimit => return Ok(None),
172 Err(e) => return Err(e),
173 };
174
175 let spans = &lexed.spans.spans;
176 let mut plan = TokenChannelPlan {
177 kinds: Vec::with_capacity(spans.len()),
178 lengths: Vec::with_capacity(spans.len()),
179 payloads: vec![Vec::new(); KIND_COUNT],
180 };
181
182 for span in spans {
183 let len = u32::try_from(span.len)
184 .map_err(|_| Error::resource_limit(format!("span length {} exceeds u32", span.len)))?;
185 let id = map(span.kind);
186 let start = span.start as usize;
187 let end = start + span.len as usize;
188 plan.kinds.push(id);
189 plan.lengths.push(len);
190 plan.payloads[id as usize].extend_from_slice(&input[start..end]);
191 }
192
193 Ok(Some(plan))
194}
195
196pub fn join(plan: &TokenChannelPlan, limits: Limits) -> Result<Vec<u8>> {
203 if plan.kinds.len() != plan.lengths.len() {
204 return Err(Error::invalid_pdf_structure(format!(
205 "kinds ({}), lengths ({}) are misaligned",
206 plan.kinds.len(),
207 plan.lengths.len()
208 )));
209 }
210 if plan.payloads.len() != KIND_COUNT {
211 return Err(Error::invalid_pdf_structure(format!(
212 "expected {KIND_COUNT} payload streams, found {}",
213 plan.payloads.len()
214 )));
215 }
216
217 let total = plan.total_len();
218 if total > limits.max_output_bytes {
219 return Err(Error::resource_limit(format!(
220 "plan total {total} exceeds max_output_bytes {}",
221 limits.max_output_bytes
222 )));
223 }
224 let capacity = usize::try_from(total)
225 .map_err(|_| Error::coverage_violation("plan total exceeds the address space"))?;
226
227 let mut cursors = [0usize; KIND_COUNT];
228 let mut out = Vec::with_capacity(capacity);
229
230 for (&id, &len) in plan.kinds.iter().zip(plan.lengths.iter()) {
231 let k = id as usize;
232 if k >= KIND_COUNT {
233 return Err(Error::invalid_pdf_structure(format!(
234 "kind id {id} is outside 0..{KIND_COUNT}"
235 )));
236 }
237 let payload = &plan.payloads[k];
238 let start = cursors[k];
239 let end = start
240 .checked_add(len as usize)
241 .ok_or_else(|| Error::coverage_violation("channel cursor overflow"))?;
242 if end > payload.len() {
243 return Err(Error::coverage_violation(format!(
244 "kind {id} needs bytes {start}..{end} but its payload holds {}",
245 payload.len()
246 )));
247 }
248 out.extend_from_slice(&payload[start..end]);
249 cursors[k] = end;
250 }
251
252 for ((id, &cursor), payload) in cursors.iter().enumerate().zip(plan.payloads.iter()) {
253 if cursor != payload.len() {
254 return Err(Error::coverage_violation(format!(
255 "kind {id} left {} payload bytes unconsumed",
256 payload.len() - cursor
257 )));
258 }
259 }
260
261 if out.len() as u64 != total {
262 return Err(Error::coverage_violation(format!(
263 "joined {} bytes but the plan declares {total}",
264 out.len()
265 )));
266 }
267
268 Ok(out)
269}
270
271#[cfg(test)]
272mod tests {
273 use super::*;
274 use crate::adapter::pdf::samples::sample_pdfs;
275
276 fn xorshift64(state: &mut u64) -> u64 {
277 let mut x = *state;
278 x ^= x << 13;
279 x ^= x >> 7;
280 x ^= x << 17;
281 *state = x;
282 x
283 }
284
285 #[test]
286 fn kind_map_is_a_bijection_onto_0_to_12() {
287 let all = [
288 SpanKind::Whitespace,
289 SpanKind::Comment,
290 SpanKind::LiteralString,
291 SpanKind::HexString,
292 SpanKind::DictOpen,
293 SpanKind::DictClose,
294 SpanKind::ArrayOpen,
295 SpanKind::ArrayClose,
296 SpanKind::BraceOpen,
297 SpanKind::BraceClose,
298 SpanKind::Name,
299 SpanKind::Regular,
300 ];
301 let mut seen = [false; KIND_COUNT];
302 for k in all {
303 let id = kind_id(k);
304 assert!((id as usize) < KIND_COUNT, "id {id} out of range");
305 assert!(!seen[id as usize], "duplicate id {id}");
306 seen[id as usize] = true;
307 assert_eq!(kind_from_id(id), Some(k));
308 }
309 assert!(seen.iter().all(|&s| s), "every id must be used");
310 assert_eq!(kind_from_id(KIND_COUNT as u8), None);
311 assert_eq!(kind_from_id(u8::MAX), None);
312 }
313
314 #[test]
315 fn split_then_join_is_identity() {
316 for (name, bytes) in sample_pdfs() {
317 let plan = split(&bytes, Limits::DEFAULT)
318 .unwrap()
319 .unwrap_or_else(|| panic!("{name} must split"));
320 let out = join(&plan, Limits::DEFAULT).unwrap();
321 assert_eq!(out, bytes, "{name} must rejoin byte-exactly");
322 }
323 }
324
325 #[test]
326 fn deterministic() {
327 for (name, bytes) in sample_pdfs() {
328 let a = split(&bytes, Limits::DEFAULT).unwrap();
329 let b = split(&bytes, Limits::DEFAULT).unwrap();
330 assert_eq!(a, b, "{name} split must be deterministic");
331 }
332 }
333
334 #[test]
335 fn payload_partition() {
336 for (name, bytes) in sample_pdfs() {
337 let plan = split(&bytes, Limits::DEFAULT).unwrap().unwrap();
338 assert_eq!(plan.kinds.len(), plan.lengths.len(), "{name} alignment");
339 assert_eq!(plan.kinds.len(), plan.token_count(), "{name} token_count");
340 assert!(
341 plan.kinds.iter().all(|&k| (k as usize) < KIND_COUNT),
342 "{name} kind ids in range"
343 );
344 let payload_sum: u64 = plan.payloads.iter().map(|p| p.len() as u64).sum();
345 assert_eq!(payload_sum, plan.total_len(), "{name} payload sum");
346 assert_eq!(plan.total_len(), bytes.len() as u64, "{name} total_len");
347 }
348 }
349
350 #[test]
351 fn random_roundtrip() {
352 let mut state: u64 = 0x243F_6A88_85A3_08D3;
353 let mut split_count = 0usize;
354 for i in 0..500 {
355 let len = (xorshift64(&mut state) % 400) as usize;
356 let mut buf = Vec::with_capacity(len);
357 for _ in 0..len {
358 buf.push((xorshift64(&mut state) & 0xFF) as u8);
359 }
360 if let Some(plan) = split(&buf, Limits::DEFAULT).unwrap() {
361 split_count += 1;
362 let out = join(&plan, Limits::DEFAULT).unwrap();
363 assert_eq!(out, buf, "case {i} must rejoin byte-exactly");
364 }
365 }
366 assert!(split_count > 0, "the court must actually exercise split");
367 }
368
369 #[test]
370 fn declines_when_too_many_tokens() {
371 let limits = Limits {
372 max_pdf_spans: 1,
373 ..Limits::DEFAULT
374 };
375 assert!(
377 split(b"a b", limits).unwrap().is_none(),
378 "a token count above max_pdf_spans must decline honestly"
379 );
380 }
381
382 #[test]
383 fn join_rejects_malformed_plan() {
384 let limits = Limits::DEFAULT;
385
386 let bad_kind = TokenChannelPlan {
388 kinds: vec![KIND_COUNT as u8],
389 lengths: vec![1],
390 payloads: vec![Vec::new(); KIND_COUNT],
391 };
392 assert_eq!(
393 join(&bad_kind, limits).unwrap_err().class(),
394 ErrorClass::InvalidPdfStructure
395 );
396
397 let mut payloads = vec![Vec::new(); KIND_COUNT];
400 payloads[0] = b"ab".to_vec();
401 let leftover = TokenChannelPlan {
402 kinds: vec![0],
403 lengths: vec![1],
404 payloads,
405 };
406 assert_eq!(
407 join(&leftover, limits).unwrap_err().class(),
408 ErrorClass::CoverageViolation
409 );
410
411 let overrun = TokenChannelPlan {
413 kinds: vec![0],
414 lengths: vec![5],
415 payloads: vec![Vec::new(); KIND_COUNT],
416 };
417 assert_eq!(
418 join(&overrun, limits).unwrap_err().class(),
419 ErrorClass::CoverageViolation
420 );
421
422 let misaligned = TokenChannelPlan {
424 kinds: vec![0, 0],
425 lengths: vec![1],
426 payloads: vec![Vec::new(); KIND_COUNT],
427 };
428 assert_eq!(
429 join(&misaligned, limits).unwrap_err().class(),
430 ErrorClass::InvalidPdfStructure
431 );
432 }
433}