1use heddle_format::compression::{
5 CompressionConfig, CompressionDictionary, compress, compress_with_dictionary, decompress,
6 decompress_with_dictionary, is_compressed,
7};
8
9use crate::{
10 object::{
11 Action, ActionId, ContentHash, State, TREE_DELTA_ANCHOR_INTERVAL, TREE_DELTA_MAX_OPS, Tree,
12 decode_tree_delta, decode_tree_delta_header, encode_tree_delta, is_canonical_tree,
13 is_delta_tree, is_lean_tree, tree_delta,
14 },
15 store::{HeddleError, Result},
16};
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21pub struct TreeLineage {
22 pub anchor: ContentHash,
23 pub depth: u8,
24}
25
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum TreeEncodingKind {
28 Lean,
29 Delta {
30 anchor: ContentHash,
31 depth: u8,
32 op_count: usize,
33 },
34}
35
36#[derive(Clone, Debug, PartialEq, Eq)]
37pub struct EncodedTree {
38 pub hash: ContentHash,
39 pub data: Vec<u8>,
40 pub kind: TreeEncodingKind,
41}
42
43pub struct TreeDeltaBase<'a> {
45 pub anchor_id: ContentHash,
46 pub anchor: &'a Tree,
47 pub parent_depth: u8,
49}
50
51pub fn encode_blob_content(content: &[u8], config: &CompressionConfig) -> Result<Vec<u8>> {
52 Ok(compress(content, config)?.unwrap_or_else(|| content.to_vec()))
53}
54
55pub fn decode_blob_content(data: &[u8]) -> Result<Vec<u8>> {
56 if is_compressed(data) {
57 Ok(decompress(data)?)
58 } else {
59 Ok(data.to_vec())
60 }
61}
62
63pub fn encode_tree(tree: &Tree, _config: &CompressionConfig) -> Result<(ContentHash, Vec<u8>)> {
64 let encoded = encode_tree_hot(tree, None)?;
65 Ok((encoded.hash, encoded.data))
66}
67
68pub fn encode_tree_hot(tree: &Tree, base: Option<TreeDeltaBase<'_>>) -> Result<EncodedTree> {
71 let hash = tree.hash();
72 let lean = tree.encode_lean()?;
73 let Some(base) = base else {
74 return Ok(EncodedTree {
75 hash,
76 data: lean,
77 kind: TreeEncodingKind::Lean,
78 });
79 };
80 if hash == base.anchor_id {
81 return Ok(EncodedTree {
82 hash,
83 data: lean,
84 kind: TreeEncodingKind::Lean,
85 });
86 }
87 let Some(depth) = base.parent_depth.checked_add(1) else {
88 return Ok(EncodedTree {
89 hash,
90 data: lean,
91 kind: TreeEncodingKind::Lean,
92 });
93 };
94 if depth >= TREE_DELTA_ANCHOR_INTERVAL {
95 return Ok(EncodedTree {
96 hash,
97 data: lean,
98 kind: TreeEncodingKind::Lean,
99 });
100 }
101 let ops = tree_delta(base.anchor, tree);
102 if ops.len() > TREE_DELTA_MAX_OPS {
103 return Ok(EncodedTree {
104 hash,
105 data: lean,
106 kind: TreeEncodingKind::Lean,
107 });
108 }
109 let delta = encode_tree_delta(base.anchor_id, base.anchor, tree, &ops)?;
110 let header = decode_tree_delta_header(&delta)?;
111 let porch_is_bounded = header.first_base_count <= 1 && header.hundred_base_count <= 100;
112 if !porch_is_bounded || delta.len() >= lean.len() {
113 return Ok(EncodedTree {
114 hash,
115 data: lean,
116 kind: TreeEncodingKind::Lean,
117 });
118 }
119 Ok(EncodedTree {
120 hash,
121 data: delta,
122 kind: TreeEncodingKind::Delta {
123 anchor: base.anchor_id,
124 depth,
125 op_count: ops.len(),
126 },
127 })
128}
129
130pub fn encode_tree_at_rest(tree: &Tree, config: &CompressionConfig) -> Result<Vec<u8>> {
132 if config.enabled && tree.len() >= crate::object::TREE_BLOCK_MIN_ENTRIES {
133 Ok(tree.encode_canonical_blocked(config.level, config.min_size)?)
134 } else {
135 Ok(tree.encode_canonical()?)
136 }
137}
138
139pub fn decode_tree(data: &[u8]) -> Result<Tree> {
140 let decoded = decode_tree_body(data)?;
141 decode_tree_serialized(&decoded)
142}
143
144pub fn decode_tree_serialized(data: &[u8]) -> Result<Tree> {
145 if !is_canonical_tree(data) {
146 return Err(HeddleError::InvalidObject(
147 "HLR1/HDC1 tree decoding requires the external object key".to_string(),
148 ));
149 }
150 Tree::decode_canonical(data).map_err(HeddleError::from)
151}
152
153pub fn decode_tree_with_key(
157 data: &[u8],
158 expected: ContentHash,
159 anchor: Option<&Tree>,
160) -> Result<Tree> {
161 let decoded = decode_tree_body(data)?;
162 decode_tree_serialized_with_key(&decoded, expected, anchor)
163}
164
165pub fn decode_tree_serialized_with_key(
166 data: &[u8],
167 expected: ContentHash,
168 anchor: Option<&Tree>,
169) -> Result<Tree> {
170 let tree = if is_lean_tree(data) {
171 Tree::decode_lean(data, expected)?
172 } else if is_delta_tree(data) {
173 let header = decode_tree_delta_header(data)?;
174 if header.anchor == expected {
175 return Err(HeddleError::InvalidObject(
176 "HDC1 result id must differ from its anchor id".to_string(),
177 ));
178 }
179 let anchor = anchor.ok_or_else(|| {
180 HeddleError::InvalidObject("HDC1 tree is missing its materialized anchor".to_string())
181 })?;
182 decode_tree_delta(data, anchor, expected)?
183 } else if is_canonical_tree(data) {
184 Tree::decode_canonical(data)?
185 } else {
186 return Err(HeddleError::InvalidObject(
187 "unsupported tree storage body".to_string(),
188 ));
189 };
190 let found = tree.hash();
191 if found != expected {
192 return Err(HeddleError::Corruption { expected, found });
193 }
194 Ok(tree)
195}
196
197pub fn decode_tree_body(data: &[u8]) -> Result<Vec<u8>> {
201 Ok(decompress_with_dictionary(data)?)
202}
203
204pub fn encode_state(state: &State, config: &CompressionConfig) -> Result<Vec<u8>> {
205 let serialized = rmp_serde::to_vec(state)?;
206 Ok(
207 compress_with_dictionary(&serialized, config, CompressionDictionary::TreeStateV1)?
208 .unwrap_or(serialized),
209 )
210}
211
212pub fn decode_state(data: &[u8]) -> Result<State> {
213 let decoded = decompress_with_dictionary(data)?;
214 let mut state: State = rmp_serde::from_slice(&decoded)?;
215 state.state_id = state.id();
216 Ok(state)
217}
218
219pub fn encode_action(
220 action: &mut Action,
221 config: &CompressionConfig,
222) -> Result<(ActionId, Vec<u8>)> {
223 let id = action.id();
224 let serialized = rmp_serde::to_vec(action)?;
225 let data = compress(&serialized, config)?.unwrap_or(serialized);
226 Ok((id, data))
227}
228
229pub fn decode_action(data: &[u8]) -> Result<Action> {
230 let decoded = decode_body(data)?;
231 Ok(rmp_serde::from_slice(&decoded)?)
232}
233
234fn decode_body(data: &[u8]) -> Result<Vec<u8>> {
235 if is_compressed(data) {
236 Ok(decompress(data)?)
237 } else {
238 Ok(data.to_vec())
239 }
240}
241
242#[cfg(test)]
243mod tests {
244 use super::*;
245 use crate::object::{Attribution, Operation, Principal, StateId, TreeEntry};
246
247 #[test]
248 fn encode_decode_blob_content_matches_old_recipe() {
249 let content = b"codec blob content ".repeat(64);
250 for config in compression_configs() {
251 let expected = old_encode_raw(&content, &config).unwrap();
252 let encoded = encode_blob_content(&content, &config).unwrap();
253 assert_eq!(encoded, expected);
254 assert_eq!(decode_blob_content(&encoded).unwrap(), content);
255 }
256 }
257
258 #[test]
259 fn encode_decode_tree() {
260 let blob_hash = ContentHash::compute(b"codec-tree-blob");
261 let tree = Tree::from_entries(vec![TreeEntry::file("file.txt", blob_hash, false).unwrap()]);
262 for config in compression_configs() {
263 let (hash, encoded) = encode_tree(&tree, &config).unwrap();
264 assert_eq!(hash, tree.hash());
265 assert!(crate::object::is_lean_tree(&encoded));
266 assert_eq!(decode_tree_with_key(&encoded, hash, None).unwrap(), tree);
267 }
268 }
269
270 #[test]
271 fn lean_and_delta_round_trip_against_external_keys() {
272 let anchor = tree_fixture(240, None);
273 let current = tree_fixture(240, Some((117, b"changed")));
274 let lean = encode_tree_hot(&anchor, None).unwrap();
275 assert_eq!(lean.kind, TreeEncodingKind::Lean);
276 assert_eq!(
277 decode_tree_with_key(&lean.data, anchor.hash(), None).unwrap(),
278 anchor
279 );
280
281 let delta = encode_tree_hot(
282 ¤t,
283 Some(TreeDeltaBase {
284 anchor_id: anchor.hash(),
285 anchor: &anchor,
286 parent_depth: 0,
287 }),
288 )
289 .unwrap();
290 assert!(matches!(
291 delta.kind,
292 TreeEncodingKind::Delta {
293 anchor: _,
294 depth: 1,
295 op_count: 1
296 }
297 ));
298 assert!(crate::object::is_delta_tree(&delta.data));
299 assert_eq!(
300 decode_tree_with_key(&delta.data, current.hash(), Some(&anchor)).unwrap(),
301 current
302 );
303 }
304
305 #[test]
306 fn result_equal_to_anchor_is_materialized_instead_of_delta_encoded() {
307 let anchor = tree_fixture(240, None);
308 let encoded = encode_tree_hot(
309 &anchor,
310 Some(TreeDeltaBase {
311 anchor_id: anchor.hash(),
312 anchor: &anchor,
313 parent_depth: 1,
314 }),
315 )
316 .unwrap();
317
318 assert_eq!(encoded.kind, TreeEncodingKind::Lean);
319 assert!(crate::object::is_lean_tree(&encoded.data));
320 }
321
322 #[test]
323 fn delta_refreshes_anchor_after_127_descendants() {
324 let anchor = tree_fixture(240, None);
325 let current = tree_fixture(240, Some((117, b"changed")));
326
327 let last_descendant = encode_tree_hot(
328 ¤t,
329 Some(TreeDeltaBase {
330 anchor_id: anchor.hash(),
331 anchor: &anchor,
332 parent_depth: TREE_DELTA_ANCHOR_INTERVAL - 2,
333 }),
334 )
335 .unwrap();
336 assert!(matches!(
337 last_descendant.kind,
338 TreeEncodingKind::Delta { depth: 127, .. }
339 ));
340
341 let refreshed = encode_tree_hot(
342 ¤t,
343 Some(TreeDeltaBase {
344 anchor_id: anchor.hash(),
345 anchor: &anchor,
346 parent_depth: TREE_DELTA_ANCHOR_INTERVAL - 1,
347 }),
348 )
349 .unwrap();
350 assert_eq!(refreshed.kind, TreeEncodingKind::Lean);
351 assert!(crate::object::is_lean_tree(&refreshed.data));
352 }
353
354 #[test]
355 fn delta_over_512_operations_refreshes_the_anchor() {
356 let anchor = tree_fixture(600, None);
357 let current = Tree::from_entries(
358 anchor
359 .entries()
360 .iter()
361 .enumerate()
362 .map(|(index, entry)| {
363 TreeEntry::file(
364 entry.name(),
365 ContentHash::compute(format!("changed-{index}").as_bytes()),
366 false,
367 )
368 .unwrap()
369 })
370 .collect(),
371 );
372 let ops = tree_delta(&anchor, ¤t);
373 assert_eq!(ops.len(), 600);
374 assert!(encode_tree_delta(anchor.hash(), &anchor, ¤t, &ops).is_err());
375 let encoded = encode_tree_hot(
376 ¤t,
377 Some(TreeDeltaBase {
378 anchor_id: anchor.hash(),
379 anchor: &anchor,
380 parent_depth: 0,
381 }),
382 )
383 .unwrap();
384 assert_eq!(encoded.kind, TreeEncodingKind::Lean);
385 }
386
387 #[test]
388 fn every_tree_form_validates_the_external_key() {
389 let anchor = tree_fixture(240, None);
390 let current = tree_fixture(240, Some((117, b"changed")));
391 let wrong = ContentHash::compute(b"wrong-tree-key");
392 let lean = anchor.encode_lean().unwrap();
393 assert!(decode_tree_with_key(&lean, wrong, None).is_err());
394
395 let ops = tree_delta(&anchor, ¤t);
396 let delta = encode_tree_delta(anchor.hash(), &anchor, ¤t, &ops).unwrap();
397 assert!(decode_tree_with_key(&delta, wrong, Some(&anchor)).is_err());
398
399 let raw = current.encode_canonical().unwrap();
400 assert!(decode_tree_with_key(&raw, wrong, None).is_err());
401 }
402
403 #[test]
404 #[cfg(feature = "zstd")]
405 fn tree_and_state_use_versioned_dictionary_frames() {
406 let tree = Tree::from_entries(
407 (0..24)
408 .map(|index| {
409 TreeEntry::file(
410 format!("module_{index:02}.rs"),
411 ContentHash::compute(format!("blob-{index}").as_bytes()),
412 false,
413 )
414 .unwrap()
415 })
416 .collect(),
417 );
418 let state = State::new(
419 tree.hash(),
420 vec![StateId::from_bytes([7; 32])],
421 sample_attribution(),
422 )
423 .with_intent("dictionary frame verification ".repeat(32));
424
425 let encoded_tree = encode_tree_at_rest(&tree, &CompressionConfig::default()).unwrap();
426 let encoded_state = encode_state(&state, &CompressionConfig::default()).unwrap();
427
428 assert!(
429 crate::object::is_canonical_tree(&encoded_tree),
430 "at-rest trees remain versioned HTR4 so resume can seek"
431 );
432 assert_eq!(&encoded_state[9..13], &1_u32.to_be_bytes());
433 }
434
435 #[test]
436 #[cfg(feature = "zstd")]
437 fn store_decoder_reads_raw_v4_and_blocked_v5() {
438 let tree = tree_fixture(600, None);
439 let raw = tree.encode_canonical().unwrap();
440 let blocked = encode_tree_at_rest(
441 &tree,
442 &CompressionConfig {
443 enabled: true,
444 level: 3,
445 min_size: 0,
446 max_delta_size: CompressionConfig::default().max_delta_size,
447 },
448 )
449 .unwrap();
450 assert_eq!(raw[4], crate::object::TREE_ENCODING_VERSION);
451 assert_eq!(blocked[4], crate::object::TREE_BLOCK_ENCODING_VERSION);
452 assert_eq!(decode_tree_with_key(&raw, tree.hash(), None).unwrap(), tree);
453 assert_eq!(
454 decode_tree_with_key(&blocked, tree.hash(), None).unwrap(),
455 tree
456 );
457 }
458
459 #[test]
460 #[cfg(feature = "zstd")]
461 fn tree_state_dictionary_corpus_roundtrips_byte_identically() {
462 let config = CompressionConfig::default();
463
464 for revision in 0..64 {
465 let tree = Tree::from_entries(
466 (0..32)
467 .map(|entry| {
468 TreeEntry::file(
469 format!("module_{entry:02}.rs"),
470 ContentHash::compute(
471 format!("revision-{revision}-blob-{entry}").as_bytes(),
472 ),
473 entry % 11 == 0,
474 )
475 .unwrap()
476 })
477 .collect(),
478 );
479 let encoded_tree = encode_tree_at_rest(&tree, &config).unwrap();
480 assert_eq!(Tree::decode_canonical(&encoded_tree).unwrap(), tree);
481
482 let state = State::new(
483 tree.hash(),
484 vec![StateId::from_bytes([revision; 32])],
485 sample_attribution(),
486 )
487 .with_intent(format!(
488 "Update the representative tree/state corpus at revision {revision}. {}",
489 "Preserve byte-identical object bodies. ".repeat(12)
490 ));
491 let serialized_state = rmp_serde::to_vec(&state).unwrap();
492 let encoded_state = encode_state(&state, &config).unwrap();
493 assert_eq!(
494 decompress_with_dictionary(&encoded_state).unwrap(),
495 serialized_state
496 );
497 }
498 }
499
500 #[test]
501 fn encode_decode_state() {
502 let attribution = sample_attribution();
503 let state = State::new(ContentHash::compute(b"codec-tree"), vec![], attribution)
504 .with_intent("codec state");
505 for config in compression_configs() {
506 let encoded = encode_state(&state, &config).unwrap();
507 assert_eq!(decode_state(&encoded).unwrap(), state);
508 }
509 }
510
511 #[test]
512 fn encode_decode_action_matches_old_recipe() {
513 let attribution = sample_attribution();
514 for config in compression_configs() {
515 let mut action = Action::new(
516 None,
517 StateId::from_bytes([1; 32]),
518 Operation::Snapshot,
519 "codec action",
520 attribution.clone(),
521 );
522 let id = action.id();
523 let serialized = rmp_serde::to_vec(&action).unwrap();
524 let expected = old_encode_raw(&serialized, &config).unwrap();
525
526 let (encoded_id, encoded) = encode_action(&mut action, &config).unwrap();
527 assert_eq!(encoded_id, id);
528 assert_eq!(encoded, expected);
529
530 let decoded = decode_action(&encoded).unwrap();
531 assert_eq!(decoded.compute_id(), id);
532 assert_eq!(decoded.from_state, action.from_state);
533 assert_eq!(decoded.to_state, action.to_state);
534 assert_eq!(decoded.operation, action.operation);
535 assert_eq!(decoded.description, action.description);
536 assert_eq!(decoded.semantic_changes, action.semantic_changes);
537 assert_eq!(decoded.attribution, action.attribution);
538 assert_eq!(decoded.timestamp, action.timestamp);
539 }
540 }
541
542 fn old_encode_raw(data: &[u8], config: &CompressionConfig) -> Result<Vec<u8>> {
543 Ok(compress(data, config)?.unwrap_or_else(|| data.to_vec()))
544 }
545
546 fn tree_fixture(entries: usize, changed: Option<(usize, &[u8])>) -> Tree {
547 Tree::from_entries(
548 (0..entries)
549 .map(|index| {
550 let payload = changed
551 .filter(|(changed_index, _)| *changed_index == index)
552 .map_or_else(
553 || format!("blob-{index}").into_bytes(),
554 |(_, payload)| payload.to_vec(),
555 );
556 TreeEntry::file(
557 format!("module_{index:04}.rs"),
558 ContentHash::compute(&payload),
559 false,
560 )
561 .unwrap()
562 })
563 .collect(),
564 )
565 }
566
567 fn compression_configs() -> Vec<CompressionConfig> {
568 #[cfg(feature = "zstd")]
569 {
570 vec![
571 CompressionConfig::default(),
572 CompressionConfig::disabled(),
573 CompressionConfig {
574 enabled: true,
575 level: 9,
576 min_size: 0,
577 max_delta_size: CompressionConfig::default().max_delta_size,
578 },
579 ]
580 }
581 #[cfg(not(feature = "zstd"))]
582 {
583 vec![CompressionConfig::default(), CompressionConfig::disabled()]
584 }
585 }
586
587 fn sample_attribution() -> Attribution {
588 Attribution::human(Principal::new("Codec Test", "codec@example.com"))
589 }
590}