feat(utils): Rework locate_workspace, introduce Workspace type (#2830)

## What

- Removes `locate_workspace` and `workspace_dir_or_current_dir` methods.
- Instead, introduces `Workspace` type that is aware of different Cargo
workspaces in the codebase.

## Why

The approach with a single `locate_workspace` doesn't work well for our
codebase, since we have multiple workspaces. It resulted in some very
implicit and convoluted code (see the removed `get_base_dir` in prover
workspace).

New approach handles all 3 workspaces _plus_ the lack of a workspace.

core/bin/contract-verifier/src/main.rs

@@ -9,14 +9,14 @@ use zksync_contract_verifier_lib::ContractVerifier;
 use zksync_core_leftovers::temp_config_store::{load_database_secrets, load_general_config};
 use zksync_dal::{ConnectionPool, Core, CoreDal};
 use zksync_queued_job_processor::JobProcessor;
-use zksync_utils::{wait_for_tasks::ManagedTasks, workspace_dir_or_current_dir};
+use zksync_utils::{env::Workspace, wait_for_tasks::ManagedTasks};
 use zksync_vlog::prometheus::PrometheusExporterConfig;
 
 async fn update_compiler_versions(connection_pool: &ConnectionPool<Core>) {
     let mut storage = connection_pool.connection().await.unwrap();
     let mut transaction = storage.start_transaction().await.unwrap();
 
-    let zksync_home = workspace_dir_or_current_dir();
+    let zksync_home = Workspace::locate().core();
 
     let zksolc_path = zksync_home.join("etc/zksolc-bin/");
     let zksolc_versions: Vec<String> = std::fs::read_dir(zksolc_path)

core/bin/system-constants-generator/src/main.rs

@@ -17,7 +17,7 @@ use zksync_types::{
     IntrinsicSystemGasConstants, ProtocolVersionId, GUARANTEED_PUBDATA_IN_TX,
     L1_GAS_PER_PUBDATA_BYTE, MAX_NEW_FACTORY_DEPS, REQUIRED_L1_TO_L2_GAS_PER_PUBDATA_BYTE,
 };
-use zksync_utils::workspace_dir_or_current_dir;
+use zksync_utils::env::Workspace;
 
 // For configs we will use the default value of `800_000` to represent the rough amount of L1 gas
 // needed to cover the batch expenses.
@@ -210,7 +210,7 @@ fn generate_rust_fee_constants(intrinsic_gas_constants: &IntrinsicSystemGasConst
 }
 
 fn save_file(path_in_repo: &str, content: String) {
-    let zksync_home = workspace_dir_or_current_dir();
+    let zksync_home = Workspace::locate().core();
     let fee_constants_path = zksync_home.join(path_in_repo);
 
     fs::write(fee_constants_path, content)

core/lib/contract_verifier/src/lib.rs

@@ -1,6 +1,6 @@
 use std::{
     collections::HashMap,
-    path::Path,
+    path::{Path, PathBuf},
     time::{Duration, Instant},
 };
 
@@ -20,7 +20,7 @@ use zksync_types::{
     },
     Address,
 };
-use zksync_utils::workspace_dir_or_current_dir;
+use zksync_utils::env::Workspace;
 
 use crate::{
     error::ContractVerifierError,
@@ -38,8 +38,8 @@ lazy_static! {
     static ref DEPLOYER_CONTRACT: Contract = zksync_contracts::deployer_contract();
 }
 
-fn home_path() -> &'static Path {
-    workspace_dir_or_current_dir()
+fn home_path() -> PathBuf {
+    Workspace::locate().core()
 }
 
 #[derive(Debug)]

core/lib/contracts/src/lib.rs

@@ -16,7 +16,7 @@ use ethabi::{
 };
 use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
-use zksync_utils::{bytecode::hash_bytecode, bytes_to_be_words, workspace_dir_or_current_dir};
+use zksync_utils::{bytecode::hash_bytecode, bytes_to_be_words, env::Workspace};
 
 pub mod test_contracts;
 
@@ -64,8 +64,8 @@ const LOADNEXT_CONTRACT_FILE: &str =
 const LOADNEXT_SIMPLE_CONTRACT_FILE: &str =
     "etc/contracts-test-data/artifacts-zk/contracts/loadnext/loadnext_contract.sol/Foo.json";
 
-fn home_path() -> &'static Path {
-    workspace_dir_or_current_dir()
+fn home_path() -> PathBuf {
+    Workspace::locate().core()
 }
 
 fn read_file_to_json_value(path: impl AsRef<Path> + std::fmt::Debug) -> serde_json::Value {

core/lib/utils/Cargo.toml

@@ -32,3 +32,4 @@ once_cell.workspace = true
 rand.workspace = true
 tokio = { workspace = true, features = ["macros", "rt"] }
 bincode.workspace = true
+assert_matches.workspace = true

core/lib/utils/src/env.rs

@@ -8,6 +8,87 @@ use once_cell::sync::OnceCell;
 
 static WORKSPACE: OnceCell<Option<PathBuf>> = OnceCell::new();
 
+/// Represents Cargo workspaces available in the repository.
+#[derive(Debug, Clone, Copy)]
+pub enum Workspace<'a> {
+    /// Workspace was not found.
+    /// Assumes that the code is running in a binary.
+    /// Will use the current directory as a fallback.
+    None,
+    /// Root folder.
+    Core(&'a Path),
+    /// `prover` folder.
+    Prover(&'a Path),
+    /// `toolbox` folder.
+    Toolbox(&'a Path),
+}
+
+impl Workspace<'static> {
+    /// Find the location of the current workspace, if this code works in workspace
+    /// then it will return the correct folder if, it's binary e.g. in docker container
+    /// you have to use fallback to another directory
+    /// The code has been inspired by `insta`
+    /// `https://github.com/mitsuhiko/insta/blob/master/insta/src/env.rs`
+    pub fn locate() -> Self {
+        // Since `locate_workspace_inner()` should be deterministic, it makes little sense to call
+        // `OnceCell::get_or_try_init()` here; the repeated calls are just as unlikely to succeed as the initial call.
+        // Instead, we store `None` in the `OnceCell` if initialization failed.
+        let path: Option<&'static Path> = WORKSPACE
+            .get_or_init(|| {
+                let result = locate_workspace_inner();
+                // If the workspace is not found, we store `None` in the `OnceCell`.
+                // It doesn't make sense to log it, since in most production cases the workspace
+                // is not present.
+                result.ok()
+            })
+            .as_deref();
+        path.map_or(Self::None, Self::from)
+    }
+}
+
+impl<'a> Workspace<'a> {
+    const PROVER_DIRECTORY_NAME: &'static str = "prover";
+    const TOOLBOX_DIRECTORY_NAME: &'static str = "zk_toolbox";
+
+    /// Returns the path of the core workspace.
+    /// For `Workspace::None`, considers the current directory to represent core workspace.
+    pub fn core(self) -> PathBuf {
+        match self {
+            Self::None => PathBuf::from("."),
+            Self::Core(path) => path.into(),
+            Self::Prover(path) | Self::Toolbox(path) => path.parent().unwrap().into(),
+        }
+    }
+
+    /// Returns the path of the `prover` workspace.
+    pub fn prover(self) -> PathBuf {
+        match self {
+            Self::Prover(path) => path.into(),
+            _ => self.core().join(Self::PROVER_DIRECTORY_NAME),
+        }
+    }
+
+    /// Returns the path of the `zk_toolbox`` workspace.
+    pub fn toolbox(self) -> PathBuf {
+        match self {
+            Self::Toolbox(path) => path.into(),
+            _ => self.core().join(Self::TOOLBOX_DIRECTORY_NAME),
+        }
+    }
+}
+
+impl<'a> From<&'a Path> for Workspace<'a> {
+    fn from(path: &'a Path) -> Self {
+        if path.ends_with(Self::PROVER_DIRECTORY_NAME) {
+            Self::Prover(path)
+        } else if path.ends_with(Self::TOOLBOX_DIRECTORY_NAME) {
+            Self::Toolbox(path)
+        } else {
+            Self::Core(path)
+        }
+    }
+}
+
 fn locate_workspace_inner() -> anyhow::Result<PathBuf> {
     let output = std::process::Command::new(
         std::env::var("CARGO")
@@ -40,31 +121,86 @@ fn locate_workspace_inner() -> anyhow::Result<PathBuf> {
         .to_path_buf())
 }
 
-/// Find the location of the current workspace, if this code works in workspace
-/// then it will return the correct folder if, it's binary e.g. in docker container
-/// you have to use fallback to another directory
-/// The code has been inspired by `insta`
-/// `https://github.com/mitsuhiko/insta/blob/master/insta/src/env.rs`
-pub fn locate_workspace() -> Option<&'static Path> {
-    // Since `locate_workspace_inner()` should be deterministic, it makes little sense to call
-    // `OnceCell::get_or_try_init()` here; the repeated calls are just as unlikely to succeed as the initial call.
-    // Instead, we store `None` in the `OnceCell` if initialization failed.
-    WORKSPACE
-        .get_or_init(|| {
-            let result = locate_workspace_inner();
-            if result.is_err() {
-                // `get_or_init()` is guaranteed to call the provided closure once per `OnceCell`;
-                // i.e., we won't spam logs here.
-                tracing::info!(
-                    "locate_workspace() failed. You are using an already compiled version"
-                );
-            }
-            result.ok()
-        })
-        .as_deref()
-}
+#[cfg(test)]
+mod tests {
+    use assert_matches::assert_matches;
+
+    use super::*;
+
+    /// Will reset the pwd on drop.
+    /// This is needed to make sure that even if the test fails, the env
+    /// for other tests is left intact.
+    struct PwdProtector(PathBuf);
+
+    impl PwdProtector {
+        fn new() -> Self {
+            let pwd = std::env::current_dir().unwrap();
+            Self(pwd)
+        }
+    }
+
+    impl Drop for PwdProtector {
+        fn drop(&mut self) {
+            std::env::set_current_dir(self.0.clone()).unwrap();
+        }
+    }
+
+    #[test]
+    fn test_workspace_locate() {
+        let _pwd_protector = PwdProtector::new();
+
+        // Core.
+
+        let workspace = Workspace::locate();
+        assert_matches!(workspace, Workspace::Core(_));
+        let core_path = workspace.core();
+        // Check if prover and toolbox directories exist.
+        assert!(workspace.prover().exists());
+        assert_matches!(
+            Workspace::from(workspace.prover().as_path()),
+            Workspace::Prover(_)
+        );
+        assert!(workspace.toolbox().exists());
+        assert_matches!(
+            Workspace::from(workspace.toolbox().as_path()),
+            Workspace::Toolbox(_)
+        );
+
+        // Prover.
+
+        // We use `cargo-nextest` for running tests, which runs each test in parallel,
+        // so we can safely alter the global env, assuming that we will restore it after
+        // the test.
+        std::env::set_current_dir(workspace.prover()).unwrap();
+        let workspace_path = locate_workspace_inner().unwrap();
+        let workspace = Workspace::from(workspace_path.as_path());
+        assert_matches!(workspace, Workspace::Prover(_));
+        let prover_path = workspace.prover();
+        assert_eq!(workspace.core(), core_path);
+        assert_matches!(
+            Workspace::from(workspace.core().as_path()),
+            Workspace::Core(_)
+        );
+        assert!(workspace.toolbox().exists());
+        assert_matches!(
+            Workspace::from(workspace.toolbox().as_path()),
+            Workspace::Toolbox(_)
+        );
 
-/// Returns [`locate_workspace()`] output with the "." fallback.
-pub fn workspace_dir_or_current_dir() -> &'static Path {
-    locate_workspace().unwrap_or_else(|| Path::new("."))
+        // Toolbox.
+        std::env::set_current_dir(workspace.toolbox()).unwrap();
+        let workspace_path = locate_workspace_inner().unwrap();
+        let workspace = Workspace::from(workspace_path.as_path());
+        assert_matches!(workspace, Workspace::Toolbox(_));
+        assert_eq!(workspace.core(), core_path);
+        assert_matches!(
+            Workspace::from(workspace.core().as_path()),
+            Workspace::Core(_)
+        );
+        assert_eq!(workspace.prover(), prover_path);
+        assert_matches!(
+            Workspace::from(workspace.prover().as_path()),
+            Workspace::Prover(_)
+        );
+    }
 }

core/lib/utils/src/lib.rs

@@ -2,12 +2,12 @@
 
 pub mod bytecode;
 mod convert;
-mod env;
+pub mod env;
 pub mod http_with_retries;
 pub mod misc;
 pub mod panic_extractor;
 mod serde_wrappers;
 pub mod time;
 pub mod wait_for_tasks;
 
-pub use self::{convert::*, env::*, misc::*, serde_wrappers::*};
+pub use self::{convert::*, misc::*, serde_wrappers::*};

core/tests/loadnext/src/config.rs

@@ -4,7 +4,7 @@ use serde::Deserialize;
 use tokio::sync::Semaphore;
 use zksync_contracts::test_contracts::LoadnextContractExecutionParams;
 use zksync_types::{network::Network, Address, L2ChainId, H160};
-use zksync_utils::workspace_dir_or_current_dir;
+use zksync_utils::env::Workspace;
 
 use crate::fs_utils::read_tokens;
 
@@ -190,7 +190,7 @@ fn default_main_token() -> H160 {
 }
 
 fn default_test_contracts_path() -> PathBuf {
-    let test_contracts_path = workspace_dir_or_current_dir().join("etc/contracts-test-data");
+    let test_contracts_path = Workspace::locate().core().join("etc/contracts-test-data");
     tracing::info!("Test contracts path: {}", test_contracts_path.display());
     test_contracts_path
 }

core/tests/loadnext/src/fs_utils.rs

@@ -5,7 +5,7 @@ use std::{fs::File, io::BufReader, path::Path};
 
 use serde::Deserialize;
 use zksync_types::{ethabi::Contract, network::Network, Address};
-use zksync_utils::workspace_dir_or_current_dir;
+use zksync_utils::env::Workspace;
 
 /// A token stored in `etc/tokens/{network}.json` files.
 #[derive(Debug, Deserialize)]
@@ -27,7 +27,7 @@ pub struct TestContract {
 }
 
 pub fn read_tokens(network: Network) -> anyhow::Result<Vec<Token>> {
-    let home = workspace_dir_or_current_dir();
+    let home = Workspace::locate().core();
     let path = home.join(format!("etc/tokens/{network}.json"));
     let file = File::open(path)?;
     let reader = BufReader::new(file);

prover/crates/bin/prover_cli/src/config/mod.rs

@@ -1,12 +1,12 @@
 use std::{io::Write, path::PathBuf};
 
-use crate::helper::core_workspace_dir_or_current_dir;
+use zksync_utils::env::Workspace;
 
 pub fn get_envfile() -> anyhow::Result<PathBuf> {
     if let Ok(envfile) = std::env::var("PLI__CONFIG") {
         return Ok(envfile.into());
     }
-    Ok(core_workspace_dir_or_current_dir().join("etc/pliconfig"))
+    Ok(Workspace::locate().core().join("etc/pliconfig"))
 }
 
 pub fn load_envfile(path: impl AsRef<std::path::Path>) -> anyhow::Result<()> {