1use std::path::PathBuf;
6use std::collections::HashMap;
7use anyhow::{Result, anyhow};
8use std::fs;
9use toml;
10use chrono::Utc;
11use colored::Colorize;
12
13use crate::vcs::{
14 ShoveId, ObjectId, ObjectStore, Tree, TreeEntry,
15 Repository, Timeline, Shove, Author
16};
17use crate::vcs::objects::{EntryType};
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum MergeStrategy {
22 Auto,
24 FastForwardOnly,
26 AlwaysCreateShove,
28 Ours,
30 Theirs,
32}
33
34#[derive(Debug)]
36pub struct MergeResult {
37 pub success: bool,
39
40 pub shove_id: Option<ShoveId>,
42
43 pub fast_forward: bool,
45
46 pub conflicts: Vec<MergeConflict>,
48}
49
50#[derive(Debug)]
52pub struct MergeConflict {
53 pub path: PathBuf,
55
56 pub base_id: Option<ObjectId>,
58
59 pub ours_id: Option<ObjectId>,
61
62 pub theirs_id: Option<ObjectId>,
64
65 pub resolution: Option<ConflictResolution>,
67}
68
69#[derive(Debug, Clone, PartialEq, Eq)]
71pub enum ConflictResolution {
72 UseOurs,
74 UseTheirs,
76 UseMerged(ObjectId),
78}
79
80pub struct Merger<'a> {
82 pub repo: &'a Repository,
84
85 pub strategy: MergeStrategy,
87}
88
89impl<'a> Merger<'a> {
90 pub fn new(repo: &'a Repository) -> Self {
92 Self {
93 repo,
94 strategy: MergeStrategy::Auto,
95 }
96 }
97
98 pub fn with_strategy(repo: &'a Repository, strategy: MergeStrategy) -> Self {
100 Self { repo, strategy }
101 }
102
103 pub fn merge_timeline(&self, other_timeline: &Timeline) -> Result<MergeResult> {
105 let current_timeline = &self.repo.current_timeline;
107 let current_head = current_timeline.head.as_ref()
108 .ok_or_else(|| anyhow!("Current timeline has no head"))?;
109
110 let other_head = other_timeline.head.as_ref()
112 .ok_or_else(|| anyhow!("Other timeline has no head"))?;
113
114 if current_head == other_head {
116 return Ok(MergeResult {
117 success: true,
118 shove_id: Some(current_head.clone()),
119 fast_forward: true,
120 conflicts: vec![],
121 });
122 }
123
124 let common_ancestor = self.find_common_ancestor(current_head, other_head)?;
126
127 if let Some(ancestor_id) = &common_ancestor {
129 if ancestor_id == current_head {
130 if self.strategy == MergeStrategy::AlwaysCreateShove {
132 return self.create_merge_shove(current_head, other_head, ancestor_id);
134 } else {
135 return Ok(MergeResult {
137 success: true,
138 shove_id: Some(other_head.clone()),
139 fast_forward: true,
140 conflicts: vec![],
141 });
142 }
143 } else if ancestor_id == other_head {
144 return Ok(MergeResult {
146 success: true,
147 shove_id: Some(current_head.clone()),
148 fast_forward: true,
149 conflicts: vec![],
150 });
151 }
152 }
153
154 if self.strategy == MergeStrategy::FastForwardOnly {
156 return Ok(MergeResult {
157 success: false,
158 shove_id: None,
159 fast_forward: false,
160 conflicts: vec![],
161 });
162 }
163
164 self.three_way_merge(current_head, other_head, common_ancestor.as_ref())
166 }
167
168 fn find_common_ancestor(&self, a: &ShoveId, b: &ShoveId) -> Result<Option<ShoveId>> {
170 let a_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", a.as_str()));
172 let b_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", b.as_str()));
173
174 if !a_path.exists() || !b_path.exists() {
175 return Err(anyhow!("One or both shoves not found"));
176 }
177
178 let a_content = fs::read_to_string(&a_path)?;
180 let b_content = fs::read_to_string(&b_path)?;
181
182 let a_shove: Shove = toml::from_str(&a_content)?;
183 let b_shove: Shove = toml::from_str(&b_content)?;
184
185 if self.is_ancestor_of(a, b)? {
187 return Ok(Some(a.clone()));
188 }
189
190 if self.is_ancestor_of(b, a)? {
191 return Ok(Some(b.clone()));
192 }
193
194 let a_ancestors = self.get_ancestors(a)?;
197
198 for b_ancestor in self.get_ancestors(b)? {
200 if a_ancestors.contains(&b_ancestor) {
201 return Ok(Some(b_ancestor));
202 }
203 }
204
205 Ok(None)
207 }
208
209 fn is_ancestor_of(&self, a: &ShoveId, b: &ShoveId) -> Result<bool> {
211 if a == b {
212 return Ok(true);
213 }
214
215 let b_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", b.as_str()));
217 let b_content = fs::read_to_string(&b_path)?;
218 let b_shove: Shove = toml::from_str(&b_content)?;
219
220 for parent_id in &b_shove.parent_ids {
222 if parent_id == a {
223 return Ok(true);
224 }
225
226 if self.is_ancestor_of(a, parent_id)? {
228 return Ok(true);
229 }
230 }
231
232 Ok(false)
233 }
234
235 fn get_ancestors(&self, id: &ShoveId) -> Result<Vec<ShoveId>> {
237 let mut ancestors = Vec::new();
238 let mut to_process = vec![id.clone()];
239
240 while let Some(current_id) = to_process.pop() {
241 if ancestors.contains(¤t_id) {
243 continue;
244 }
245
246 ancestors.push(current_id.clone());
248
249 let shove_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", current_id.as_str()));
251 if !shove_path.exists() {
252 continue;
253 }
254
255 let shove_content = fs::read_to_string(&shove_path)?;
256 let shove: Shove = toml::from_str(&shove_content)?;
257
258 for parent_id in shove.parent_ids {
260 to_process.push(parent_id);
261 }
262 }
263
264 Ok(ancestors)
265 }
266
267 fn three_way_merge(
269 &self,
270 ours: &ShoveId,
271 theirs: &ShoveId,
272 base: Option<&ShoveId>,
273 ) -> Result<MergeResult> {
274 let base_shove = if let Some(base_id) = base {
276 let base_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", base_id.as_str()));
278 if base_path.exists() {
279 let base_content = fs::read_to_string(&base_path)?;
280 Some(toml::from_str::<Shove>(&base_content)?)
281 } else {
282 return Err(anyhow!("Base shove not found: {}", base_id.as_str()));
283 }
284 } else {
285 None
286 };
287
288 let our_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", ours.as_str()));
290 let our_content = fs::read_to_string(&our_path)?;
291 let our_shove: Shove = toml::from_str(&our_content)?;
292
293 let their_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", theirs.as_str()));
295 let their_content = fs::read_to_string(&their_path)?;
296 let their_shove: Shove = toml::from_str(&their_content)?;
297
298 let our_tree_path = self.repo.path.join(".pocket").join("objects").join(our_shove.root_tree_id.as_str());
300 let our_tree_content = fs::read_to_string(&our_tree_path)?;
301 let our_tree: Tree = toml::from_str(&our_tree_content)?;
302
303 let their_tree_path = self.repo.path.join(".pocket").join("objects").join(their_shove.root_tree_id.as_str());
304 let their_tree_content = fs::read_to_string(&their_tree_path)?;
305 let their_tree: Tree = toml::from_str(&their_tree_content)?;
306
307 let base_tree = if let Some(base_shove) = &base_shove {
308 let base_tree_path = self.repo.path.join(".pocket").join("objects").join(base_shove.root_tree_id.as_str());
309 let base_tree_content = fs::read_to_string(&base_tree_path)?;
310 Some(toml::from_str::<Tree>(&base_tree_content)?)
311 } else {
312 None
313 };
314
315 let mut our_entries = HashMap::new();
317 for entry in our_tree.entries {
318 our_entries.insert(entry.name.clone(), entry);
319 }
320
321 let mut their_entries = HashMap::new();
322 for entry in their_tree.entries {
323 their_entries.insert(entry.name.clone(), entry);
324 }
325
326 let base_entries = if let Some(base_tree) = base_tree {
327 let mut entries = HashMap::new();
328 for entry in base_tree.entries {
329 entries.insert(entry.name.clone(), entry);
330 }
331 Some(entries)
332 } else {
333 None
334 };
335
336 let mut merged_entries = Vec::new();
338 let mut conflicts = Vec::new();
339
340 for (name, our_entry) in &our_entries {
342 if our_entry.entry_type != EntryType::File {
343 merged_entries.push(our_entry.clone());
345 continue;
346 }
347
348 if let Some(their_entry) = their_entries.get(name) {
349 if our_entry.id == their_entry.id {
351 merged_entries.push(our_entry.clone());
353 } else {
354 if let Some(base_entries) = &base_entries {
356 if let Some(base_entry) = base_entries.get(name) {
357 if our_entry.id == base_entry.id {
358 merged_entries.push(their_entry.clone());
360 } else if their_entry.id == base_entry.id {
361 merged_entries.push(our_entry.clone());
363 } else {
364 match self.strategy {
366 MergeStrategy::Ours => {
367 merged_entries.push(our_entry.clone());
369 },
370 MergeStrategy::Theirs => {
371 merged_entries.push(their_entry.clone());
373 },
374 _ => {
375 conflicts.push(MergeConflict {
377 path: PathBuf::from(name),
378 base_id: Some(base_entry.id.clone()),
379 ours_id: Some(our_entry.id.clone()),
380 theirs_id: Some(their_entry.id.clone()),
381 resolution: None,
382 });
383
384 merged_entries.push(our_entry.clone());
386 }
387 }
388 }
389 } else {
390 match self.strategy {
392 MergeStrategy::Ours => {
393 merged_entries.push(our_entry.clone());
395 },
396 MergeStrategy::Theirs => {
397 merged_entries.push(their_entry.clone());
399 },
400 _ => {
401 conflicts.push(MergeConflict {
403 path: PathBuf::from(name),
404 base_id: None,
405 ours_id: Some(our_entry.id.clone()),
406 theirs_id: Some(their_entry.id.clone()),
407 resolution: None,
408 });
409
410 merged_entries.push(our_entry.clone());
412 }
413 }
414 }
415 } else {
416 match self.strategy {
418 MergeStrategy::Ours => {
419 merged_entries.push(our_entry.clone());
421 },
422 MergeStrategy::Theirs => {
423 merged_entries.push(their_entry.clone());
425 },
426 _ => {
427 conflicts.push(MergeConflict {
429 path: PathBuf::from(name),
430 base_id: None,
431 ours_id: Some(our_entry.id.clone()),
432 theirs_id: Some(their_entry.id.clone()),
433 resolution: None,
434 });
435
436 merged_entries.push(our_entry.clone());
438 }
439 }
440 }
441 }
442 } else {
443 if let Some(base_entries) = &base_entries {
445 if let Some(_) = base_entries.get(name) {
446 match self.strategy {
448 MergeStrategy::Ours => {
449 merged_entries.push(our_entry.clone());
451 },
452 MergeStrategy::Theirs => {
453 },
455 _ => {
456 conflicts.push(MergeConflict {
458 path: PathBuf::from(name),
459 base_id: Some(base_entries.get(name).unwrap().id.clone()),
460 ours_id: Some(our_entry.id.clone()),
461 theirs_id: None,
462 resolution: None,
463 });
464
465 merged_entries.push(our_entry.clone());
467 }
468 }
469 } else {
470 merged_entries.push(our_entry.clone());
472 }
473 } else {
474 merged_entries.push(our_entry.clone());
476 }
477 }
478 }
479
480 for (name, their_entry) in &their_entries {
482 if their_entry.entry_type != EntryType::File {
483 if !our_entries.contains_key(name) {
485 merged_entries.push(their_entry.clone());
486 }
487 continue;
488 }
489
490 if !our_entries.contains_key(name) {
491 if let Some(base_entries) = &base_entries {
493 if let Some(_) = base_entries.get(name) {
494 match self.strategy {
496 MergeStrategy::Ours => {
497 },
499 MergeStrategy::Theirs => {
500 merged_entries.push(their_entry.clone());
502 },
503 _ => {
504 conflicts.push(MergeConflict {
506 path: PathBuf::from(name),
507 base_id: Some(base_entries.get(name).unwrap().id.clone()),
508 ours_id: None,
509 theirs_id: Some(their_entry.id.clone()),
510 resolution: None,
511 });
512
513 }
515 }
516 } else {
517 merged_entries.push(their_entry.clone());
519 }
520 } else {
521 merged_entries.push(their_entry.clone());
523 }
524 }
525 }
526
527 if !conflicts.is_empty() && self.strategy != MergeStrategy::Ours && self.strategy != MergeStrategy::Theirs {
530 return Ok(MergeResult {
531 success: false,
532 shove_id: None,
533 fast_forward: false,
534 conflicts,
535 });
536 }
537
538 let merged_tree = Tree {
540 entries: merged_entries,
541 };
542
543 let object_store = ObjectStore::new(self.repo.path.clone());
545 let tree_id = object_store.store_tree(&merged_tree)?;
546
547 let author = Author {
549 name: self.repo.config.user.name.clone(),
550 email: self.repo.config.user.email.clone(),
551 timestamp: Utc::now(),
552 };
553
554 let mut parent_ids = vec![ours.clone()];
555 if ours != theirs {
556 parent_ids.push(theirs.clone());
557 }
558
559 let message = format!("Merge {} into {}",
560 their_shove.message.lines().next().unwrap_or(""),
561 our_shove.message.lines().next().unwrap_or(""));
562
563 let shove = Shove::new(&self.repo.pile, parent_ids, author, &message, tree_id);
564
565 let shove_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", shove.id.as_str()));
567 shove.save(&shove_path)?;
568
569 Ok(MergeResult {
570 success: true,
571 shove_id: Some(shove.id.clone()),
572 fast_forward: false,
573 conflicts: Vec::new(),
574 })
575 }
576
577 fn create_merge_shove(
579 &self,
580 ours: &ShoveId,
581 theirs: &ShoveId,
582 base: &ShoveId,
583 ) -> Result<MergeResult> {
584 let ours_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", ours.as_str()));
586 let theirs_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", theirs.as_str()));
587
588 let ours_content = fs::read_to_string(&ours_path)?;
589 let theirs_content = fs::read_to_string(&theirs_path)?;
590
591 let our_shove: Shove = toml::from_str(&ours_content)?;
592 let their_shove: Shove = toml::from_str(&theirs_content)?;
593
594 let merged_tree_id = self.merge_trees(&our_shove.root_tree_id, &their_shove.root_tree_id)?;
596
597 let mut parent_ids = vec![our_shove.id.clone()];
599 parent_ids.push(their_shove.id.clone());
600
601 let message = format!("Merge {} into {}", theirs.as_str(), ours.as_str());
603
604 let author = Author {
606 name: self.repo.config.user.name.clone(),
607 email: self.repo.config.user.email.clone(),
608 timestamp: Utc::now(),
609 };
610
611 let new_shove = Shove::new(
613 &self.repo.pile,
614 parent_ids,
615 author,
616 &message,
617 merged_tree_id,
618 );
619
620 let shove_path = self.repo.path.join(".pocket").join("shoves").join(format!("{}.toml", new_shove.id.as_str()));
622 new_shove.save(&shove_path)?;
623
624 let timeline_path = self.repo.path.join(".pocket").join("timelines").join("current");
626 let mut timeline = Timeline::load(&timeline_path)?;
627 timeline.head = Some(new_shove.id.clone());
628 timeline.save(&timeline_path)?;
629
630 Ok(MergeResult {
632 success: true,
633 shove_id: Some(new_shove.id),
634 fast_forward: false,
635 conflicts: Vec::new(),
636 })
637 }
638
639 fn merge_trees(&self, ours_id: &ObjectId, theirs_id: &ObjectId) -> Result<ObjectId> {
640 let ours_path = self.repo.path.join(".pocket").join("objects").join(ours_id.as_str());
642 let theirs_path = self.repo.path.join(".pocket").join("objects").join(theirs_id.as_str());
643
644 let ours_content = fs::read_to_string(&ours_path)?;
645 let theirs_content = fs::read_to_string(&theirs_path)?;
646
647 let ours_tree: Tree = toml::from_str(&ours_content)?;
648 let theirs_tree: Tree = toml::from_str(&theirs_content)?;
649
650 let mut merged_entries = Vec::new();
652 let mut our_entries_map = std::collections::HashMap::new();
653
654 for entry in &ours_tree.entries {
656 our_entries_map.insert(entry.name.clone(), entry.clone());
657 }
658
659 for entry in &ours_tree.entries {
661 merged_entries.push(entry.clone());
662 }
663
664 for entry in &theirs_tree.entries {
666 if !our_entries_map.contains_key(&entry.name) {
667 merged_entries.push(entry.clone());
668 }
669 }
670
671 let merged_tree = Tree {
673 entries: merged_entries,
674 };
675
676 let object_store = ObjectStore::new(self.repo.path.clone());
678 let tree_id = object_store.store_tree(&merged_tree)?;
679
680 Ok(tree_id)
681 }
682
683 pub fn resolve_conflicts_interactively(&self, conflicts: &[MergeConflict]) -> Result<Vec<ConflictResolution>> {
685 let mut resolutions = Vec::new();
686
687 println!("\n{} Resolving {} conflicts interactively", "🔄".bright_blue(), conflicts.len());
688
689 for (i, conflict) in conflicts.iter().enumerate() {
690 println!("\n{} Conflict {}/{}: {}", "⚠️".yellow(), i + 1, conflicts.len(), conflict.path.display());
691
692 self.display_conflict(conflict)?;
694
695 let resolution = self.prompt_for_resolution(conflict)?;
697 resolutions.push(resolution);
698
699 println!("{} Conflict resolved", "✅".green());
700 }
701
702 println!("\n{} All conflicts resolved", "🎉".green());
703
704 Ok(resolutions)
705 }
706
707 fn display_conflict(&self, conflict: &MergeConflict) -> Result<()> {
709 let base_content = if let Some(id) = &conflict.base_id {
711 self.load_object_content(id)?
712 } else {
713 String::new()
714 };
715
716 let ours_content = if let Some(id) = &conflict.ours_id {
717 self.load_object_content(id)?
718 } else {
719 String::new()
720 };
721
722 let theirs_content = if let Some(id) = &conflict.theirs_id {
723 self.load_object_content(id)?
724 } else {
725 String::new()
726 };
727
728 println!("\n{} Base version:", "⚪".bright_black());
730 self.print_content(&base_content, " ");
731
732 println!("\n{} Our version (current timeline):", "🟢".green());
733 self.print_content(&ours_content, " ");
734
735 println!("\n{} Their version (incoming timeline):", "🔵".blue());
736 self.print_content(&theirs_content, " ");
737
738 Ok(())
739 }
740
741 fn print_content(&self, content: &str, prefix: &str) {
743 for (i, line) in content.lines().enumerate() {
744 println!("{}{:3} | {}", prefix, i + 1, line);
745 }
746
747 if content.is_empty() {
748 println!("{} | (empty file)", prefix);
749 }
750 }
751
752 fn prompt_for_resolution(&self, conflict: &MergeConflict) -> Result<ConflictResolution> {
754 println!("\nHow would you like to resolve this conflict?");
755 println!(" 1. {} Use our version (current timeline)", "🟢".green());
756 println!(" 2. {} Use their version (incoming timeline)", "🔵".blue());
757 println!(" 3. {} Edit and merge manually", "✏️".yellow());
758
759 println!("\nSelected: 1. Use our version");
762
763 if let Some(id) = &conflict.ours_id {
764 Ok(ConflictResolution::UseOurs)
765 } else {
766 Ok(ConflictResolution::UseTheirs)
768 }
769 }
770
771 fn load_object_content(&self, id: &ObjectId) -> Result<String> {
773 let object_path = self.repo.path.join(".pocket").join("objects").join(id.as_str());
774 let content = fs::read_to_string(object_path)?;
775 Ok(content)
776 }
777
778 pub fn apply_resolutions(&self, conflicts: &[MergeConflict], resolutions: &[ConflictResolution]) -> Result<Tree> {
780 let mut merged_tree = Tree { entries: Vec::new() };
782
783 for (conflict, resolution) in conflicts.iter().zip(resolutions.iter()) {
785 match resolution {
786 ConflictResolution::UseOurs => {
787 if let Some(id) = &conflict.ours_id {
788 merged_tree.entries.push(TreeEntry {
789 name: conflict.path.to_string_lossy().to_string(),
790 id: id.clone(),
791 entry_type: EntryType::File,
792 permissions: 0o644,
793 });
794 }
795 },
796 ConflictResolution::UseTheirs => {
797 if let Some(id) = &conflict.theirs_id {
798 merged_tree.entries.push(TreeEntry {
799 name: conflict.path.to_string_lossy().to_string(),
800 id: id.clone(),
801 entry_type: EntryType::File,
802 permissions: 0o644,
803 });
804 }
805 },
806 ConflictResolution::UseMerged(id) => {
807 merged_tree.entries.push(TreeEntry {
808 name: conflict.path.to_string_lossy().to_string(),
809 id: id.clone(),
810 entry_type: EntryType::File,
811 permissions: 0o644,
812 });
813 },
814 }
815 }
816
817 Ok(merged_tree)
818 }
819}