copies: implement copy support in `MergedTree::diff_stream()` as adapter

The support for copy tracing is already simply added to the stream
just before yielding the item, so we can easily implement it as a
stream adapter. That ensures that we use the same logic for the
iterator- and stream-based versions. More importantly, it enables
further cleanups and a simpler interface.

lib/src/copies.rs

···
 //! Code for working with copies and renames.
 
 use std::collections::HashMap;
+use std::pin::Pin;
+use std::task::{Context, Poll};
 
 use futures::executor::block_on_stream;
 use futures::stream::BoxStream;
+use futures::Stream;
 
 use crate::backend::{BackendResult, CopyRecord};
+use crate::merged_tree::{MergedTree, TreeDiffEntry, TreeDiffStream};
 use crate::repo_path::{RepoPath, RepoPathBuf};
 
 /// A collection of CopyRecords.
···
         self.records.iter()
     }
 }
+
+/// Wraps a `TreeDiffStream`, adding support for copies and renames.
+pub struct CopiesTreeDiffStream<'a> {
+    inner: TreeDiffStream<'a>,
+    source_tree: MergedTree,
+    copy_records: &'a CopyRecords,
+}
+
+impl<'a> CopiesTreeDiffStream<'a> {
+    /// Create a new diff stream with copy information.
+    pub fn new(
+        inner: TreeDiffStream<'a>,
+        source_tree: MergedTree,
+        copy_records: &'a CopyRecords,
+    ) -> Self {
+        Self {
+            inner,
+            source_tree,
+            copy_records,
+        }
+    }
+}
+
+impl Stream for CopiesTreeDiffStream<'_> {
+    type Item = TreeDiffEntry;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.inner.as_mut().poll_next(cx).map(|option| {
+            option.map(|diff_entry| {
+                diff_entry.adjust_for_copy_tracking(&self.source_tree, self.copy_records)
+            })
+        })
+    }
+}

lib/src/merged_tree.rs

···
 
 use crate::backend;
 use crate::backend::{BackendResult, CopyRecord, MergedTreeId, TreeId, TreeValue};
-use crate::copies::CopyRecords;
+use crate::copies::{CopiesTreeDiffStream, CopyRecords};
 use crate::matchers::{EverythingMatcher, Matcher};
 use crate::merge::{Merge, MergeBuilder, MergedTreeVal, MergedTreeValue};
 use crate::repo_path::{RepoPath, RepoPathBuf, RepoPathComponent};
···
         copy_records: &'matcher CopyRecords,
     ) -> TreeDiffStream<'matcher> {
         let concurrency = self.store().concurrency();
-        if concurrency <= 1 {
+        let stream: TreeDiffStream<'matcher> = if concurrency <= 1 {
             Box::pin(futures::stream::iter(TreeDiffIterator::new(
                 &self.trees,
                 &other.trees,
                 matcher,
-                copy_records,
             )))
         } else {
             Box::pin(TreeDiffStreamImpl::new(
                 self.clone(),
                 other.clone(),
                 matcher,
-                copy_records,
                 concurrency,
             ))
-        }
+        };
+        Box::pin(CopiesTreeDiffStream::new(
+            stream,
+            self.clone(),
+            copy_records,
+        ))
     }
 
     /// Merges this tree with `other`, using `base` as base. Any conflicts will
···
 }
 
 impl TreeDiffEntry {
-    fn adjust_for_copy_tracking(
+    pub(crate) fn adjust_for_copy_tracking(
         self,
         source_tree: &MergedTree,
         copy_records: &CopyRecords,
···
 pub struct TreeDiffIterator<'matcher> {
     store: Arc<Store>,
     stack: Vec<TreeDiffItem>,
-    source: MergedTree,
     matcher: &'matcher dyn Matcher,
-    copy_records: &'matcher CopyRecords,
 }
 
 struct TreeDiffDirItem {
···
 impl<'matcher> TreeDiffIterator<'matcher> {
     /// Creates a iterator over the differences between two trees. Generally
     /// prefer `MergedTree::diff()` of calling this directly.
-    pub fn new(
-        trees1: &Merge<Tree>,
-        trees2: &Merge<Tree>,
-        matcher: &'matcher dyn Matcher,
-        copy_records: &'matcher CopyRecords,
-    ) -> Self {
+    pub fn new(trees1: &Merge<Tree>, trees2: &Merge<Tree>, matcher: &'matcher dyn Matcher) -> Self {
         assert!(Arc::ptr_eq(trees1.first().store(), trees2.first().store()));
         let root_dir = RepoPath::root();
         let mut stack = Vec::new();
···
         Self {
             store: trees1.first().store().clone(),
             stack,
-            source: MergedTree::new(trees1.clone()),
             matcher,
-            copy_records,
         }
     }
 
···
 
     fn next(&mut self) -> Option<Self::Item> {
         self.next_impl()
-            .map(|diff_entry| diff_entry.adjust_for_copy_tracking(&self.source, self.copy_records))
     }
 }
 
 /// Stream of differences between two trees.
 pub struct TreeDiffStreamImpl<'matcher> {
     matcher: &'matcher dyn Matcher,
-    copy_records: &'matcher CopyRecords,
-    source_tree: MergedTree,
     /// Pairs of tree values that may or may not be ready to emit, sorted in the
     /// order we want to emit them. If either side is a tree, there will be
     /// a corresponding entry in `pending_trees`.
···
         tree1: MergedTree,
         tree2: MergedTree,
         matcher: &'matcher dyn Matcher,
-        copy_records: &'matcher CopyRecords,
         max_concurrent_reads: usize,
     ) -> Self {
         let mut stream = Self {
             matcher,
-            copy_records,
-            source_tree: tree1.clone(),
             items: BTreeMap::new(),
             pending_trees: VecDeque::new(),
             max_concurrent_reads,
···
     type Item = TreeDiffEntry;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        self.as_mut().poll_next_impl(cx).map(|option| {
-            option.map(|diff_entry| {
-                diff_entry.adjust_for_copy_tracking(&self.source_tree, self.copy_records)
-            })
-        })
+        self.as_mut().poll_next_impl(cx)
     }
 }
 

lib/tests/test_merged_tree.rs

···
     (diff.target, diff.value.unwrap())
 }
 
-fn diff_stream_equals_iter(
-    tree1: &MergedTree,
-    tree2: &MergedTree,
-    matcher: &dyn Matcher,
-    copy_records: &CopyRecords,
-) {
-    let iter_diff: Vec<_> =
-        TreeDiffIterator::new(tree1.as_merge(), tree2.as_merge(), matcher, copy_records)
+fn diff_stream_equals_iter(tree1: &MergedTree, tree2: &MergedTree, matcher: &dyn Matcher) {
+    let iter_diff: Vec<_> = TreeDiffIterator::new(tree1.as_merge(), tree2.as_merge(), matcher)
+        .map(|diff| (diff.source, diff.target, diff.value.unwrap()))
+        .collect();
+    let max_concurrent_reads = 10;
+    let stream_diff: Vec<_> =
+        TreeDiffStreamImpl::new(tree1.clone(), tree2.clone(), matcher, max_concurrent_reads)
             .map(|diff| (diff.source, diff.target, diff.value.unwrap()))
-            .collect();
-    let max_concurrent_reads = 10;
-    let stream_diff: Vec<_> = TreeDiffStreamImpl::new(
-        tree1.clone(),
-        tree2.clone(),
-        matcher,
-        copy_records,
-        max_concurrent_reads,
-    )
-    .map(|diff| (diff.source, diff.target, diff.value.unwrap()))
-    .collect()
-    .block_on();
+            .collect()
+            .block_on();
     assert_eq!(stream_diff, iter_diff);
 }
 
···
             ),
         )
     );
-    diff_stream_equals_iter(
-        &before_merged,
-        &after_merged,
-        &EverythingMatcher,
-        &CopyRecords::default(),
-    );
+    diff_stream_equals_iter(&before_merged, &after_merged, &EverythingMatcher);
 }
 
 fn create_copy_records(paths: &[(&RepoPath, &RepoPath)]) -> CopyRecords {
···
             ),
         )
     );
-    diff_stream_equals_iter(
-        &before_merged,
-        &after_merged,
-        &EverythingMatcher,
-        &copy_records,
-    );
+    diff_stream_equals_iter(&before_merged, &after_merged, &EverythingMatcher);
 }
 
 /// Diff two conflicted trees
···
         })
         .collect_vec();
     assert_eq!(actual_diff, expected_diff);
-    diff_stream_equals_iter(
-        &left_merged,
-        &right_merged,
-        &EverythingMatcher,
-        &CopyRecords::default(),
-    );
+    diff_stream_equals_iter(&left_merged, &right_merged, &EverythingMatcher);
     // Test the reverse diff
     let actual_diff: Vec<_> = right_merged
         .diff_stream(&left_merged, &EverythingMatcher, &Default::default())
···
         })
         .collect_vec();
     assert_eq!(actual_diff, expected_diff);
-    diff_stream_equals_iter(
-        &right_merged,
-        &left_merged,
-        &EverythingMatcher,
-        &CopyRecords::default(),
-    );
+    diff_stream_equals_iter(&right_merged, &left_merged, &EverythingMatcher);
 }
 
 #[test]
···
             (path6.to_owned(), (Merge::absent(), right_value(path6))),
         ];
         assert_eq!(actual_diff, expected_diff);
-        diff_stream_equals_iter(
-            &left_merged,
-            &right_merged,
-            &EverythingMatcher,
-            &CopyRecords::default(),
-        );
+        diff_stream_equals_iter(&left_merged, &right_merged, &EverythingMatcher);
     }
 
     // Test the reverse diff
···
             ),
         ];
         assert_eq!(actual_diff, expected_diff);
-        diff_stream_equals_iter(
-            &right_merged,
-            &left_merged,
-            &EverythingMatcher,
-            &CopyRecords::default(),
-        );
+        diff_stream_equals_iter(&right_merged, &left_merged, &EverythingMatcher);
     }
 
     // Diff while filtering by `path1` (file1 -> directory1) as a file
···
             (path1.to_owned(), (left_value(path1), Merge::absent())),
         ];
         assert_eq!(actual_diff, expected_diff);
-        diff_stream_equals_iter(
-            &left_merged,
-            &right_merged,
-            &matcher,
-            &CopyRecords::default(),
-        );
+        diff_stream_equals_iter(&left_merged, &right_merged, &matcher);
     }
 
     // Diff while filtering by `path1/file` (file1 -> directory1) as a file
···
             ),
         ];
         assert_eq!(actual_diff, expected_diff);
-        diff_stream_equals_iter(
-            &left_merged,
-            &right_merged,
-            &matcher,
-            &CopyRecords::default(),
-        );
+        diff_stream_equals_iter(&left_merged, &right_merged, &matcher);
     }
 
     // Diff while filtering by `path1` (file1 -> directory1) as a prefix
···
             ),
         ];
         assert_eq!(actual_diff, expected_diff);
-        diff_stream_equals_iter(
-            &left_merged,
-            &right_merged,
-            &matcher,
-            &CopyRecords::default(),
-        );
+        diff_stream_equals_iter(&left_merged, &right_merged, &matcher);
     }
 
     // Diff while filtering by `path6` (directory1 -> file1+(directory1-absent)) as
···
             .block_on();
         let expected_diff = vec![(path6.to_owned(), (Merge::absent(), right_value(path6)))];
         assert_eq!(actual_diff, expected_diff);
-        diff_stream_equals_iter(
-            &left_merged,
-            &right_merged,
-            &matcher,
-            &CopyRecords::default(),
-        );
+        diff_stream_equals_iter(&left_merged, &right_merged, &matcher);
     }
 }