transaction: make sure set of heads has only heads

`Transaction::add_head()` and others would let the caller add
non-heads to the set (i.e. ancestors of others heads) and the the
non-heads were filterd out when the transaction was committed. That's
a little surprising, so let's try to keep the set valid even within a
transaction. That will surely make commands that add many commits
noticeably slower in large repos. Hopefully we can improve that
later.

lib/src/view.rs

@@ -369,11 +369,19 @@ impl MutableView {
         self.data.checkout = id;
     }
 
+    // TODO: This is surely terribly slow on large repos, at least in its current
+    // form. We should make it faster (using the index) and avoid calling it in
+    // most cases (avoid adding a head that's already reachable in the view).
+    fn remove_non_heads(&mut self) {
+        self.data.head_ids = heads_of_set(&self.store, self.heads().cloned());
+    }
+
     pub fn add_head(&mut self, head: &Commit) {
         self.data.head_ids.insert(head.id().clone());
         for parent in head.parents() {
             self.data.head_ids.remove(parent.id());
         }
+        self.remove_non_heads();
     }
 
     pub fn remove_head(&mut self, head: &Commit) {
@@ -381,18 +389,19 @@ impl MutableView {
         for parent in head.parents() {
             self.data.head_ids.insert(parent.id().clone());
         }
+        self.remove_non_heads();
     }
 
     pub fn set_view(&mut self, data: op_store::View) {
         self.data = data;
+        self.remove_non_heads();
     }
 
-    pub fn save(mut self, description: String, operation_start_time: Timestamp) -> Operation {
+    pub fn save(self, description: String, operation_start_time: Timestamp) -> Operation {
         let op_heads_dir = self.path.join("op_heads");
 
         // First write the current view whether or not there have been any concurrent
         // operations. We'll later create a merge operation if necessary.
-        self.data.head_ids = heads_of_set(&self.store, self.heads().cloned());
         let view_id = self.op_store.write_view(&self.data).unwrap();
         let mut operation_metadata = OperationMetadata::new(description);
         operation_metadata.start_time = operation_start_time;

lib/tests/test_transaction.rs

@@ -18,6 +18,7 @@ use jj_lib::repo_path::FileRepoPath;
 use jj_lib::store::{Conflict, ConflictId, ConflictPart, TreeValue};
 use jj_lib::store_wrapper::StoreWrapper;
 use jj_lib::testutils;
+use std::collections::HashSet;
 use std::sync::Arc;
 use test_case::test_case;
 
@@ -301,3 +302,64 @@ fn test_checkout_previous_empty_and_pruned(use_git: bool) {
         .is_empty());
     tx.discard();
 }
+
+#[test_case(false ; "local store")]
+// #[test_case(true ; "git store")]
+fn test_add_head_success(use_git: bool) {
+    // Test that Transaction::add_head() adds the head, and that it's still there
+    // after commit.
+    let settings = testutils::user_settings();
+    let (_temp_dir, mut repo) = testutils::init_repo(&settings, use_git);
+
+    // Create a commit outside of the repo by using a temporary transaction. Then
+    // add that as a head.
+    let mut tx = repo.start_transaction("test");
+    let new_commit = CommitBuilder::for_new_commit(
+        &settings,
+        repo.store(),
+        repo.store().empty_tree_id().clone(),
+    )
+    .write_to_transaction(&mut tx);
+    tx.discard();
+
+    let mut tx = repo.start_transaction("test");
+    let heads: HashSet<_> = tx.as_repo().view().heads().cloned().collect();
+    assert!(!heads.contains(new_commit.id()));
+    tx.add_head(&new_commit);
+    let heads: HashSet<_> = tx.as_repo().view().heads().cloned().collect();
+    assert!(heads.contains(new_commit.id()));
+    tx.commit();
+    Arc::get_mut(&mut repo).unwrap().reload();
+    let heads: HashSet<_> = repo.view().heads().cloned().collect();
+    assert!(heads.contains(new_commit.id()));
+}
+
+#[test_case(false ; "local store")]
+// #[test_case(true ; "git store")]
+fn test_add_head_ancestor(use_git: bool) {
+    // Test that Transaction::add_head() does not add a head if it's an ancestor of
+    // an existing head.
+    let settings = testutils::user_settings();
+    let (_temp_dir, mut repo) = testutils::init_repo(&settings, use_git);
+    let store = repo.store();
+
+    // Create a commit outside of the repo by using a temporary transaction. Then
+    // add that as a head.
+    let mut tx = repo.start_transaction("test");
+    let commit1 = CommitBuilder::for_new_commit(&settings, store, store.empty_tree_id().clone())
+        .write_to_transaction(&mut tx);
+    let commit2 = CommitBuilder::for_new_commit(&settings, store, store.empty_tree_id().clone())
+        .set_parents(vec![commit1.id().clone()])
+        .write_to_transaction(&mut tx);
+    let _commit3 = CommitBuilder::for_new_commit(&settings, store, store.empty_tree_id().clone())
+        .set_parents(vec![commit2.id().clone()])
+        .write_to_transaction(&mut tx);
+    tx.commit();
+    Arc::get_mut(&mut repo).unwrap().reload();
+
+    let mut tx = repo.start_transaction("test");
+    tx.add_head(&commit1);
+    let heads: HashSet<_> = tx.as_repo().view().heads().cloned().collect();
+    assert!(!heads.contains(commit1.id()));
+    tx.discard();
+}