index: keep up to date within transaction

With tons of groundwork done, wee can now finally keep the index up to
date within a transaction! That means that we can start relying on the
index to always be valid, so we can use it e.g. for finding common
ancestors within a transaction. That should help speed up `jj evolve`
immensely on large repos.

We still don't write the updated index to disk when the transaction
closes. That will come later.

lib/src/dag_walk.rs

@@ -124,10 +124,10 @@ where
 }
 
 /// Returns neighbors before the node itself.
-pub fn topo_order_reverse<T, ID, II, NI>(
+pub fn topo_order_reverse<'a, T, ID, II, NI>(
     start: II,
-    id_fn: Box<dyn Fn(&T) -> ID>,
+    id_fn: Box<dyn Fn(&T) -> ID + 'a>,
-    mut neighbors_fn: Box<dyn FnMut(&T) -> NI>,
+    mut neighbors_fn: Box<dyn FnMut(&T) -> NI + 'a>,
 ) -> Vec<T>
 where
     T: Hash + Eq + Clone,

lib/src/index.rs

@@ -375,11 +375,14 @@ impl MutableIndex {
         IndexRef::Mutable(self)
     }
 
-    fn add_commit(&mut self, commit: &Commit) {
+    pub fn add_commit(&mut self, commit: &Commit) {
         self.add_commit_data(commit.id().clone(), commit.parent_ids());
     }
 
     fn add_commit_data(&mut self, id: CommitId, parent_ids: Vec<CommitId>) {
+        if self.lookup.contains_key(&id) {
+            return;
+        }
         let mut entry = MutableGraphEntry {
             commit_id: id,
             generation_number: 0,

lib/src/repo.rs

@@ -24,7 +24,7 @@ use thiserror::Error;
 use crate::commit_builder::{new_change_id, signature};
 use crate::evolution::{EvolutionRef, MutableEvolution, ReadonlyEvolution};
 use crate::git_store::GitStore;
-use crate::index::ReadonlyIndex;
+use crate::index::{IndexRef, MutableIndex, ReadonlyIndex};
 use crate::local_store::LocalStore;
 use crate::operation::Operation;
 use crate::settings::{RepoSettings, UserSettings};
@@ -70,6 +70,13 @@ impl<'a, 'r> RepoRef<'a, 'r> {
         }
     }
 
+    pub fn index(&self) -> IndexRef {
+        match self {
+            RepoRef::Readonly(repo) => IndexRef::Readonly(repo.index()),
+            RepoRef::Mutable(repo) => IndexRef::Mutable(repo.index()),
+        }
+    }
+
     pub fn view(&self) -> ViewRef<'a> {
         match self {
             RepoRef::Readonly(repo) => ViewRef::Readonly(repo.view()),
@@ -327,7 +334,12 @@ impl ReadonlyRepo {
     }
 
     pub fn start_transaction(&self, description: &str) -> Transaction {
-        let mut_repo = MutableRepo::new(self, &self.view, &self.evolution.as_ref().unwrap());
+        let mut_repo = MutableRepo::new(
+            self,
+            self.index(),
+            &self.view,
+            &self.evolution.as_ref().unwrap(),
+        );
         Transaction::new(mut_repo, description)
     }
 
@@ -356,6 +368,7 @@ impl ReadonlyRepo {
 
 pub struct MutableRepo<'r> {
     repo: &'r ReadonlyRepo,
+    index: Option<MutableIndex>,
     view: Option<MutableView>,
     evolution: Option<MutableEvolution<'static, 'static>>,
 }
@@ -363,12 +376,15 @@ pub struct MutableRepo<'r> {
 impl<'r> MutableRepo<'r> {
     pub fn new(
         repo: &'r ReadonlyRepo,
+        index: Arc<ReadonlyIndex>,
         view: &ReadonlyView,
         evolution: &ReadonlyEvolution<'r>,
     ) -> Arc<MutableRepo<'r>> {
         let mut_view = view.start_modification();
+        let mut_index = MutableIndex::incremental(index);
         let mut mut_repo = Arc::new(MutableRepo {
             repo,
+            index: Some(mut_index),
             view: Some(mut_view),
             evolution: None,
         });
@@ -390,6 +406,14 @@ impl<'r> MutableRepo<'r> {
         self.repo.store()
     }
 
+    pub fn index(&self) -> &MutableIndex {
+        self.index.as_ref().unwrap()
+    }
+
+    pub fn index_mut(&mut self) -> &mut MutableIndex {
+        self.index.as_mut().unwrap()
+    }
+
     pub fn base_repo(&self) -> &'r ReadonlyRepo {
         self.repo
     }

lib/src/transaction.rs

@@ -15,7 +15,9 @@
 use crate::commit::Commit;
 use crate::commit_builder::CommitBuilder;
 use crate::conflicts;
+use crate::dag_walk::topo_order_reverse;
 use crate::evolution::MutableEvolution;
+use crate::index::MutableIndex;
 use crate::op_store;
 use crate::operation::Operation;
 use crate::repo::{MutableRepo, ReadonlyRepo, RepoRef};
@@ -59,6 +61,10 @@ impl<'r> Transaction<'r> {
         Arc::get_mut(self.repo.as_mut().unwrap()).unwrap()
     }
 
+    pub fn index(&self) -> &MutableIndex {
+        self.repo.as_ref().unwrap().index()
+    }
+
     pub fn view(&self) -> &MutableView {
         self.repo.as_ref().unwrap().view()
     }
@@ -135,9 +141,25 @@ impl<'r> Transaction<'r> {
             .iter()
             .all(|parent_id| current_heads.contains(parent_id))
         {
+            mut_repo.index_mut().add_commit(head);
             mut_repo.view_mut().add_head(head);
             mut_repo.evolution_mut().add_commit(head);
         } else {
+            let missing_commits = topo_order_reverse(
+                vec![head.clone()],
+                Box::new(|commit: &Commit| commit.id().clone()),
+                Box::new(|commit: &Commit| -> Vec<Commit> {
+                    commit
+                        .parents()
+                        .into_iter()
+                        .filter(|parent| !current_heads.contains(parent.id()))
+                        .collect()
+                }),
+            );
+            let mut_index = mut_repo.index_mut();
+            for missing_commit in missing_commits.iter().rev() {
+                mut_index.add_commit(missing_commit);
+            }
             mut_repo.view_mut().add_head(head);
             mut_repo.evolution_mut().invalidate();
         }

lib/tests/test_index.rs

@@ -337,10 +337,13 @@ fn test_index_commits_incremental(use_git: bool) {
     assert_eq!(stats.num_commits, 2 + 3);
     assert_eq!(stats.num_merges, 0);
     assert_eq!(stats.max_generation_number, 3);
-    assert_eq!(stats.levels.len(), 2);
+    assert_eq!(stats.levels.len(), 3);
     assert_eq!(stats.levels[0].num_commits, 2);
-    assert_eq!(stats.levels[1].num_commits, 3);
+    assert_eq!(stats.levels[1].num_commits, 1);
     assert_ne!(stats.levels[1].name, stats.levels[0].name);
+    assert_eq!(stats.levels[2].num_commits, 2);
+    assert_ne!(stats.levels[2].name, stats.levels[0].name);
+    assert_ne!(stats.levels[2].name, stats.levels[1].name);
 
     assert_eq!(generation_number(index.clone(), root_commit.id()), 0);
     assert_eq!(generation_number(index.clone(), commit_a.id()), 1);
@@ -383,9 +386,10 @@ fn test_index_commits_incremental_empty_transaction(use_git: bool) {
     assert_eq!(stats.num_commits, 2 + 1);
     assert_eq!(stats.num_merges, 0);
     assert_eq!(stats.max_generation_number, 1);
-    assert_eq!(stats.levels.len(), 2);
+    assert_eq!(stats.levels.len(), 3);
     assert_eq!(stats.levels[0].num_commits, 0);
-    assert_eq!(stats.levels[1].num_commits, 3);
+    assert_eq!(stats.levels[1].num_commits, 1);
+    assert_eq!(stats.levels[2].num_commits, 2);
 
     assert_eq!(generation_number(index.clone(), root_commit.id()), 0);
     assert_eq!(generation_number(index.clone(), commit_a.id()), 1);

lib/tests/test_transaction.rs

@@ -298,7 +298,7 @@ fn test_checkout_previous_empty_and_pruned(use_git: bool) {
 // #[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.
+    // after commit. It should also be indexed.
     let settings = testutils::user_settings();
     let (_temp_dir, mut repo) = testutils::init_repo(&settings, use_git);
 
@@ -311,11 +311,14 @@ fn test_add_head_success(use_git: bool) {
 
     let mut tx = repo.start_transaction("test");
     assert!(!tx.view().heads().contains(new_commit.id()));
+    assert!(!tx.as_repo_ref().index().has_id(new_commit.id()));
     tx.add_head(&new_commit);
     assert!(tx.view().heads().contains(new_commit.id()));
+    assert!(tx.as_repo_ref().index().has_id(new_commit.id()));
     tx.commit();
     Arc::get_mut(&mut repo).unwrap().reload();
     assert!(repo.view().heads().contains(new_commit.id()));
+    assert!(repo.index().has_id(new_commit.id()));
 }
 
 #[test_case(false ; "local store")]
@@ -381,7 +384,12 @@ fn test_add_head_not_immediate_child(use_git: bool) {
 // #[test_case(true ; "git store")]
 fn test_remove_head(use_git: bool) {
     // Test that Transaction::remove_head() removes the head, and that it's still
-    // removed after commit.
+    // removed after commit. It should remain in the index, since we otherwise would
+    // have to reindex everything.
+    // TODO: Consider if it's better to have the index be exactly the commits
+    // reachable from the view's heads. We would probably want to add tombstones
+    // for commits no longer visible in that case so we don't have to reindex e.g.
+    // when the user does `jj op undo`.
     let settings = testutils::user_settings();
     let (_temp_dir, mut repo) = testutils::init_repo(&settings, use_git);
 
@@ -403,6 +411,9 @@ fn test_remove_head(use_git: bool) {
     assert!(!heads.contains(commit3.id()));
     assert!(!heads.contains(commit2.id()));
     assert!(!heads.contains(commit1.id()));
+    assert!(tx.index().has_id(commit1.id()));
+    assert!(tx.index().has_id(commit2.id()));
+    assert!(tx.index().has_id(commit3.id()));
     tx.commit();
 
     Arc::get_mut(&mut repo).unwrap().reload();
@@ -410,6 +421,9 @@ fn test_remove_head(use_git: bool) {
     assert!(!heads.contains(commit3.id()));
     assert!(!heads.contains(commit2.id()));
     assert!(!heads.contains(commit1.id()));
+    assert!(repo.index().has_id(commit1.id()));
+    assert!(repo.index().has_id(commit2.id()));
+    assert!(repo.index().has_id(commit3.id()));
 }
 
 #[test_case(false ; "local store")]