From 1c55c02106025647f2ab426c9c3dd045c1031f95 Mon Sep 17 00:00:00 2001
From: Martin von Zweigbergk <martinvonz@google.com>
Date: Sat, 2 Oct 2021 22:39:49 -0700
Subject: [PATCH] Transaction: remove hidden heads on commit

I recently made the CLI remove hidden heads when a transaction is
committed (38474a9). Let's move that to `Transaction::commit()`, so
the library crate becomes more similar to how the CLI behaves and more
similar to our evolution-less future (#32).
---
 lib/src/repo.rs              | 29 +++++++++++++++++++++++++++++
 lib/src/transaction.rs       |  3 ++-
 lib/tests/test_operations.rs |  2 +-
 src/commands.rs              | 27 ---------------------------
 4 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/lib/src/repo.rs b/lib/src/repo.rs
index 0294c1d18..6b135307f 100644
--- a/lib/src/repo.rs
+++ b/lib/src/repo.rs
@@ -20,6 +20,7 @@ use std::io::Write;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, MutexGuard};
 
+use itertools::Itertools;
 use thiserror::Error;
 
 use crate::backend::{Backend, BackendError, CommitId};
@@ -34,6 +35,7 @@ use crate::local_backend::LocalBackend;
 use crate::op_heads_store::OpHeadsStore;
 use crate::op_store::{BranchTarget, OpStore, OperationId, RefTarget};
 use crate::operation::Operation;
+use crate::revset::RevsetExpression;
 use crate::rewrite::DescendantRebaser;
 use crate::settings::{RepoSettings, UserSettings};
 use crate::simple_op_store::SimpleOpStore;
@@ -799,6 +801,33 @@ impl MutableRepo {
         self.invalidate_evolution();
     }
 
+    pub fn remove_hidden_heads(&mut self) {
+        let mut view = self.view().store_view().clone();
+        let heads_expression =
+            RevsetExpression::commits(view.head_ids.iter().cloned().collect_vec())
+                .non_obsolete_heads();
+        let public_heads_expression =
+            RevsetExpression::commits(view.public_head_ids.iter().cloned().collect_vec())
+                .non_obsolete_heads();
+        view.head_ids.clear();
+        view.public_head_ids.clear();
+        for index_entry in heads_expression
+            .evaluate(self.as_repo_ref())
+            .unwrap()
+            .iter()
+        {
+            view.head_ids.insert(index_entry.commit_id());
+        }
+        for index_entry in public_heads_expression
+            .evaluate(self.as_repo_ref())
+            .unwrap()
+            .iter()
+        {
+            view.public_head_ids.insert(index_entry.commit_id());
+        }
+        self.set_view(view)
+    }
+
     pub fn get_branch(&self, name: &str) -> Option<&BranchTarget> {
         self.view.get_branch(name)
     }
diff --git a/lib/src/transaction.rs b/lib/src/transaction.rs
index dc0b6f417..3fcee355e 100644
--- a/lib/src/transaction.rs
+++ b/lib/src/transaction.rs
@@ -72,7 +72,8 @@ impl Transaction {
     /// That means that a repo can be loaded at the operation, but the
     /// operation will not be seen when loading the repo at head.
     pub fn write(mut self) -> UnpublishedOperation {
-        let mut_repo = self.repo.take().unwrap();
+        let mut mut_repo = self.repo.take().unwrap();
+        mut_repo.remove_hidden_heads();
         let base_repo = mut_repo.base_repo().clone();
         let (mut_index, view, maybe_mut_evolution) = mut_repo.consume();
         let maybe_evolution =
diff --git a/lib/tests/test_operations.rs b/lib/tests/test_operations.rs
index ae4147d7c..7cdcd95ad 100644
--- a/lib/tests/test_operations.rs
+++ b/lib/tests/test_operations.rs
@@ -188,6 +188,6 @@ fn test_isolation(use_git: bool) {
     let repo = repo.reload();
     assert_heads(
         repo.as_repo_ref(),
-        vec![&wc_id, initial.id(), rewrite1.id(), rewrite2.id()],
+        vec![&wc_id, rewrite1.id(), rewrite2.id()],
     );
 }
diff --git a/src/commands.rs b/src/commands.rs
index e868b15d3..f9e89879b 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -453,38 +453,11 @@ impl RepoCommandHelper {
         if self.auto_update_branches {
             update_branches_after_rewrite(mut_repo);
         }
-        remove_hidden_heads(mut_repo);
         self.repo = tx.commit();
         update_working_copy(ui, &self.repo, &self.repo.working_copy_locked())
     }
 }
 
-fn remove_hidden_heads(mut_repo: &mut MutableRepo) {
-    let mut view = mut_repo.view().store_view().clone();
-    let heads_expression =
-        RevsetExpression::commits(view.head_ids.iter().cloned().collect_vec()).non_obsolete_heads();
-    let public_heads_expression =
-        RevsetExpression::commits(view.public_head_ids.iter().cloned().collect_vec())
-            .non_obsolete_heads();
-    view.head_ids.clear();
-    view.public_head_ids.clear();
-    for index_entry in heads_expression
-        .evaluate(mut_repo.as_repo_ref())
-        .unwrap()
-        .iter()
-    {
-        view.head_ids.insert(index_entry.commit_id());
-    }
-    for index_entry in public_heads_expression
-        .evaluate(mut_repo.as_repo_ref())
-        .unwrap()
-        .iter()
-    {
-        view.public_head_ids.insert(index_entry.commit_id());
-    }
-    mut_repo.set_view(view)
-}
-
 fn rev_arg<'a, 'b>() -> Arg<'a, 'b> {
     Arg::with_name("revision")
         .long("revision")