Add support overriding block id/tag/hash with a named fork binding

crates/cheatnet/src/runtime_extensions/forge_config_extension/config.rs

@@ -25,10 +25,17 @@ pub struct InlineForkConfig {
     pub block: BlockId,
 }
 
+#[derive(Debug, Clone, CairoDeserialize, PartialEq)]
+pub struct OverriddenForkConfig {
+    pub name: ByteArray,
+    pub block: BlockId,
+}
+
 #[derive(Debug, Clone, CairoDeserialize, PartialEq)]
 pub enum RawForkConfig {
     Inline(InlineForkConfig),
     Named(ByteArray),
+    Overridden(OverriddenForkConfig),
 }
 
 // fuzzer

crates/forge/src/run_tests/resolve_config.rs

@@ -3,6 +3,7 @@ use anyhow::{anyhow, Result};
 use cheatnet::runtime_extensions::forge_config_extension::config::{
     BlockId, InlineForkConfig, RawForkConfig,
 };
+use conversions::byte_array::ByteArray;
 use forge_runner::package_tests::{
     with_config::TestTargetWithConfig,
     with_config_resolved::{
@@ -98,21 +99,28 @@ fn parse_block_id(fork_target: &ForkTarget) -> Result<BlockId> {
     Ok(block_id)
 }
 
+fn get_fork_target_from_runner_config<'a>(
+    fork_targets: &'a [ForkTarget],
+    name: &ByteArray,
+) -> Result<&'a ForkTarget> {
+    fork_targets
+        .iter()
+        .find(|fork| fork.name == String::from(name.clone()))
+        .ok_or_else(|| {
+            let name = String::from(name.clone());
+            anyhow!("Fork configuration named = {name} not found in the Scarb.toml")
+        })
+}
+
 fn replace_id_with_params(
     raw_fork_config: RawForkConfig,
     fork_targets: &[ForkTarget],
 ) -> Result<InlineForkConfig> {
     match raw_fork_config {
         RawForkConfig::Inline(raw_fork_params) => Ok(raw_fork_params),
         RawForkConfig::Named(name) => {
-            let fork_target_from_runner_config = fork_targets
-                .iter()
-                .find(|fork| fork.name == String::from(name.clone()))
-                .ok_or_else(|| {
-                    let name = String::from(name);
-
-                    anyhow!("Fork configuration named = {name} not found in the Scarb.toml")
-                })?;
+            let fork_target_from_runner_config =
+                get_fork_target_from_runner_config(fork_targets, &name)?;
 
             let block_id = parse_block_id(fork_target_from_runner_config)?;
 
@@ -121,6 +129,15 @@ fn replace_id_with_params(
                 block: block_id,
             })
         }
+        RawForkConfig::Overridden(overridden) => {
+            let fork_target_from_runner_config =
+                get_fork_target_from_runner_config(fork_targets, &overridden.name)?;
+
+            Ok(InlineForkConfig {
+                url: fork_target_from_runner_config.url.parse()?,
+                block: overridden.block,
+            })
+        }
     }
 }
 

crates/snforge-scarb-plugin/src/args.rs

@@ -84,6 +84,11 @@ impl Arguments {
         }
     }
 
+    #[inline]
+    pub fn unnamed(&self) -> UnnamedArgs {
+        UnnamedArgs::new(&self.unnamed)
+    }
+
     #[inline]
     pub fn assert_is_empty<T: AttributeInfo>(&self) -> Result<(), Diagnostic> {
         if self.is_empty() {

crates/snforge-scarb-plugin/src/args/unnamed.rs

@@ -3,10 +3,10 @@ use cairo_lang_macro::Diagnostic;
 use cairo_lang_syntax::node::ast::Expr;
 use std::{collections::HashMap, ops::Deref};
 
-pub struct UnnamedArgs<'a>(Vec<&'a Expr>);
+pub struct UnnamedArgs<'a>(Vec<(usize, &'a Expr)>);
 
 impl<'a> Deref for UnnamedArgs<'a> {
-    type Target = Vec<&'a Expr>;
+    type Target = Vec<(usize, &'a Expr)>;
 
     fn deref(&self) -> &Self::Target {
         &self.0
@@ -19,7 +19,7 @@ impl UnnamedArgs<'_> {
 
         args.sort_by(|(a, _), (b, _)| a.cmp(b));
 
-        let args = args.into_iter().map(|(_, expr)| expr).collect();
+        let args = args.into_iter().map(|(&pos, expr)| (pos, expr)).collect();
 
         UnnamedArgs(args)
     }
@@ -28,7 +28,7 @@ impl UnnamedArgs<'_> {
 impl<'a> UnnamedArgs<'a> {
     pub fn of_length<const N: usize, T: AttributeInfo>(
         &self,
-    ) -> Result<&[&'a Expr; N], Diagnostic> {
+    ) -> Result<&[(usize, &'a Expr); N], Diagnostic> {
         self.as_slice()
             .try_into()
             .map_err(|_| T::error(format!("expected {} arguments, got: {}", N, self.len())))

crates/snforge-scarb-plugin/src/attributes/available_gas.rs

@@ -26,7 +26,7 @@ impl AttributeCollector for AvailableGasCollector {
     ) -> Result<String, Diagnostics> {
         let &[arg] = args.unnamed_only::<Self>()?.of_length::<1, Self>()?;
 
-        let gas = Number::parse_from_expr::<Self>(db, arg, "0")?;
+        let gas = Number::parse_from_expr::<Self>(db, arg.1, arg.0.to_string().as_str())?;
 
         let gas = gas.as_cairo_expression();
 

crates/snforge-scarb-plugin/src/attributes/fork.rs

@@ -2,12 +2,12 @@ use self::block_id::{BlockId, BlockIdVariants};
 use crate::{
     args::Arguments,
     attributes::{AttributeCollector, AttributeInfo, AttributeTypeData},
+    branch,
     cairo_expression::CairoExpression,
     config_statement::extend_with_config_cheatcodes,
     types::ParseFromExpr,
-    utils::branch,
 };
-use cairo_lang_macro::{Diagnostic, Diagnostics, ProcMacroResult, TokenStream};
+use cairo_lang_macro::{Diagnostic, Diagnostics, ProcMacroResult, Severity, TokenStream};
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use indoc::formatdoc;
 use url::Url;
@@ -30,7 +30,11 @@ impl AttributeCollector for ForkCollector {
         args: Arguments,
         _warns: &mut Vec<Diagnostic>,
     ) -> Result<String, Diagnostics> {
-        let expr = branch(inline_args(db, &args), || from_file_args(db, &args))?;
+        let expr = branch!(
+            inline_args(db, &args),
+            overridden_args(db, &args),
+            from_file_args(db, &args)
+        )?;
 
         Ok(expr)
     }
@@ -69,7 +73,7 @@ fn from_file_args(db: &dyn SyntaxGroup, args: &Arguments) -> Result<String, Diag
         .unnamed_only::<ForkCollector>()?
         .of_length::<1, ForkCollector>()?;
 
-    let name = String::parse_from_expr::<ForkCollector>(db, arg, "0")?;
+    let name = String::parse_from_expr::<ForkCollector>(db, arg.1, arg.0.to_string().as_str())?;
 
     let name = name.as_cairo_expression();
 
@@ -78,6 +82,31 @@ fn from_file_args(db: &dyn SyntaxGroup, args: &Arguments) -> Result<String, Diag
     ))
 }
 
+fn overridden_args(db: &dyn SyntaxGroup, args: &Arguments) -> Result<String, Diagnostic> {
+    let &[arg] = args.unnamed().of_length::<1, ForkCollector>()?;
+
+    let block_id = args.named.one_of_once(&[
+        BlockIdVariants::Hash,
+        BlockIdVariants::Number,
+        BlockIdVariants::Tag,
+    ])?;
+
+    let block_id = BlockId::parse_from_expr::<ForkCollector>(db, &block_id, block_id.0.as_ref())?;
+    let name = String::parse_from_expr::<ForkCollector>(db, arg.1, arg.0.to_string().as_str())?;
+
+    let block_id = block_id.as_cairo_expression();
+    let name = name.as_cairo_expression();
+
+    Ok(formatdoc!(
+        "
+            snforge_std::_config_types::ForkConfig::Overridden(
+                block: {block_id},
+                name: {name}
+            )
+        "
+    ))
+}
+
 #[must_use]
 pub fn fork(args: TokenStream, item: TokenStream) -> ProcMacroResult {
     extend_with_config_cheatcodes::<ForkCollector>(args, item)

crates/snforge-scarb-plugin/src/utils.rs

@@ -1,30 +1,51 @@
 use cairo_lang_macro::{Diagnostic, Severity};
 use indoc::formatdoc;
 
-fn higher_severity(a: Severity, b: Severity) -> Severity {
+pub fn higher_severity(a: Severity, b: Severity) -> Severity {
     match (a, b) {
         (Severity::Warning, Severity::Warning) => Severity::Warning,
         _ => Severity::Error,
     }
 }
+pub fn format_error_message(variants: &[Diagnostic]) -> String {
+    let formatted_variants: Vec<String> = variants
+        .iter()
+        .map(|variant| format!("- variant: {}", variant.message))
+        .collect();
 
-pub fn branch(
-    left: Result<String, Diagnostic>,
-    right: impl Fn() -> Result<String, Diagnostic>,
-) -> Result<String, Diagnostic> {
-    left.or_else(|error| {
-        right().map_err(|next_error| Diagnostic {
-            severity: higher_severity(error.severity, next_error.severity),
-            message: formatdoc!(
-                "
-                    Both options failed
-                    First variant: {}
-                    Second variant: {}
-                    Resolve at least one of them
-                ",
-                error.message,
-                next_error.message
-            ),
-        })
-    })
+    formatdoc! {"
+        All options failed
+        {}
+        Resolve at least one of them
+    ", formatted_variants.join("\n")}
+}
+
+#[macro_export]
+macro_rules! branch {
+    ($($result:expr),+) => {{
+        let mut messages = Vec::new();
+        let mut result = None;
+
+        $(
+            if result.is_none() {
+                match $result {
+                    Ok(val) => {
+                        result = Some(val);
+                    },
+                    Err(err) => {
+                        messages.push(err);
+                    },
+                }
+            }
+        )+
+
+        if let Some(result) = result {
+            Ok(result)
+        } else {
+            Err(Diagnostic {
+                message: $crate::utils::format_error_message(&messages),
+                severity: messages.into_iter().fold(Severity::Warning, |acc, x| $crate::utils::higher_severity(acc, x.severity))
+            })
+        }
+    }};
 }