feat(effects): add with_progress method to iterators

In situations where we `continue` a loop, the naive approach of incrementing the iterator progress at the end of the loop won't work as it'll skip the increment. This extension method aims to improve the ergonomics and correctness of that use-case.

Author: arxanas

git-branchless-lib/src/core/effects.rs

@@ -1000,7 +1000,7 @@ impl Drop for ProgressHandle<'_> {
     }
 }
 
-impl ProgressHandle<'_> {
+impl<'a> ProgressHandle<'a> {
     /// Notify the progress meter that the current operation has `total`
     /// discrete units of work, and it's currently `current` units of the way
     /// through the operation.
@@ -1025,6 +1025,46 @@ impl ProgressHandle<'_> {
     }
 }
 
+/// A wrapper around an iterator that reports progress as it iterates.
+pub struct ProgressIter<'a, TItem, TInner: Iterator<Item = TItem>> {
+    index: usize,
+    inner: TInner,
+    progress: ProgressHandle<'a>,
+}
+
+impl<TItem, TInner: Iterator<Item = TItem>> Iterator for ProgressIter<'_, TItem, TInner> {
+    type Item = TItem;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let (lower, upper) = self.inner.size_hint();
+        let size_guess = upper.unwrap_or(lower);
+        self.progress.notify_progress(self.index, size_guess);
+        self.index += 1;
+        self.inner.next()
+    }
+}
+
+/// Extension trait for iterators that adds a `with_progress` method.
+pub trait WithProgress<'a, TItem>: Iterator<Item = TItem> {
+    /// The type of the iterator returned by `with_progress`.
+    type Iter: Iterator<Item = TItem> + 'a;
+
+    /// Wrap the iterator into an iterator that reports progress as it consumes items.
+    fn with_progress(self, progress: ProgressHandle<'a>) -> Self::Iter;
+}
+
+impl<'a, TItem: 'a, TIter: Iterator<Item = TItem> + 'a> WithProgress<'a, TItem> for TIter {
+    type Iter = ProgressIter<'a, TItem, TIter>;
+
+    fn with_progress(self, progress: ProgressHandle<'a>) -> Self::Iter {
+        ProgressIter {
+            index: 0,
+            inner: self,
+            progress,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

git-branchless-lib/src/core/rewrite/plan.rs

@@ -11,7 +11,7 @@ use rayon::{prelude::*, ThreadPool};
 use tracing::{instrument, warn};
 
 use crate::core::dag::{sorted_commit_set, union_all, CommitSet, Dag};
-use crate::core::effects::{Effects, OperationType};
+use crate::core::effects::{Effects, OperationType, WithProgress};
 use crate::core::formatting::Pluralize;
 use crate::core::rewrite::{RepoPool, RepoResource};
 use crate::core::task::ResourcePool;
@@ -1368,12 +1368,15 @@ impl<'a> RebasePlanBuilder<'a> {
             let path: Vec<CacheLookupResult<Option<NonZeroOid>, NonZeroOid>> = {
                 let (effects, progress) = effects.start_operation(OperationType::ReadingFromCache);
                 let _effects = effects;
-                progress.notify_progress(0, path.len());
                 if touched_paths_cache.is_empty() {
                     // Fast path for when the cache hasn't been populated.
-                    path.into_iter().map(CacheLookupResult::NotCached).collect()
+                    path.into_iter()
+                        .with_progress(progress)
+                        .map(CacheLookupResult::NotCached)
+                        .collect()
                 } else {
                     path.into_iter()
+                        .with_progress(progress)
                         .map(|commit_oid| match touched_paths_cache.get(&commit_oid) {
                             Some(upstream_touched_paths) => {
                                 if Self::should_check_patch_id(
@@ -1387,15 +1390,14 @@ impl<'a> RebasePlanBuilder<'a> {
                             }
                             None => CacheLookupResult::NotCached(commit_oid),
                         })
-                        .inspect(|_| progress.notify_progress_inc(1))
                         .collect()
                 }
             };
 
-            let (_effects, progress) = effects.start_operation(OperationType::FilterByTouchedPaths);
-            progress.notify_progress(0, path.len());
-
+            let (effects, progress) = effects.start_operation(OperationType::FilterByTouchedPaths);
+            let _effects = effects;
             pool.install(|| {
+                progress.notify_progress(0, path.len());
                 path.into_par_iter()
                     .map(|commit_oid| {
                         let commit_oid = match commit_oid {

git-branchless-submit/src/phabricator.rs

@@ -19,7 +19,7 @@ use itertools::Itertools;
 use lazy_static::lazy_static;
 use lib::core::check_out::CheckOutCommitOptions;
 use lib::core::dag::{CommitSet, Dag};
-use lib::core::effects::{Effects, OperationType};
+use lib::core::effects::{Effects, OperationType, WithProgress};
 use lib::core::eventlog::EventLogDb;
 use lib::core::formatting::StyledStringBuilder;
 use lib::core::rewrite::{
@@ -882,8 +882,7 @@ impl PhabricatorForge<'_> {
         // Newly-created commits won't have been observed by the DAG, so add them in manually here.
         let draft_commits = self.dag.query_draft_commits()?.union(newly_created_commits);
 
-        progress.notify_progress(0, commit_oids.len());
-        for commit_oid in commit_oids {
+        for commit_oid in commit_oids.into_iter().with_progress(progress) {
             let id = match self.get_revision_id(commit_oid)? {
                 Some(id) => id,
                 None => {
@@ -929,7 +928,6 @@ impl PhabricatorForge<'_> {
                 Ok(()) => {}
                 Err(exit_code) => return Ok(Err(exit_code)),
             }
-            progress.notify_progress_inc(1);
         }
         Ok(Ok(()))
     }

git-branchless/src/commands/repair.rs

@@ -2,6 +2,7 @@ use std::fmt::Write;
 use std::{collections::HashSet, time::SystemTime};
 
 use itertools::Itertools;
+use lib::core::effects::WithProgress;
 use lib::git::{CategorizedReferenceName, MaybeZeroOid};
 use lib::util::EyreExitOr;
 use lib::{
@@ -24,13 +25,11 @@ pub fn repair(effects: &Effects, dry_run: bool) -> EyreExitOr<()> {
         let (effects, progress) = effects.start_operation(OperationType::RepairCommits);
         let _effects = effects;
         let cursor_oids = event_replayer.get_cursor_oids(event_cursor);
-        progress.notify_progress(0, cursor_oids.len());
         let mut result = HashSet::new();
-        for oid in cursor_oids {
+        for oid in cursor_oids.into_iter().with_progress(progress) {
             if repo.find_commit(oid)?.is_none() {
                 result.insert(oid);
             }
-            progress.notify_progress_inc(1);
         }
         result
     };
@@ -48,13 +47,11 @@ pub fn repair(effects: &Effects, dry_run: bool) -> EyreExitOr<()> {
                     .map(move |reference_name| (oid, reference_name))
             })
             .collect_vec();
-        progress.notify_progress(0, branch_names.len());
         let mut result = HashSet::new();
-        for (oid, reference_name) in branch_names {
+        for (oid, reference_name) in branch_names.into_iter().with_progress(progress) {
             if repo.find_reference(&reference_name)?.is_none() {
                 result.insert((oid, reference_name));
             }
-            progress.notify_progress_inc(1);
         }
         result
     };

git-branchless/src/commands/sync.rs

@@ -15,7 +15,7 @@ use git_branchless_opts::{MoveOptions, ResolveRevsetOptions, Revset};
 use git_branchless_revset::{check_revset_syntax, resolve_commits};
 use lib::core::config::get_restack_preserve_timestamps;
 use lib::core::dag::{sorted_commit_set, union_all, CommitSet, Dag};
-use lib::core::effects::{Effects, OperationType};
+use lib::core::effects::{Effects, OperationType, WithProgress};
 use lib::core::eventlog::{EventLogDb, EventReplayer};
 use lib::core::formatting::{Pluralize, StyledStringBuilder};
 use lib::core::rewrite::{
@@ -403,9 +403,9 @@ fn execute_plans(
         let mut skipped_commits: Vec<Commit> = Vec::new();
 
         let (effects, progress) = effects.start_operation(OperationType::SyncCommits);
-        progress.notify_progress(0, root_commit_and_plans.len());
-
-        for (root_commit_oid, rebase_plan) in root_commit_and_plans {
+        for (root_commit_oid, rebase_plan) in
+            root_commit_and_plans.into_iter().with_progress(progress)
+        {
             let root_commit = repo.find_commit_or_fail(root_commit_oid)?;
             let rebase_plan = match rebase_plan {
                 Some(rebase_plan) => rebase_plan,
@@ -423,7 +423,6 @@ fn execute_plans(
                 &rebase_plan,
                 execute_options,
             )?;
-            progress.notify_progress_inc(1);
             match result {
                 ExecuteRebasePlanResult::Succeeded { rewritten_oids: _ } => {
                     success_commits.push(root_commit);