multi_sequence.rs

use alloy_network::Network;
use eyre::{ContextCompat, Result, WrapErr};
use forge_script_sequence::{
    DRY_RUN_DIR, ScriptSequence, SensitiveScriptSequence, now, sig_to_file_name,
};
use foundry_common::{fs, shell};
use foundry_compilers::ArtifactId;
use foundry_config::Config;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;

/// Holds the sequences of multiple chain deployments.
#[derive(Clone, Default, Serialize, Deserialize)]
#[serde(bound(
    serialize = "N::TransactionRequest: Serialize, N::TxEnvelope: Serialize",
    deserialize = "N::TransactionRequest: for<'de2> Deserialize<'de2>, N::TxEnvelope: for<'de2> Deserialize<'de2>"
))]
pub struct MultiChainSequence<N: Network> {
    pub deployments: Vec<ScriptSequence<N>>,
    #[serde(skip)]
    pub path: PathBuf,
    #[serde(skip)]
    pub sensitive_path: PathBuf,
    pub timestamp: u128,
}

/// Sensitive values from script sequences.
#[derive(Clone, Default, Serialize, Deserialize)]
pub struct SensitiveMultiChainSequence {
    pub deployments: Vec<SensitiveScriptSequence>,
}

impl SensitiveMultiChainSequence {
    fn from_multi_sequence<N: Network>(sequence: &MultiChainSequence<N>) -> Self {
        Self {
            deployments: sequence.deployments.iter().map(SensitiveScriptSequence::from).collect(),
        }
    }
}

impl<N: Network> MultiChainSequence<N> {
    pub fn new(
        deployments: Vec<ScriptSequence<N>>,
        sig: &str,
        target: &ArtifactId,
        config: &Config,
        dry_run: bool,
    ) -> Result<Self> {
        let (path, sensitive_path) = Self::get_paths(config, sig, target, dry_run)?;

        Ok(Self { deployments, path, sensitive_path, timestamp: now().as_millis() })
    }

    /// Gets paths in the formats
    /// ./broadcast/multi/contract_filename[-timestamp]/sig.json and
    /// ./cache/multi/contract_filename[-timestamp]/sig.json
    pub fn get_paths(
        config: &Config,
        sig: &str,
        target: &ArtifactId,
        dry_run: bool,
    ) -> Result<(PathBuf, PathBuf)> {
        let mut broadcast = config.broadcast.clone();
        let mut cache = config.cache_path.clone();
        let mut common = PathBuf::new();

        common.push("multi");

        if dry_run {
            common.push(DRY_RUN_DIR);
        }

        let target_fname = target
            .source
            .file_name()
            .wrap_err_with(|| format!("No filename for {:?}", target.source))?
            .to_string_lossy();

        common.push(format!("{target_fname}-latest"));

        broadcast.push(common.clone());
        cache.push(common);

        fs::create_dir_all(&broadcast)?;
        fs::create_dir_all(&cache)?;

        let filename = format!("{}.json", sig_to_file_name(sig));

        broadcast.push(filename.clone());
        cache.push(filename);

        Ok((broadcast, cache))
    }

    /// Loads the sequences for the multi chain deployment.
    pub fn load(config: &Config, sig: &str, target: &ArtifactId, dry_run: bool) -> Result<Self>
    where
        N::TxEnvelope: for<'d> Deserialize<'d>,
    {
        let (path, sensitive_path) = Self::get_paths(config, sig, target, dry_run)?;
        let mut sequence: Self = foundry_compilers::utils::read_json_file(&path)
            .wrap_err("Multi-chain deployment not found.")?;
        let sensitive_sequence: SensitiveMultiChainSequence =
            foundry_compilers::utils::read_json_file(&sensitive_path)
                .wrap_err("Multi-chain deployment sensitive details not found.")?;

        let deployments_len = sequence.deployments.len();
        let sensitive_deployments_len = sensitive_sequence.deployments.len();
        if deployments_len != sensitive_deployments_len {
            eyre::bail!(
                "sensitive-cache deployment count ({sensitive_deployments_len}) does not match \
                 deployment count ({deployments_len}); the multi-chain deployment and its \
                 sensitive-cache counterpart are out of sync"
            );
        }
        for (i, deployment) in sequence.deployments.iter_mut().enumerate() {
            deployment.fill_sensitive(&sensitive_sequence.deployments[i])?;
        }

        sequence.path = path;
        sequence.sensitive_path = sensitive_path;

        Ok(sequence)
    }

    /// Saves the transactions as file if it's a standalone deployment.
    pub fn save(&mut self, silent: bool, save_ts: bool) -> Result<()>
    where
        N::TxEnvelope: Serialize,
    {
        self.deployments.iter_mut().for_each(|sequence| sequence.sort_receipts());

        self.timestamp = now().as_millis();

        let sensitive_sequence = SensitiveMultiChainSequence::from_multi_sequence(&*self);

        // broadcast writes
        //../Contract-latest/run.json
        fs::write_pretty_json_file(&self.path, self)?;

        if save_ts {
            //../Contract-[timestamp]/run.json
            let path = self.path.to_string_lossy();
            let file = PathBuf::from(&path.replace("-latest", &format!("-{}", self.timestamp)));
            fs::create_dir_all(file.parent().unwrap())?;
            fs::copy(&self.path, &file)?;
        }

        // cache writes
        //../Contract-latest/run.json
        fs::write_sensitive_json_file(&self.sensitive_path, &sensitive_sequence)?;

        if save_ts {
            //../Contract-[timestamp]/run.json
            let path = self.sensitive_path.to_string_lossy();
            let file = PathBuf::from(&path.replace("-latest", &format!("-{}", self.timestamp)));
            fs::create_dir_all(file.parent().unwrap())?;
            fs::copy(&self.sensitive_path, &file)?;
        }

        if !silent {
            if shell::is_json() {
                sh_println!(
                    "{}",
                    serde_json::json!({
                        "status": "success",
                        "transactions": self.path.display().to_string(),
                        "sensitive": self.sensitive_path.display().to_string(),
                    })
                )?;
            } else {
                sh_println!("\nTransactions saved to: {}\n", self.path.display())?;
                sh_println!("Sensitive details saved to: {}\n", self.sensitive_path.display())?;
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloy_network::Ethereum;

    #[test]
    fn load_rejects_mismatched_deployment_counts() {
        let dir = tempfile::tempdir().unwrap();
        let config = Config {
            broadcast: dir.path().join("broadcast"),
            cache_path: dir.path().join("cache"),
            ..Default::default()
        };
        let target = ArtifactId {
            path: PathBuf::from("Script.json"),
            name: "Script".to_string(),
            source: PathBuf::from("Script.sol"),
            version: "0.8.30".parse().unwrap(),
            build_id: String::new(),
            profile: "default".to_string(),
        };
        let (path, sensitive_path) =
            MultiChainSequence::<Ethereum>::get_paths(&config, "run()", &target, false).unwrap();
        let sequence = MultiChainSequence::<Ethereum> {
            deployments: vec![ScriptSequence::default()],
            path: PathBuf::new(),
            sensitive_path: PathBuf::new(),
            timestamp: 0,
        };
        fs::write_pretty_json_file(&path, &sequence).unwrap();
        for count in [0, 2] {
            let sensitive = SensitiveMultiChainSequence {
                deployments: vec![SensitiveScriptSequence::default(); count],
            };
            fs::write_sensitive_json_file(&sensitive_path, &sensitive).unwrap();
            let err = MultiChainSequence::<Ethereum>::load(&config, "run()", &target, false)
                .err()
                .expect("mismatched counts must fail");
            assert_eq!(
                err.to_string(),
                format!(
                    "sensitive-cache deployment count ({count}) does not match deployment count (1); \
                     the multi-chain deployment and its sensitive-cache counterpart are out of sync"
                )
            );
        }
    }
}