1use std::{
5 collections::BTreeMap,
6 path::{Path, PathBuf},
7};
8
9use objects::{
10 object::{
11 Annotation, AnnotationScope, Blob, ContentHash, ContextBlob, ContextTarget, EntryType,
12 State, Tree, TreeEntry,
13 },
14 store::ObjectStore,
15};
16
17use super::{HeddleError, Repository, Result};
18
19#[derive(Clone, Debug, PartialEq, Eq)]
20pub struct ContextEntry {
21 pub target: ContextTarget,
22 pub blob: ContextBlob,
23}
24
25impl Repository {
26 pub fn get_context_blob(
28 &self,
29 context_root: &ContentHash,
30 target: &ContextTarget,
31 ) -> Result<Option<ContextBlob>> {
32 let Some(blob_hash) = self.lookup_context_leaf_for_target(context_root, target)? else {
33 return Ok(None);
34 };
35 let Some(blob) = self.store.get_blob(&blob_hash)? else {
36 return Ok(None);
37 };
38 ContextBlob::decode(blob.content())
39 .map(Some)
40 .map_err(|e| HeddleError::InvalidObject(format!("invalid context blob: {e}")))
41 }
42
43 pub fn set_context_blob(
47 &self,
48 context_root: Option<&ContentHash>,
49 target: &ContextTarget,
50 blob: &ContextBlob,
51 ) -> Result<ContentHash> {
52 let bytes = blob
53 .encode()
54 .map_err(|e| HeddleError::InvalidObject(format!("encode context: {e}")))?;
55 let blob_hash = self.store.put_blob(&Blob::new(bytes))?;
56
57 let current_tree = match context_root {
58 Some(root) => self.require_tree(root)?,
59 None => Tree::new(),
60 };
61
62 self.insert_leaf_at_path(¤t_tree, &target.storage_path(), blob_hash)
63 }
64
65 pub fn remove_context_at_target(
69 &self,
70 context_root: &ContentHash,
71 target: &ContextTarget,
72 scope: Option<&AnnotationScope>,
73 ) -> Result<Option<ContentHash>> {
74 if let Some(scope) = scope {
75 if let Some(mut blob) = self.get_context_blob(context_root, target)? {
76 blob.annotations.retain(|a| !a.scope.matches(scope));
77 if blob.annotations.is_empty() {
78 return self.remove_context_target(context_root, target);
79 }
80 let new_root = self.set_context_blob(Some(context_root), target, &blob)?;
81 return Ok(Some(new_root));
82 }
83 return Ok(Some(*context_root));
84 }
85
86 self.remove_context_target(context_root, target)
87 }
88
89 pub fn remove_context_target(
90 &self,
91 context_root: &ContentHash,
92 target: &ContextTarget,
93 ) -> Result<Option<ContentHash>> {
94 self.remove_leaf_at_path(context_root, &target.storage_path())
95 }
96
97 pub fn list_context_entries(
99 &self,
100 context_root: &ContentHash,
101 prefix: Option<&Path>,
102 ) -> Result<Vec<ContextEntry>> {
103 let tree = match self.store.get_tree(context_root)? {
104 Some(t) => t,
105 None => return Ok(Vec::new()),
106 };
107 let mut results = BTreeMap::new();
108 self.walk_context_tree(
109 &tree,
110 &PathBuf::new(),
111 prefix,
112 &mut results,
113 ContextWalkMode::CanonicalOnly,
114 )?;
115 Ok(results
116 .into_iter()
117 .map(|(_, (target, blob))| ContextEntry { target, blob })
118 .collect())
119 }
120
121 pub fn find_annotation(
122 &self,
123 context_root: &ContentHash,
124 annotation_id: &str,
125 ) -> Result<Option<(ContextTarget, ContextBlob, usize)>> {
126 for entry in self.list_context_entries(context_root, None)? {
127 if let Some(index) = entry
128 .blob
129 .annotations
130 .iter()
131 .position(|annotation| annotation.annotation_id == annotation_id)
132 {
133 return Ok(Some((entry.target, entry.blob, index)));
134 }
135 }
136 Ok(None)
137 }
138
139 fn lookup_context_leaf_for_target(
142 &self,
143 root: &ContentHash,
144 target: &ContextTarget,
145 ) -> Result<Option<ContentHash>> {
146 self.lookup_context_leaf(root, &target.storage_path())
147 }
148
149 fn lookup_context_leaf(&self, root: &ContentHash, path: &Path) -> Result<Option<ContentHash>> {
150 let Some((name, rest)) = split_path(path) else {
151 return Ok(None);
152 };
153 let Some(tree) = self.store.get_tree(root)? else {
154 return Ok(None);
155 };
156 let Some(entry) = tree.get(name) else {
157 return Ok(None);
158 };
159 if rest.as_os_str().is_empty() {
160 return Ok(entry.blob_hash());
161 }
162 if !entry.is_tree() {
163 return Ok(None);
164 }
165 let Some(tree_hash) = entry.tree_hash() else {
166 return Ok(None);
167 };
168 self.lookup_context_leaf(&tree_hash, rest)
169 }
170
171 fn insert_leaf_at_path(
172 &self,
173 tree: &Tree,
174 path: &Path,
175 blob_hash: ContentHash,
176 ) -> Result<ContentHash> {
177 let Some((name, rest)) = split_path(path) else {
178 return Err(HeddleError::InvalidObject("empty path".to_string()));
179 };
180
181 let mut new_tree = tree.clone();
182
183 if rest.as_os_str().is_empty() {
184 new_tree.insert(TreeEntry::file(name, blob_hash, false)?);
185 } else {
186 let subtree = tree
187 .get(name)
188 .filter(|e| e.is_tree())
189 .and_then(|e| e.tree_hash())
190 .and_then(|hash| self.store.get_tree(&hash).ok().flatten())
191 .unwrap_or_default();
192
193 let sub_hash = self.insert_leaf_at_path(&subtree, rest, blob_hash)?;
194 new_tree.insert(TreeEntry::directory(name, sub_hash)?);
195 }
196
197 self.store.put_tree(&new_tree)
198 }
199
200 fn remove_leaf_at_path(&self, root: &ContentHash, path: &Path) -> Result<Option<ContentHash>> {
201 let Some(tree) = self.store.get_tree(root)? else {
202 return Ok(None);
203 };
204 let Some((name, rest)) = split_path(path) else {
205 return Ok(None);
206 };
207
208 let mut new_tree = tree.clone();
209
210 if rest.as_os_str().is_empty() {
211 new_tree.remove(name);
212 } else {
213 let Some(entry) = tree.get(name) else {
214 return Ok(Some(*root));
215 };
216 if !entry.is_tree() {
217 return Ok(Some(*root));
218 }
219 let Some(tree_hash) = entry.tree_hash() else {
220 return Ok(Some(*root));
221 };
222 match self.remove_leaf_at_path(&tree_hash, rest)? {
223 Some(sub_hash) => {
224 new_tree.insert(TreeEntry::directory(name, sub_hash)?);
225 }
226 None => {
227 new_tree.remove(name);
228 }
229 }
230 }
231
232 if new_tree.is_empty() {
233 Ok(None)
234 } else {
235 Ok(Some(self.store.put_tree(&new_tree)?))
236 }
237 }
238
239 fn walk_context_tree(
240 &self,
241 tree: &Tree,
242 current_path: &Path,
243 prefix: Option<&Path>,
244 results: &mut BTreeMap<String, (ContextTarget, ContextBlob)>,
245 mode: ContextWalkMode,
246 ) -> Result<()> {
247 for entry in tree.entries() {
248 let entry_path = current_path.join(entry.name());
249 match entry.entry_type() {
250 EntryType::Tree => {
251 if let Some(prefix) = prefix
252 && !prefix.starts_with(&entry_path)
253 && !entry_path.starts_with(prefix)
254 && !entry_path.starts_with("__files")
255 && !entry_path.starts_with("__states")
256 {
257 continue;
258 }
259 if let Some(tree_hash) = entry.tree_hash()
260 && let Some(subtree) = self.store.get_tree(&tree_hash)?
261 {
262 self.walk_context_tree(&subtree, &entry_path, prefix, results, mode)?;
263 }
264 }
265 EntryType::Blob => {
266 let Some(target) = context_target_from_entry_path(&entry_path, mode) else {
267 continue;
268 };
269 if let Some(prefix) = prefix
270 && let Some(path) = target.path()
271 && !Path::new(path).starts_with(prefix)
272 {
273 continue;
274 }
275 if let Some(blob_hash) = entry.blob_hash()
276 && let Some(blob) = self.store.get_blob(&blob_hash)?
277 && let Ok(context) = ContextBlob::decode(blob.content())
278 {
279 results.insert(context_entry_key(&target), (target, context));
280 }
281 }
282 EntryType::Symlink | EntryType::Gitlink | EntryType::Spoollink => {}
283 }
284 }
285 Ok(())
286 }
287
288 pub fn inherit_parent_context(&self, parent: &State) -> Result<Option<ContentHash>> {
290 Ok(self
291 .latest_state_attachment(&parent.id(), crate::StateAttachmentKind::Context)?
292 .and_then(|attachment| match attachment.body {
293 objects::object::StateAttachmentBody::Context(hash) => Some(hash),
294 _ => None,
295 }))
296 }
297
298 pub fn union_parent_contexts(&self, parents: &[&State]) -> Result<Option<ContentHash>> {
312 let mut roots = Vec::new();
314 for parent in parents {
315 if let Some(root) = self.inherit_parent_context(parent)? {
316 roots.push(root);
317 }
318 }
319 if roots.is_empty() {
320 return Ok(None);
321 }
322 if roots.len() == 1 {
323 return Ok(Some(roots.pop().expect("len == 1")));
324 }
325 if roots.iter().all(|r| *r == roots[0]) {
326 return Ok(Some(roots[0]));
328 }
329
330 let mut merged: BTreeMap<String, (ContextTarget, ContextBlob)> = BTreeMap::new();
334 for parent_root in &roots {
335 for entry in self.list_context_entries(parent_root, None)? {
336 let key = context_entry_key(&entry.target);
337 match merged.remove(&key) {
338 None => {
339 merged.insert(key, (entry.target, entry.blob));
340 }
341 Some((target, existing)) => {
342 let merged_blob = merge_context_blobs(existing, entry.blob);
343 merged.insert(key, (target, merged_blob));
344 }
345 }
346 }
347 }
348
349 if merged.is_empty() {
350 return Ok(None);
351 }
352
353 let mut root: Option<ContentHash> = None;
355 for (_, (target, blob)) in merged {
356 if blob.annotations.is_empty() {
357 continue;
358 }
359 let new_root = self.set_context_blob(root.as_ref(), &target, &blob)?;
360 root = Some(new_root);
361 }
362
363 Ok(root)
364 }
365}
366
367fn merge_context_blobs(left: ContextBlob, right: ContextBlob) -> ContextBlob {
373 let format_version = left.format_version;
374 let mut by_id: BTreeMap<String, Annotation> = BTreeMap::new();
375 for annotation in left.annotations.into_iter().chain(right.annotations) {
376 match by_id.remove(&annotation.annotation_id) {
377 None => {
378 by_id.insert(annotation.annotation_id.clone(), annotation);
379 }
380 Some(existing) => {
381 let winner = pick_newer_annotation(existing, annotation);
382 by_id.insert(winner.annotation_id.clone(), winner);
383 }
384 }
385 }
386 ContextBlob {
387 format_version,
388 annotations: by_id.into_values().collect(),
389 }
390}
391
392fn pick_newer_annotation(a: Annotation, b: Annotation) -> Annotation {
393 let ts_a = a
394 .current_revision()
395 .map(|r| r.created_at)
396 .unwrap_or(i64::MIN);
397 let ts_b = b
398 .current_revision()
399 .map(|r| r.created_at)
400 .unwrap_or(i64::MIN);
401 if ts_a > ts_b {
402 a
403 } else if ts_b > ts_a {
404 b
405 } else {
406 let rev_a = a
409 .current_revision()
410 .map(|r| r.revision_id.as_str())
411 .unwrap_or("");
412 let rev_b = b
413 .current_revision()
414 .map(|r| r.revision_id.as_str())
415 .unwrap_or("");
416 if rev_a >= rev_b { a } else { b }
417 }
418}
419
420fn context_entry_key(target: &ContextTarget) -> String {
421 match target {
422 ContextTarget::File { path } => format!("file:{path}"),
423 ContextTarget::State { state_id } => format!("state:{}", state_id.to_string_full()),
424 }
425}
426
427#[derive(Clone, Copy)]
428enum ContextWalkMode {
429 CanonicalOnly,
430}
431
432fn context_target_from_entry_path(path: &Path, mode: ContextWalkMode) -> Option<ContextTarget> {
433 match mode {
434 ContextWalkMode::CanonicalOnly => ContextTarget::from_storage_path(path),
435 }
436}
437
438fn split_path(path: &Path) -> Option<(&str, &Path)> {
439 let mut components = path.components();
440 let first = components.next()?;
441 let std::path::Component::Normal(name) = first else {
442 return None;
443 };
444 Some((name.to_str()?, components.as_path()))
445}
446
447#[cfg(test)]
448mod tests {
449 use objects::object::{Annotation, AnnotationKind, StateAttachment, StateAttachmentBody};
450 use tempfile::TempDir;
451
452 use super::{Repository, *};
453
454 fn setup() -> (TempDir, Repository) {
455 let dir = TempDir::new().unwrap();
456 let repo = Repository::init_default(dir.path()).unwrap();
457 (dir, repo)
458 }
459
460 fn make_annotation(scope: AnnotationScope, content: &str) -> Annotation {
461 Annotation::new(
462 scope,
463 AnnotationKind::Rationale,
464 content.to_string(),
465 vec![],
466 "test@example.com".to_string(),
467 1700000000,
468 None,
469 None,
470 )
471 }
472
473 #[test]
474 fn get_and_set_context_blob_for_file_target() {
475 let (_dir, repo) = setup();
476 let target = ContextTarget::file("src/main.rs").unwrap();
477 let blob = ContextBlob::new(vec![make_annotation(AnnotationScope::File, "Entry point")]);
478
479 let root = repo.set_context_blob(None, &target, &blob).unwrap();
480 let retrieved = repo.get_context_blob(&root, &target).unwrap().unwrap();
481
482 assert_eq!(retrieved, blob);
483 }
484
485 #[test]
486 fn supports_state_targets() {
487 let (_dir, repo) = setup();
488 let target = ContextTarget::state(crate::test_state_id());
489 let blob = ContextBlob::new(vec![make_annotation(AnnotationScope::File, "review note")]);
490
491 let root = repo.set_context_blob(None, &target, &blob).unwrap();
492 let retrieved = repo.get_context_blob(&root, &target).unwrap().unwrap();
493 assert_eq!(retrieved, blob);
494 }
495
496 #[test]
497 fn remove_context_blob_by_scope() {
498 let (_dir, repo) = setup();
499 let target = ContextTarget::file("src/lib.rs").unwrap();
500 let blob = ContextBlob::new(vec![
501 make_annotation(AnnotationScope::File, "file-level"),
502 make_annotation(AnnotationScope::Lines(1, 10), "range-level"),
503 ]);
504
505 let root = repo.set_context_blob(None, &target, &blob).unwrap();
506 let new_root = repo
507 .remove_context_at_target(&root, &target, Some(&AnnotationScope::Lines(1, 10)))
508 .unwrap()
509 .unwrap();
510 let remaining = repo.get_context_blob(&new_root, &target).unwrap().unwrap();
511
512 assert_eq!(remaining.annotations.len(), 1);
513 assert_eq!(
514 remaining
515 .annotations
516 .first()
517 .unwrap()
518 .current_revision()
519 .unwrap()
520 .content,
521 "file-level"
522 );
523 }
524
525 #[test]
526 fn list_context_entries_filters_by_prefix() {
527 let (_dir, repo) = setup();
528 let target1 = ContextTarget::file("src/main.rs").unwrap();
529 let target2 = ContextTarget::file("src/lib.rs").unwrap();
530 let target3 = ContextTarget::file("tests/test.rs").unwrap();
531 let blob1 = ContextBlob::new(vec![make_annotation(AnnotationScope::File, "first")]);
532 let blob2 = ContextBlob::new(vec![make_annotation(AnnotationScope::File, "second")]);
533 let blob3 = ContextBlob::new(vec![make_annotation(AnnotationScope::File, "third")]);
534
535 let root1 = repo.set_context_blob(None, &target1, &blob1).unwrap();
536 let root2 = repo
537 .set_context_blob(Some(&root1), &target2, &blob2)
538 .unwrap();
539 let root3 = repo
540 .set_context_blob(Some(&root2), &target3, &blob3)
541 .unwrap();
542
543 let all = repo.list_context_entries(&root3, None).unwrap();
544 assert_eq!(all.len(), 3);
545
546 let src_only = repo
547 .list_context_entries(&root3, Some(Path::new("src")))
548 .unwrap();
549 assert_eq!(src_only.len(), 2);
550
551 let exact_root_file = repo
552 .list_context_entries(&root3, Some(Path::new("tests/test.rs")))
553 .unwrap();
554 assert_eq!(exact_root_file.len(), 1);
555 }
556
557 #[test]
558 fn find_annotation_returns_target_and_index() {
559 let (_dir, repo) = setup();
560 let target = ContextTarget::file("src/main.rs").unwrap();
561 let blob = ContextBlob::new(vec![make_annotation(AnnotationScope::File, "first")]);
562 let annotation_id = blob.annotations[0].annotation_id.clone();
563 let root = repo.set_context_blob(None, &target, &blob).unwrap();
564
565 let found = repo
566 .find_annotation(&root, &annotation_id)
567 .unwrap()
568 .unwrap();
569 assert_eq!(found.0, target);
570 assert_eq!(found.2, 0);
571 }
572
573 fn state_with_context(repo: &Repository, path: &str, anns: Vec<Annotation>) -> State {
578 let target = ContextTarget::file(path).unwrap();
579 let blob = ContextBlob::new(anns);
580 let root = repo.set_context_blob(None, &target, &blob).unwrap();
581 let state = State::new_snapshot(
582 ContentHash::compute(b""),
583 vec![],
584 objects::object::Attribution::human(objects::object::Principal::new(
585 "test",
586 "test@example.com",
587 )),
588 );
589 repo.store().put_state(&state).unwrap();
590 repo.put_state_attachment(&StateAttachment {
591 state_id: state.id(),
592 body: StateAttachmentBody::Context(root),
593 attribution: state.attribution.clone(),
594 created_at: chrono::Utc::now(),
595 supersedes: None,
596 })
597 .unwrap();
598 state
599 }
600
601 fn ann_with_id(id: &str, content: &str, created_at: i64) -> Annotation {
602 let mut a = Annotation::new(
603 AnnotationScope::File,
604 AnnotationKind::Rationale,
605 content.to_string(),
606 vec![],
607 "test@example.com".to_string(),
608 created_at,
609 None,
610 None,
611 );
612 a.annotation_id = id.to_string();
613 a
614 }
615
616 #[test]
617 fn inherit_parent_context_passes_through_pointer() {
618 let (_dir, repo) = setup();
619 let parent = state_with_context(
620 &repo,
621 "src/lib.rs",
622 vec![make_annotation(AnnotationScope::File, "first")],
623 );
624 let inherited = repo.inherit_parent_context(&parent).unwrap();
625 assert!(inherited.is_some());
626 }
627
628 #[test]
629 fn inherit_parent_context_yields_none_when_parent_has_none() {
630 let (_dir, repo) = setup();
631 let parent = State::new_snapshot(
632 ContentHash::compute(b""),
633 vec![],
634 objects::object::Attribution::human(objects::object::Principal::new(
635 "test",
636 "test@example.com",
637 )),
638 );
639 assert_eq!(repo.inherit_parent_context(&parent).unwrap(), None);
640 }
641
642 #[test]
643 fn union_parent_contexts_returns_none_for_empty_parents() {
644 let (_dir, repo) = setup();
645 let p = State::new_snapshot(
646 ContentHash::compute(b""),
647 vec![],
648 objects::object::Attribution::human(objects::object::Principal::new(
649 "test",
650 "test@example.com",
651 )),
652 );
653 let merged = repo.union_parent_contexts(&[&p, &p]).unwrap();
654 assert_eq!(merged, None);
655 }
656
657 #[test]
658 fn union_parent_contexts_pointer_copies_when_one_side_has_context() {
659 let (_dir, repo) = setup();
660 let parent_with = state_with_context(
661 &repo,
662 "src/lib.rs",
663 vec![make_annotation(AnnotationScope::File, "first")],
664 );
665 let parent_without = State::new_snapshot(
666 ContentHash::compute(b""),
667 vec![],
668 objects::object::Attribution::human(objects::object::Principal::new(
669 "test",
670 "test@example.com",
671 )),
672 );
673 let merged = repo
674 .union_parent_contexts(&[&parent_with, &parent_without])
675 .unwrap();
676 assert_eq!(merged, repo.inherit_parent_context(&parent_with).unwrap());
677 }
678
679 #[test]
680 fn union_parent_contexts_carries_disjoint_annotations() {
681 let (_dir, repo) = setup();
682 let left = state_with_context(
683 &repo,
684 "src/lib.rs",
685 vec![ann_with_id("ann-a", "left side", 1)],
686 );
687 let right = state_with_context(
688 &repo,
689 "src/main.rs",
690 vec![ann_with_id("ann-b", "right side", 1)],
691 );
692 let merged = repo
693 .union_parent_contexts(&[&left, &right])
694 .unwrap()
695 .expect("merged context root");
696 let entries = repo.list_context_entries(&merged, None).unwrap();
697 assert_eq!(entries.len(), 2);
698 let mut ids: Vec<String> = entries
699 .iter()
700 .flat_map(|e| e.blob.annotations.iter().map(|a| a.annotation_id.clone()))
701 .collect();
702 ids.sort();
703 assert_eq!(ids, vec!["ann-a".to_string(), "ann-b".to_string()]);
704 }
705
706 #[test]
707 fn union_parent_contexts_dedupes_same_id_with_newest_revision_wins() {
708 let (_dir, repo) = setup();
709 let older = ann_with_id("ann-shared", "older content", 1);
710 let newer = ann_with_id("ann-shared", "newer content", 9);
711 let left = state_with_context(&repo, "src/lib.rs", vec![older]);
712 let right = state_with_context(&repo, "src/lib.rs", vec![newer]);
713 let merged = repo
714 .union_parent_contexts(&[&left, &right])
715 .unwrap()
716 .expect("merged context root");
717 let entries = repo.list_context_entries(&merged, None).unwrap();
718 assert_eq!(entries.len(), 1);
719 let blob = &entries[0].blob;
720 assert_eq!(blob.annotations.len(), 1);
721 let revision = blob.annotations[0]
722 .current_revision()
723 .expect("annotation has a revision");
724 assert_eq!(revision.content, "newer content");
725 }
726}