1423 lines · 52.8 KB · Rust
RawDownload
use std::{    future::Future,    path::{Path, PathBuf},    sync::{        Arc,        atomic::{AtomicU64, AtomicUsize},    },}; use anyhow::{Context as _, Result, anyhow, bail};use collections::HashMap;use fs::{Fs, copy_recursive};use futures::{FutureExt, future::Shared};use gpui::{    App, AppContext as _, AsyncApp, Context, Entity, EntityId, EventEmitter, Global, Task,    WeakEntity,};use itertools::Either;use postage::{prelude::Stream as _, watch};use rpc::{    AnyProtoClient, ErrorExt, TypedEnvelope,    proto::{self, REMOTE_SERVER_PROJECT_ID},};use text::ReplicaId;use util::{    ResultExt,    path_list::PathList,    paths::{PathStyle, RemotePathBuf, SanitizedPath},    rel_path::RelPath,};use worktree::{    CreatedEntry, Entry, ProjectEntryId, UpdatedEntriesSet, UpdatedGitRepositoriesSet, Worktree,    WorktreeId,}; use crate::{ProjectPath, trusted_worktrees::TrustedWorktrees}; /// The current paths for a project's worktrees. Each folder path has a corresponding/// main worktree path at the same position. The two lists are always the/// same length and are modified together via `add_path` / `remove_main_path`.////// For non-linked worktrees, the main path and folder path are identical./// For linked worktrees, the main path is the original repo and the folder/// path is the linked worktree location.#[derive(Default, Debug, Clone)]pub struct WorktreePaths {    paths: PathList,    main_paths: PathList,} impl PartialEq for WorktreePaths {    fn eq(&self, other: &Self) -> bool {        self.paths == other.paths && self.main_paths == other.main_paths    }} impl WorktreePaths {    /// Build from two parallel `PathList`s that already share the same    /// insertion order. Used for deserialization from DB.    ///    /// Returns an error if the two lists have different lengths, which    /// indicates corrupted data from a prior migration bug.    pub fn from_path_lists(        main_worktree_paths: PathList,        folder_paths: PathList,    ) -> anyhow::Result<Self> {        anyhow::ensure!(            main_worktree_paths.paths().len() == folder_paths.paths().len(),            "main_worktree_paths has {} entries but folder_paths has {}",            main_worktree_paths.paths().len(),            folder_paths.paths().len(),        );        Ok(Self {            paths: folder_paths,            main_paths: main_worktree_paths,        })    }     /// Build for non-linked worktrees where main == folder for every path.    pub fn from_folder_paths(folder_paths: &PathList) -> Self {        Self {            paths: folder_paths.clone(),            main_paths: folder_paths.clone(),        }    }     pub fn is_empty(&self) -> bool {        self.paths.is_empty()    }     /// The folder paths (for workspace matching / `threads_by_paths` index).    pub fn folder_path_list(&self) -> &PathList {        &self.paths    }     /// The main worktree paths (for group key / `threads_by_main_paths` index).    pub fn main_worktree_path_list(&self) -> &PathList {        &self.main_paths    }     /// Iterate the (main_worktree_path, folder_path) pairs in insertion order.    pub fn ordered_pairs(&self) -> impl Iterator<Item = (&PathBuf, &PathBuf)> {        self.main_paths            .ordered_paths()            .zip(self.paths.ordered_paths())    }     /// Add a new path pair. If the exact (main, folder) pair already exists,    /// this is a no-op. Rebuilds both internal `PathList`s to maintain    /// consistent ordering.    pub fn add_path(&mut self, main_path: &Path, folder_path: &Path) {        let already_exists = self            .ordered_pairs()            .any(|(m, f)| m.as_path() == main_path && f.as_path() == folder_path);        if already_exists {            return;        }        let (mut mains, mut folders): (Vec<PathBuf>, Vec<PathBuf>) = self            .ordered_pairs()            .map(|(m, f)| (m.clone(), f.clone()))            .unzip();        mains.push(main_path.to_path_buf());        folders.push(folder_path.to_path_buf());        self.main_paths = PathList::new(&mains);        self.paths = PathList::new(&folders);    }     /// Remove all pairs whose main worktree path matches the given path.    /// This removes the corresponding entries from both lists.    pub fn remove_main_path(&mut self, main_path: &Path) {        let (mains, folders): (Vec<PathBuf>, Vec<PathBuf>) = self            .ordered_pairs()            .filter(|(m, _)| m.as_path() != main_path)            .map(|(m, f)| (m.clone(), f.clone()))            .unzip();        self.main_paths = PathList::new(&mains);        self.paths = PathList::new(&folders);    }     /// Remove all pairs whose folder path matches the given path.    /// This removes the corresponding entries from both lists.    pub fn remove_folder_path(&mut self, folder_path: &Path) {        let (mains, folders): (Vec<PathBuf>, Vec<PathBuf>) = self            .ordered_pairs()            .filter(|(_, f)| f.as_path() != folder_path)            .map(|(m, f)| (m.clone(), f.clone()))            .unzip();        self.main_paths = PathList::new(&mains);        self.paths = PathList::new(&folders);    }} enum WorktreeStoreState {    Local {        fs: Arc<dyn Fs>,    },    Remote {        upstream_client: AnyProtoClient,        upstream_project_id: u64,        path_style: PathStyle,    },} #[derive(Clone)]pub struct WorktreeIdCounter(Arc<AtomicU64>); impl Default for WorktreeIdCounter {    fn default() -> Self {        Self(Arc::new(AtomicU64::new(1)))    }} impl WorktreeIdCounter {    pub fn get(cx: &mut App) -> Self {        cx.default_global::<Self>().clone()    }     fn next(&self) -> u64 {        self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed)    }} impl Global for WorktreeIdCounter {} pub struct WorktreeStore {    next_entry_id: Arc<AtomicUsize>,    next_worktree_id: WorktreeIdCounter,    downstream_client: Option<(AnyProtoClient, u64)>,    retain_worktrees: bool,    worktrees: Vec<WorktreeHandle>,    worktrees_reordered: bool,    scanning_enabled: bool,    #[allow(clippy::type_complexity)]    loading_worktrees:        HashMap<Arc<SanitizedPath>, Shared<Task<Result<Entity<Worktree>, Arc<anyhow::Error>>>>>,    initial_scan_complete: (watch::Sender<bool>, watch::Receiver<bool>),    state: WorktreeStoreState,} #[derive(Debug)]pub enum WorktreeStoreEvent {    WorktreeAdded(Entity<Worktree>),    WorktreeRemoved(EntityId, WorktreeId),    WorktreeReleased(EntityId, WorktreeId),    WorktreeOrderChanged,    WorktreeUpdateSent(Entity<Worktree>),    WorktreeUpdatedEntries(WorktreeId, UpdatedEntriesSet),    WorktreeUpdatedGitRepositories(WorktreeId, UpdatedGitRepositoriesSet),    WorktreeDeletedEntry(WorktreeId, ProjectEntryId),    WorktreeUpdatedRootRepoCommonDir(WorktreeId),} impl EventEmitter<WorktreeStoreEvent> for WorktreeStore {} impl WorktreeStore {    pub fn init(client: &AnyProtoClient) {        client.add_entity_request_handler(Self::handle_create_project_entry);        client.add_entity_request_handler(Self::handle_copy_project_entry);        client.add_entity_request_handler(Self::handle_delete_project_entry);        client.add_entity_request_handler(Self::handle_expand_project_entry);        client.add_entity_request_handler(Self::handle_expand_all_for_project_entry);    }     pub fn init_remote(client: &AnyProtoClient) {        client.add_entity_request_handler(Self::handle_allocate_worktree_id);    }     pub fn local(        retain_worktrees: bool,        fs: Arc<dyn Fs>,        next_worktree_id: WorktreeIdCounter,    ) -> Self {        Self {            next_entry_id: Default::default(),            next_worktree_id,            loading_worktrees: Default::default(),            downstream_client: None,            worktrees: Vec::new(),            worktrees_reordered: false,            scanning_enabled: true,            retain_worktrees,            initial_scan_complete: watch::channel_with(true),            state: WorktreeStoreState::Local { fs },        }    }     pub fn remote(        retain_worktrees: bool,        upstream_client: AnyProtoClient,        upstream_project_id: u64,        path_style: PathStyle,        next_worktree_id: WorktreeIdCounter,    ) -> Self {        Self {            next_entry_id: Default::default(),            next_worktree_id,            loading_worktrees: Default::default(),            downstream_client: None,            worktrees: Vec::new(),            worktrees_reordered: false,            scanning_enabled: true,            retain_worktrees,            initial_scan_complete: watch::channel_with(true),            state: WorktreeStoreState::Remote {                upstream_client,                upstream_project_id,                path_style,            },        }    }     pub fn next_worktree_id(&self) -> impl Future<Output = Result<WorktreeId>> + use<> {        let strategy = match (&self.state, &self.downstream_client) {            // we are a remote server, the client is in charge of assigning worktree ids            (WorktreeStoreState::Local { .. }, Some((client, REMOTE_SERVER_PROJECT_ID))) => {                Either::Left(client.clone())            }            // we are just a local zed project, we can assign ids            (WorktreeStoreState::Local { .. }, _) => Either::Right(self.next_worktree_id.next()),            // we are connected to a remote server, we are in charge of assigning worktree ids            (WorktreeStoreState::Remote { .. }, _) => Either::Right(self.next_worktree_id.next()),        };        async move {            match strategy {                Either::Left(client) => Ok(client                    .request(proto::AllocateWorktreeId {                        project_id: REMOTE_SERVER_PROJECT_ID,                    })                    .await?                    .worktree_id),                Either::Right(id) => Ok(id),            }            .map(WorktreeId::from_proto)        }    }     pub fn disable_scanner(&mut self) {        self.scanning_enabled = false;        *self.initial_scan_complete.0.borrow_mut() = true;    }     /// Returns a future that resolves when all visible worktrees have completed    /// their initial scan (entries populated, git repos detected).    pub fn wait_for_initial_scan(&self) -> impl Future<Output = ()> + use<> {        let mut rx = self.initial_scan_complete.1.clone();        async move {            let mut done = *rx.borrow();            while !done {                if let Some(value) = rx.recv().await {                    done = value;                } else {                    break;                }            }        }    }     /// Returns whether all visible worktrees have completed their initial scan.    pub fn initial_scan_completed(&self) -> bool {        *self.initial_scan_complete.1.borrow()    }     /// Checks whether all visible worktrees have completed their initial scan    /// and no worktree creations are pending, and updates the watch channel accordingly.    fn update_initial_scan_state(&mut self, cx: &App) {        let complete = self.loading_worktrees.is_empty()            && self                .visible_worktrees(cx)                .all(|wt| wt.read(cx).completed_scan_id() >= 1);        *self.initial_scan_complete.0.borrow_mut() = complete;    }     /// Spawns a detached task that waits for a worktree's initial scan to complete,    /// then rechecks and updates the aggregate initial scan state.    fn observe_worktree_scan_completion(        &mut self,        worktree: &Entity<Worktree>,        cx: &mut Context<Self>,    ) {        let await_scan = worktree.update(cx, |worktree, _cx| worktree.wait_for_snapshot(1));        cx.spawn(async move |this, cx| {            await_scan.await.ok();            this.update(cx, |this, cx| {                this.update_initial_scan_state(cx);            })            .ok();            anyhow::Ok(())        })        .detach();    }     /// Iterates through all worktrees, including ones that don't appear in the project panel    pub fn worktrees(&self) -> impl '_ + DoubleEndedIterator<Item = Entity<Worktree>> {        self.worktrees            .iter()            .filter_map(move |worktree| worktree.upgrade())    }     /// Iterates through all user-visible worktrees, the ones that appear in the project panel.    pub fn visible_worktrees<'a>(        &'a self,        cx: &'a App,    ) -> impl 'a + DoubleEndedIterator<Item = Entity<Worktree>> {        self.worktrees()            .filter(|worktree| worktree.read(cx).is_visible())    }     /// Iterates through all user-visible worktrees (directories and files that appear in the project panel) and other, invisible single files that could appear e.g. due to drag and drop.    pub fn visible_worktrees_and_single_files<'a>(        &'a self,        cx: &'a App,    ) -> impl 'a + DoubleEndedIterator<Item = Entity<Worktree>> {        self.worktrees()            .filter(|worktree| worktree.read(cx).is_visible() || worktree.read(cx).is_single_file())    }     pub fn worktree_for_id(&self, id: WorktreeId, cx: &App) -> Option<Entity<Worktree>> {        self.worktrees()            .find(|worktree| worktree.read(cx).id() == id)    }     pub fn worktree_for_entry(        &self,        entry_id: ProjectEntryId,        cx: &App,    ) -> Option<Entity<Worktree>> {        self.worktrees()            .find(|worktree| worktree.read(cx).contains_entry(entry_id))    }     pub fn find_worktree(        &self,        abs_path: impl AsRef<Path>,        cx: &App,    ) -> Option<(Entity<Worktree>, Arc<RelPath>)> {        let abs_path = SanitizedPath::new(abs_path.as_ref());        for tree in self.worktrees() {            let path_style = tree.read(cx).path_style();            if let Some(relative_path) =                path_style.strip_prefix(abs_path.as_ref(), tree.read(cx).abs_path().as_ref())            {                return Some((tree.clone(), relative_path.into_arc()));            }        }        None    }     pub fn project_path_for_absolute_path(&self, abs_path: &Path, cx: &App) -> Option<ProjectPath> {        self.find_worktree(abs_path, cx)            .map(|(worktree, relative_path)| ProjectPath {                worktree_id: worktree.read(cx).id(),                path: relative_path,            })    }     pub fn absolutize(&self, project_path: &ProjectPath, cx: &App) -> Option<PathBuf> {        let worktree = self.worktree_for_id(project_path.worktree_id, cx)?;        Some(worktree.read(cx).absolutize(&project_path.path))    }     pub fn path_style(&self) -> PathStyle {        match &self.state {            WorktreeStoreState::Local { .. } => PathStyle::local(),            WorktreeStoreState::Remote { path_style, .. } => *path_style,        }    }     pub fn find_or_create_worktree(        &mut self,        abs_path: impl AsRef<Path>,        visible: bool,        cx: &mut Context<Self>,    ) -> Task<Result<(Entity<Worktree>, Arc<RelPath>)>> {        let abs_path = abs_path.as_ref();        if let Some((tree, relative_path)) = self.find_worktree(abs_path, cx) {            Task::ready(Ok((tree, relative_path)))        } else {            let worktree = self.create_worktree(abs_path, visible, cx);            cx.background_spawn(async move { Ok((worktree.await?, RelPath::empty().into())) })        }    }     pub fn entry_for_id<'a>(&'a self, entry_id: ProjectEntryId, cx: &'a App) -> Option<&'a Entry> {        self.worktrees()            .find_map(|worktree| worktree.read(cx).entry_for_id(entry_id))    }     pub fn worktree_and_entry_for_id<'a>(        &'a self,        entry_id: ProjectEntryId,        cx: &'a App,    ) -> Option<(Entity<Worktree>, &'a Entry)> {        self.worktrees().find_map(|worktree| {            worktree                .read(cx)                .entry_for_id(entry_id)                .map(|e| (worktree.clone(), e))        })    }     pub fn entry_for_path<'a>(&'a self, path: &ProjectPath, cx: &'a App) -> Option<&'a Entry> {        self.worktree_for_id(path.worktree_id, cx)?            .read(cx)            .entry_for_path(&path.path)    }     pub fn copy_entry(        &mut self,        entry_id: ProjectEntryId,        new_project_path: ProjectPath,        cx: &mut Context<Self>,    ) -> Task<Result<Option<Entry>>> {        let Some(old_worktree) = self.worktree_for_entry(entry_id, cx) else {            return Task::ready(Err(anyhow!("no such worktree")));        };        let Some(old_entry) = old_worktree.read(cx).entry_for_id(entry_id) else {            return Task::ready(Err(anyhow!("no such entry")));        };        let Some(new_worktree) = self.worktree_for_id(new_project_path.worktree_id, cx) else {            return Task::ready(Err(anyhow!("no such worktree")));        };         match &self.state {            WorktreeStoreState::Local { fs } => {                let old_abs_path = old_worktree.read(cx).absolutize(&old_entry.path);                let new_abs_path = new_worktree.read(cx).absolutize(&new_project_path.path);                let fs = fs.clone();                let copy = cx.background_spawn(async move {                    copy_recursive(                        fs.as_ref(),                        &old_abs_path,                        &new_abs_path,                        Default::default(),                    )                    .await                });                 cx.spawn(async move |_, cx| {                    copy.await?;                    new_worktree

Lines 1–500 of 1423.