Colocated workspaces, minimal

cli/src/cli_util.rs

@@ -386,6 +386,10 @@ impl CommandHelper {
         WorkspaceCommandHelper::new(ui, workspace, repo, env, self.is_at_head_operation())
     }
 
+    pub fn get_store_factories(&self) -> &StoreFactories {
+        &self.data.store_factories
+    }
+
     pub fn get_working_copy_factory(&self) -> Result<&dyn WorkingCopyFactory, CommandError> {
         let loader = self.workspace_loader()?;
 

cli/src/command_error.rs

@@ -304,6 +304,7 @@ impl From<WorkspaceInitError> for CommandError {
             }
             WorkspaceInitError::SignInit(err @ SignInitError::UnknownBackend(_)) => user_error(err),
             WorkspaceInitError::SignInit(err) => internal_error(err),
+            WorkspaceInitError::LoadNewlyCreated(err) => internal_error(err),
         }
     }
 }

cli/src/commands/workspace/add.rs

@@ -114,6 +114,7 @@ pub fn cmd_workspace_add(
         repo_path,
         repo,
         working_copy_factory,
+        command.get_store_factories(),
         workspace_id,
     )?;
     writeln!(

cli/tests/test_git_colocated.rs

@@ -791,6 +791,7 @@ fn get_log_output(test_env: &TestEnvironment, workspace_root: &Path) -> String {
       commit_id,
       bookmarks,
       if(git_head, "git_head()"),
+      working_copies,
       description,
     )
     "#;
@@ -806,6 +807,7 @@ fn get_log_output_with_stderr(
       commit_id,
       bookmarks,
       if(git_head, "git_head()"),
+      working_copies,
       description,
     )
     "#;
@@ -888,3 +890,34 @@ fn get_bookmark_output(test_env: &TestEnvironment, repo_path: &Path) -> String {
     // --quiet to suppress deleted bookmarks hint
     test_env.jj_cmd_success(repo_path, &["bookmark", "list", "--all-remotes", "--quiet"])
 }
+
+#[test]
+fn test_git_colocated_create_workspace_moving_wc() {
+    let test_env = TestEnvironment::default();
+    let repo_path = test_env.env_root().join("repo");
+    test_env.jj_cmd_ok(test_env.env_root(), &["git", "init", "--colocate", "repo"]);
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "second_wc_parent"]);
+    let second_wc_parent =
+        test_env.jj_cmd_success(&repo_path, &["log", "-Tcommit_id", "-r@-", "--no-graph"]);
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "should be git head"]);
+    test_env.jj_cmd_ok(
+        &repo_path,
+        &["workspace", "add", "../second", "-r", &second_wc_parent],
+    );
+
+    // The second workspace is not colocated. So creating it should not
+    // move git_head(). This tests whether we managed to reload the workspace
+    // command effectively in the middle of `jj workspace add`; if we did not,
+    // it will still think it's colocated in the default workspace.
+    //
+    // This test should survive switching to multi-git-head views, where each
+    // workspace gets its own git_head().
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r"
+    ○  3ee64208816f8264850a1cfcb99ac75e878c207a second@
+    │ @  076d128fc75cb291f908b77ffe9ff24c80de4fd7 default@
+    │ ○  005cf6c5b5e5fd032b804d98d11212a5317b4f3a git_head() should be git head
+    ├─╯
+    ○  ad5ff99a9ee8726ffe511430e4569f1d57bd550e second_wc_parent
+    ◆  0000000000000000000000000000000000000000
+    ");
+}

lib/src/workspace.rs

@@ -33,6 +33,7 @@ use crate::file_util::PathError;
 use crate::local_backend::LocalBackend;
 use crate::local_working_copy::LocalWorkingCopy;
 use crate::local_working_copy::LocalWorkingCopyFactory;
+use crate::op_store::OpStoreError;
 use crate::op_store::OperationId;
 use crate::op_store::WorkspaceId;
 use crate::repo::read_store_type;
@@ -45,6 +46,7 @@ use crate::repo::ReadonlyRepo;
 use crate::repo::Repo;
 use crate::repo::RepoInitError;
 use crate::repo::RepoLoader;
+use crate::repo::RepoLoaderError;
 use crate::repo::StoreFactories;
 use crate::repo::StoreLoadError;
 use crate::repo::SubmoduleStoreInitializer;
@@ -75,6 +77,18 @@ pub enum WorkspaceInitError {
     Backend(#[from] BackendInitError),
     #[error(transparent)]
     SignInit(#[from] SignInitError),
+    #[error("Could not load newly created workspace: {0}")]
+    LoadNewlyCreated(#[from] WorkspaceLoadNewlyCreatedError),
+}
+
+#[derive(Error, Debug)]
+pub enum WorkspaceLoadNewlyCreatedError {
+    #[error(transparent)]
+    LoadStore(#[from] StoreLoadError),
+    #[error(transparent)]
+    LoadOperation(#[from] OpStoreError),
+    #[error(transparent)]
+    LoadRepo(#[from] RepoLoaderError),
 }
 
 #[derive(Error, Debug)]
@@ -348,8 +362,9 @@ impl Workspace {
         user_settings: &UserSettings,
         workspace_root: &Path,
         repo_path: &Path,
-        repo: &Arc<ReadonlyRepo>,
+        existing_repo: &Arc<ReadonlyRepo>,
         working_copy_factory: &dyn WorkingCopyFactory,
+        store_factories: &StoreFactories,
         workspace_id: WorkspaceId,
     ) -> Result<(Self, Arc<ReadonlyRepo>), WorkspaceInitError> {
         let jj_dir = create_jj_dir(workspace_root)?;
@@ -366,20 +381,37 @@ impl Workspace {
             )
             .context(&repo_file_path)?;
 
+        // Load the new repo from the file we just wrote.
+        // This gives e.g. GitBackend has an opportunity to check if we
+        // are colocated, given we have a new workspace root.
+        let repo_loader = RepoLoader::init_from_file_system(
+            user_settings,
+            repo_path,
+            store_factories,
+            Some(workspace_root),
+        )
+        .map_err(WorkspaceLoadNewlyCreatedError::LoadStore)?;
+
+        // Now load a repo at the current operation. We go through
+        // OperationId as the Operation struct contains a reference
+        // to the old workspace's Store.
+        let current_op = repo_loader
+            .load_operation(existing_repo.operation().id())
+            .map_err(WorkspaceLoadNewlyCreatedError::LoadOperation)?;
+        let repo = repo_loader
+            .load_at(&current_op)
+            .map_err(WorkspaceLoadNewlyCreatedError::LoadRepo)?;
+
         let (working_copy, repo) = init_working_copy(
             user_settings,
-            repo,
+            &repo,
             workspace_root,
             &jj_dir,
             working_copy_factory,
             workspace_id,
         )?;
-        let workspace = Workspace::new(
-            workspace_root,
-            repo_dir,
-            working_copy,
-            repo.loader().clone(),
-        )?;
+        let workspace = Workspace::new(workspace_root, repo_dir, working_copy, repo_loader)?;
+
         Ok((workspace, repo))
     }
 

lib/tests/test_workspace.rs

@@ -57,6 +57,7 @@ fn test_init_additional_workspace() {
         test_workspace.repo_path(),
         &test_workspace.repo,
         &*default_working_copy_factory(),
+        &test_workspace.env.default_store_factories(),
         ws2_id.clone(),
     )
     .unwrap();