From 6a5615c9334a0202ef7f3852fbfaaee560056248 Mon Sep 17 00:00:00 2001 From: Martin von Zweigbergk Date: Fri, 3 Nov 2023 14:20:11 -0700 Subject: [PATCH] rewrite: use `MergedTree::diff_stream()` when restoring from tree --- lib/src/rewrite.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/lib/src/rewrite.rs b/lib/src/rewrite.rs index 36bad4c74..bd97856af 100644 --- a/lib/src/rewrite.rs +++ b/lib/src/rewrite.rs @@ -17,7 +17,9 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use futures::StreamExt; use itertools::Itertools; +use pollster::FutureExt; use tracing::instrument; use crate::backend::{BackendError, BackendResult, CommitId, MergedTreeId, ObjectId}; @@ -83,10 +85,15 @@ pub fn restore_tree( // TODO: We should be able to not traverse deeper in the diff if the matcher // matches an entire subtree. let mut tree_builder = MergedTreeBuilder::new(destination.id().clone()); - for (repo_path, diff) in source.diff(destination, matcher) { - let (source_value, _destination_value) = diff?; - tree_builder.set_or_remove(repo_path, source_value); + async { + let mut diff_stream = source.diff_stream(destination, matcher); + while let Some((repo_path, diff)) = diff_stream.next().await { + let (source_value, _destination_value) = diff?; + tree_builder.set_or_remove(repo_path, source_value); + } + Ok::<(), BackendError>(()) } + .block_on()?; tree_builder.write_tree(destination.store()) } }