WIP: re-arranging the RepositoryEntry representation

Added branches to the randomized test to check the git branch
Added the remaining database integrations in collab

Co-authored-by: Max <max@zed.dev>
Co-authored-by: Petros <petros@zed.dev>

crates/collab/src/db.rs

@@ -1490,6 +1490,8 @@ impl Database {
                         visible: db_worktree.visible,
                         updated_entries: Default::default(),
                         removed_entries: Default::default(),
+                        updated_repositories: Default::default(),
+                        removed_repositories: Default::default(),
                         diagnostic_summaries: Default::default(),
                         scan_id: db_worktree.scan_id as u64,
                         completed_scan_id: db_worktree.completed_scan_id as u64,
@@ -1499,38 +1501,77 @@ impl Database {
                         .worktrees
                         .iter()
                         .find(|worktree| worktree.id == db_worktree.id as u64);
-                    let entry_filter = if let Some(rejoined_worktree) = rejoined_worktree {
-                        worktree_entry::Column::ScanId.gt(rejoined_worktree.scan_id)
-                    } else {
-                        worktree_entry::Column::IsDeleted.eq(false)
-                    };
 
-                    let mut db_entries = worktree_entry::Entity::find()
+                    // File entries
-                        .filter(
+                    {
-                            Condition::all()
+                        let entry_filter = if let Some(rejoined_worktree) = rejoined_worktree {
-                                .add(worktree_entry::Column::WorktreeId.eq(worktree.id))
+                            worktree_entry::Column::ScanId.gt(rejoined_worktree.scan_id)
-                                .add(entry_filter),
-                        )
-                        .stream(&*tx)
-                        .await?;
-
-                    while let Some(db_entry) = db_entries.next().await {
-                        let db_entry = db_entry?;
-                        if db_entry.is_deleted {
-                            worktree.removed_entries.push(db_entry.id as u64);
                         } else {
-                            worktree.updated_entries.push(proto::Entry {
+                            worktree_entry::Column::IsDeleted.eq(false)
-                                id: db_entry.id as u64,
+                        };
-                                is_dir: db_entry.is_dir,
+
-                                path: db_entry.path,
+                        let mut db_entries = worktree_entry::Entity::find()
-                                inode: db_entry.inode as u64,
+                            .filter(
-                                mtime: Some(proto::Timestamp {
+                                Condition::all()
-                                    seconds: db_entry.mtime_seconds as u64,
+                                    .add(worktree_entry::Column::WorktreeId.eq(worktree.id))
-                                    nanos: db_entry.mtime_nanos as u32,
+                                    .add(entry_filter),
-                                }),
+                            )
-                                is_symlink: db_entry.is_symlink,
+                            .stream(&*tx)
-                                is_ignored: db_entry.is_ignored,
+                            .await?;
-                            });
+
+                        while let Some(db_entry) = db_entries.next().await {
+                            let db_entry = db_entry?;
+                            if db_entry.is_deleted {
+                                worktree.removed_entries.push(db_entry.id as u64);
+                            } else {
+                                worktree.updated_entries.push(proto::Entry {
+                                    id: db_entry.id as u64,
+                                    is_dir: db_entry.is_dir,
+                                    path: db_entry.path,
+                                    inode: db_entry.inode as u64,
+                                    mtime: Some(proto::Timestamp {
+                                        seconds: db_entry.mtime_seconds as u64,
+                                        nanos: db_entry.mtime_nanos as u32,
+                                    }),
+                                    is_symlink: db_entry.is_symlink,
+                                    is_ignored: db_entry.is_ignored,
+                                });
+                            }
+                        }
+                    }
+
+                    // Repository Entries
+                    {
+                        let repository_entry_filter =
+                            if let Some(rejoined_worktree) = rejoined_worktree {
+                                worktree_repository::Column::ScanId.gt(rejoined_worktree.scan_id)
+                            } else {
+                                worktree_repository::Column::IsDeleted.eq(false)
+                            };
+
+                        let mut db_repositories = worktree_repository::Entity::find()
+                            .filter(
+                                Condition::all()
+                                    .add(worktree_repository::Column::WorktreeId.eq(worktree.id))
+                                    .add(repository_entry_filter),
+                            )
+                            .stream(&*tx)
+                            .await?;
+
+                        while let Some(db_repository) = db_repositories.next().await {
+                            let db_repository = db_repository?;
+                            if db_repository.is_deleted {
+                                worktree
+                                    .removed_repositories
+                                    .push(db_repository.dot_git_entry_id as u64);
+                            } else {
+                                worktree.updated_repositories.push(proto::RepositoryEntry {
+                                    dot_git_entry_id: db_repository.dot_git_entry_id as u64,
+                                    scan_id: db_repository.scan_id as u64,
+                                    work_directory: db_repository.work_directory_path,
+                                    branch: db_repository.branch,
+                                });
+                            }
                         }
                     }
 
@@ -2555,6 +2596,7 @@ impl Database {
                             root_name: db_worktree.root_name,
                             visible: db_worktree.visible,
                             entries: Default::default(),
+                            repository_entries: Default::default(),
                             diagnostic_summaries: Default::default(),
                             scan_id: db_worktree.scan_id as u64,
                             completed_scan_id: db_worktree.completed_scan_id as u64,
@@ -2592,6 +2634,31 @@ impl Database {
                 }
             }
 
+            // Populate repository entries.
+            {
+                let mut db_repository_entries = worktree_repository::Entity::find()
+                    .filter(
+                        Condition::all()
+                            .add(worktree_repository::Column::ProjectId.eq(project_id))
+                            .add(worktree_repository::Column::IsDeleted.eq(false)),
+                    )
+                    .stream(&*tx)
+                    .await?;
+                while let Some(db_repository_entry) = db_repository_entries.next().await {
+                    let db_repository_entry = db_repository_entry?;
+                    if let Some(worktree) =
+                        worktrees.get_mut(&(db_repository_entry.worktree_id as u64))
+                    {
+                        worktree.repository_entries.push(proto::RepositoryEntry {
+                            dot_git_entry_id: db_repository_entry.dot_git_entry_id as u64,
+                            scan_id: db_repository_entry.scan_id as u64,
+                            work_directory: db_repository_entry.work_directory_path,
+                            branch: db_repository_entry.branch,
+                        });
+                    }
+                }
+            }
+
             // Populate worktree diagnostic summaries.
             {
                 let mut db_summaries = worktree_diagnostic_summary::Entity::find()
@@ -3273,6 +3340,8 @@ pub struct RejoinedWorktree {
     pub visible: bool,
     pub updated_entries: Vec<proto::Entry>,
     pub removed_entries: Vec<u64>,
+    pub updated_repositories: Vec<proto::RepositoryEntry>,
+    pub removed_repositories: Vec<u64>,
     pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
     pub scan_id: u64,
     pub completed_scan_id: u64,
@@ -3327,6 +3396,7 @@ pub struct Worktree {
     pub root_name: String,
     pub visible: bool,
     pub entries: Vec<proto::Entry>,
+    pub repository_entries: Vec<proto::RepositoryEntry>,
     pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
     pub scan_id: u64,
     pub completed_scan_id: u64,

crates/collab/src/rpc.rs

@@ -1386,9 +1386,8 @@ async fn join_project(
             removed_entries: Default::default(),
             scan_id: worktree.scan_id,
             is_last_update: worktree.scan_id == worktree.completed_scan_id,
-            // TODO repo
+            updated_repositories: worktree.repository_entries,
-            updated_repositories: vec![],
+            removed_repositories: Default::default(),
-            removed_repositories: vec![],
         };
         for update in proto::split_worktree_update(message, MAX_CHUNK_SIZE) {
             session.peer.send(session.connection_id, update.clone())?;

crates/collab/src/tests/randomized_integration_tests.rs

@@ -785,6 +785,28 @@ async fn apply_client_operation(
             }
             client.fs.set_index_for_repo(&dot_git_dir, &contents).await;
         }
+
+        ClientOperation::WriteGitBranch {
+            repo_path,
+            new_branch,
+        } => {
+            if !client.fs.directories().contains(&repo_path) {
+                return Err(TestError::Inapplicable);
+            }
+
+            log::info!(
+                "{}: writing git branch for repo {:?}: {:?}",
+                client.username,
+                repo_path,
+                new_branch
+            );
+
+            let dot_git_dir = repo_path.join(".git");
+            if client.fs.metadata(&dot_git_dir).await?.is_none() {
+                client.fs.create_dir(&dot_git_dir).await?;
+            }
+            client.fs.set_branch_name(&dot_git_dir, new_branch).await;
+        }
     }
     Ok(())
 }
@@ -859,6 +881,12 @@ fn check_consistency_between_clients(clients: &[(Rc<TestClient>, TestAppContext)
                                 host_snapshot.abs_path(),
                                 guest_project.remote_id(),
                             );
+                            assert_eq!(guest_snapshot.repositories().collect::<Vec<_>>(), host_snapshot.repositories().collect::<Vec<_>>(),
+                                "{} has different repositories than the host for worktree {:?} and project {:?}",
+                                client.username,
+                                host_snapshot.abs_path(),
+                                guest_project.remote_id(),
+                            );
                             assert_eq!(guest_snapshot.scan_id(), host_snapshot.scan_id(),
                                 "{} has different scan id than the host for worktree {:?} and project {:?}",
                                 client.username,
@@ -1151,6 +1179,10 @@ enum ClientOperation {
         repo_path: PathBuf,
         contents: Vec<(PathBuf, String)>,
     },
+    WriteGitBranch {
+        repo_path: PathBuf,
+        new_branch: Option<String>,
+    },
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -1664,7 +1696,7 @@ impl TestPlan {
                 }
 
                 // Update a git index
-                91..=95 => {
+                91..=93 => {
                     let repo_path = client
                         .fs
                         .directories()
@@ -1698,6 +1730,24 @@ impl TestPlan {
                     };
                 }
 
+                // Update a git branch
+                94..=95 => {
+                    let repo_path = client
+                        .fs
+                        .directories()
+                        .choose(&mut self.rng)
+                        .unwrap()
+                        .clone();
+
+                    let new_branch = (self.rng.gen_range(0..10) > 3)
+                        .then(|| Alphanumeric.sample_string(&mut self.rng, 8));
+
+                    break ClientOperation::WriteGitBranch {
+                        repo_path,
+                        new_branch,
+                    };
+                }
+
                 // Create or update a file or directory
                 96.. => {
                     let is_dir = self.rng.gen::<bool>();

crates/project/src/worktree.rs

@@ -117,13 +117,10 @@ pub struct Snapshot {
     completed_scan_id: usize,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub struct RepositoryEntry {
     pub(crate) scan_id: usize,
-    pub(crate) dot_git_entry_id: ProjectEntryId,
+    pub(crate) work_directory_id: ProjectEntryId,
-    /// Relative to the worktree, the repository for the root will have
-    /// a work directory equal to: ""
-    pub(crate) work_directory: RepositoryWorkDirectory,
     pub(crate) branch: Option<Arc<str>>,
 }
 
@@ -132,17 +129,16 @@ impl RepositoryEntry {
         self.branch.clone()
     }
 
-    pub fn work_directory(&self) -> Arc<Path> {
+    pub fn work_directory_id(&self) -> ProjectEntryId {
-        self.work_directory.0.clone()
+        self.work_directory_id
     }
 }
 
 impl From<&RepositoryEntry> for proto::RepositoryEntry {
     fn from(value: &RepositoryEntry) -> Self {
         proto::RepositoryEntry {
-            dot_git_entry_id: value.dot_git_entry_id.to_proto(),
             scan_id: value.scan_id as u64,
-            work_directory: value.work_directory.to_string_lossy().to_string(),
+            work_directory_id: value.work_directory_id.to_proto(),
             branch: value.branch.as_ref().map(|str| str.to_string()),
         }
     }
@@ -212,6 +208,7 @@ impl AsRef<Path> for RepositoryWorkDirectory {
 pub struct LocalSnapshot {
     ignores_by_parent_abs_path: HashMap<Arc<Path>, (Arc<Gitignore>, usize)>,
     // The ProjectEntryId corresponds to the entry for the .git dir
+    // work_directory_id
     git_repositories: TreeMap<ProjectEntryId, LocalRepositoryEntry>,
     removed_entry_ids: HashMap<u64, ProjectEntryId>,
     next_entry_id: Arc<AtomicUsize>,
@@ -701,12 +698,12 @@ impl LocalWorktree {
         ) {
             for a_repo in a {
                 let matched = b.find(|b_repo| {
-                    a_repo.dot_git_entry_id == b_repo.dot_git_entry_id
+                    a_repo.work_directory_id == b_repo.work_directory_id
                         && a_repo.scan_id == b_repo.scan_id
                 });
 
                 if matched.is_none() {
-                    updated.insert(a_repo.dot_git_entry_id, a_repo.clone());
+                    updated.insert(a_repo.work_directory_id, a_repo.clone());
                 }
             }
         }
@@ -1158,7 +1155,7 @@ impl LocalWorktree {
         repo_path: RepoPath,
         cx: &mut ModelContext<Worktree>,
     ) -> Task<Option<String>> {
-        let Some(git_ptr) = self.git_repositories.get(&repo.dot_git_entry_id).map(|git_ptr| git_ptr.to_owned()) else {
+        let Some(git_ptr) = self.git_repositories.get(&repo.work_directory_id).map(|git_ptr| git_ptr.to_owned()) else {
             return Task::Ready(Some(None))
         };
         let git_ptr = git_ptr.repo_ptr;
@@ -1476,6 +1473,10 @@ impl Snapshot {
         self.traverse_from_offset(true, include_ignored, 0)
     }
 
+    pub fn repositories(&self) -> impl Iterator<Item = &RepositoryEntry> {
+        self.repository_entries.values()
+    }
+
     pub fn paths(&self) -> impl Iterator<Item = &Arc<Path>> {
         let empty_path = Path::new("");
         self.entries_by_path

crates/rpc/proto/zed.proto

@@ -982,10 +982,9 @@ message Entry {
 }
 
 message RepositoryEntry {
-    uint64 dot_git_entry_id = 1;
+    uint64 scan_id = 1;
-    uint64 scan_id = 2;
+    uint64 work_directory_id = 2;
-    string work_directory = 4;
+    optional string branch = 3;
-    optional string branch = 5;
 }
 
 message BufferState {