Merge remote-tracking branch 'origin/main' into feature/account-allocator

Cargo.lock

@@ -4467,6 +4467,8 @@ dependencies = [
 name = "revive-dt-common"
 version = "0.1.0"
 dependencies = [
+ "alloy",
+ "alloy-primitives",
  "anyhow",
  "clap",
  "moka",

crates/common/Cargo.toml

@@ -9,6 +9,8 @@ repository.workspace = true
 rust-version.workspace = true
 
 [dependencies]
+alloy = { workspace = true }
+alloy-primitives = { workspace = true }
 anyhow = { workspace = true }
 clap = { workspace = true }
 moka = { workspace = true, features = ["sync"] }

crates/common/src/types/mod.rs

@@ -1,7 +1,9 @@
 mod identifiers;
 mod mode;
+mod private_key_allocator;
 mod version_or_requirement;
 
 pub use identifiers::*;
 pub use mode::*;
+pub use private_key_allocator::*;
 pub use version_or_requirement::*;

crates/common/src/types/private_key_allocator.rs

@@ -0,0 +1,35 @@
+use alloy::signers::local::PrivateKeySigner;
+use alloy_primitives::U256;
+use anyhow::{Result, bail};
+
+/// This is a sequential private key allocator. When instantiated, it allocated private keys in
+/// sequentially and in order until the maximum private key specified is reached.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct PrivateKeyAllocator {
+    /// The next private key to be returned by the allocator when requested.
+    next_private_key: U256,
+
+    /// The highest private key (exclusive) that can be returned by this allocator.
+    highest_private_key_exclusive: U256,
+}
+
+impl PrivateKeyAllocator {
+    /// Creates a new instance of the private key allocator.
+    pub fn new(highest_private_key_exclusive: U256) -> Self {
+        Self {
+            next_private_key: U256::ZERO,
+            highest_private_key_exclusive,
+        }
+    }
+
+    /// Allocates a new private key and errors out if the maximum private key has been reached.
+    pub fn allocate(&mut self) -> Result<PrivateKeySigner> {
+        if self.next_private_key >= self.highest_private_key_exclusive {
+            bail!("Attempted to allocate a private key but failed since all have been allocated");
+        };
+        let private_key =
+            PrivateKeySigner::from_slice(self.next_private_key.to_be_bytes::<32>().as_slice())?;
+        self.next_private_key += U256::ONE;
+        Ok(private_key)
+    }
+}

crates/config/src/lib.rs

@@ -490,6 +490,10 @@ impl WalletConfiguration {
             })
             .clone()
     }
+
+    pub fn highest_private_key_exclusive(&self) -> U256 {
+        U256::try_from(self.additional_keys).unwrap()
+    }
 }
 
 fn serialize_private_key<S>(value: &PrivateKeySigner, serializer: S) -> Result<S::Ok, S::Error>

crates/core/src/driver/mod.rs

@@ -2,6 +2,7 @@
 
 use std::collections::HashMap;
 use std::path::PathBuf;
+use std::sync::Arc;
 
 use alloy::consensus::EMPTY_ROOT_HASH;
 use alloy::hex;
@@ -17,22 +18,23 @@ use alloy::{
     primitives::Address,
     rpc::types::{TransactionRequest, trace::geth::DiffMode},
 };
-use anyhow::Context as _;
+use anyhow::{Context as _, bail};
 use futures::{TryStreamExt, future::try_join_all};
 use indexmap::IndexMap;
-use revive_dt_common::types::PlatformIdentifier;
+use revive_dt_common::types::{PlatformIdentifier, PrivateKeyAllocator};
 use revive_dt_format::traits::{ResolutionContext, ResolverApi};
 use revive_dt_report::ExecutionSpecificReporter;
 use semver::Version;
 
 use revive_dt_format::case::Case;
-use revive_dt_format::input::{
+use revive_dt_format::metadata::{ContractIdent, ContractInstance, ContractPathAndIdent};
+use revive_dt_format::steps::{
     BalanceAssertionStep, Calldata, EtherValue, Expected, ExpectedOutput, FunctionCallStep, Method,
     StepIdx, StorageEmptyAssertionStep,
 };
-use revive_dt_format::metadata::{ContractIdent, ContractInstance, ContractPathAndIdent};
-use revive_dt_format::{input::Step, metadata::Metadata};
+use revive_dt_format::{metadata::Metadata, steps::Step};
 use revive_dt_node_interaction::EthereumNode;
+use tokio::sync::Mutex;
 use tokio::try_join;
 use tracing::{Instrument, info, info_span, instrument};
 
@@ -53,6 +55,10 @@ pub struct CaseState {
 
     /// The execution reporter.
     execution_reporter: ExecutionSpecificReporter,
+
+    /// The private key allocator used for this case state. This is an Arc Mutex to allow for the
+    /// state to be cloned and for all of the clones to refer to the same allocator.
+    private_key_allocator: Arc<Mutex<PrivateKeyAllocator>>,
 }
 
 impl CaseState {
@@ -61,6 +67,7 @@ impl CaseState {
         compiled_contracts: HashMap<PathBuf, HashMap<String, (String, JsonAbi)>>,
         deployed_contracts: HashMap<ContractInstance, (ContractIdent, Address, JsonAbi)>,
         execution_reporter: ExecutionSpecificReporter,
+        private_key_allocator: Arc<Mutex<PrivateKeyAllocator>>,
     ) -> Self {
         Self {
             compiled_contracts,
@@ -68,6 +75,7 @@ impl CaseState {
             variables: Default::default(),
             compiler_version,
             execution_reporter,
+            private_key_allocator,
         }
     }
 
@@ -108,6 +116,12 @@ impl CaseState {
                 .context("Failed to handle the repetition step")?;
                 Ok(StepOutput::Repetition)
             }
+            Step::AllocateAccount(account_allocation) => {
+                self.handle_account_allocation(account_allocation.variable_name.as_str())
+                    .await
+                    .context("Failed to allocate account")?;
+                Ok(StepOutput::AccountAllocation)
+            }
         }
         .inspect(|_| info!("Step Succeeded"))
     }
@@ -201,6 +215,21 @@ impl CaseState {
         Ok(())
     }
 
+    #[instrument(level = "info", name = "Handling Account Allocation", skip_all)]
+    pub async fn handle_account_allocation(&mut self, variable_name: &str) -> anyhow::Result<()> {
+        let Some(variable_name) = variable_name.strip_prefix("$VARIABLE:") else {
+            bail!("Account allocation must start with $VARIABLE:");
+        };
+
+        let private_key = self.private_key_allocator.lock().await.allocate()?;
+        let account = private_key.address();
+        let variable = U256::from_be_slice(account.0.as_slice());
+
+        self.variables.insert(variable_name.to_string(), variable);
+
+        Ok(())
+    }
+
     /// Handles the contract deployment for a given input performing it if it needs to be performed.
     #[instrument(level = "info", skip_all)]
     async fn handle_input_contract_deployment(
@@ -227,15 +256,16 @@ impl CaseState {
                 .then_some(input.value)
                 .flatten();
 
+            let caller = {
+                let context = self.default_resolution_context();
+                let resolver = node.resolver().await?;
+                input
+                    .caller
+                    .resolve_address(resolver.as_ref(), context)
+                    .await?
+            };
             if let (_, _, Some(receipt)) = self
-                .get_or_deploy_contract_instance(
-                    &instance,
-                    metadata,
-                    input.caller,
-                    calldata,
-                    value,
-                    node,
-                )
+                .get_or_deploy_contract_instance(&instance, metadata, caller, calldata, value, node)
                 .await
                 .context("Failed to get or deploy contract instance during input execution")?
             {
@@ -465,13 +495,9 @@ impl CaseState {
             {
                 // Handling the emitter assertion.
                 if let Some(ref expected_address) = expected_event.address {
-                    let expected = Address::from_slice(
-                        Calldata::new_compound([expected_address])
-                            .calldata(resolver, resolution_context)
-                            .await?
-                            .get(12..32)
-                            .expect("Can't fail"),
-                    );
+                    let expected = expected_address
+                        .resolve_address(resolver, resolution_context)
+                        .await?;
                     let actual = actual_event.address();
                     if actual != expected {
                         tracing::error!(
@@ -568,17 +594,17 @@ impl CaseState {
         balance_assertion: &BalanceAssertionStep,
         node: &dyn EthereumNode,
     ) -> anyhow::Result<()> {
-        let Some(instance) = balance_assertion
-            .address
-            .strip_suffix(".address")
-            .map(ContractInstance::new)
-        else {
+        let Some(address) = balance_assertion.address.as_resolvable_address() else {
             return Ok(());
         };
+        let Some(instance) = address.strip_suffix(".address").map(ContractInstance::new) else {
+            return Ok(());
+        };
+
         self.get_or_deploy_contract_instance(
             &instance,
             metadata,
-            FunctionCallStep::default_caller(),
+            FunctionCallStep::default_caller_address(),
             None,
             None,
             node,
@@ -591,20 +617,16 @@ impl CaseState {
     pub async fn handle_balance_assertion_execution(
         &mut self,
         BalanceAssertionStep {
-            address: address_string,
+            address,
             expected_balance: amount,
             ..
         }: &BalanceAssertionStep,
         node: &dyn EthereumNode,
     ) -> anyhow::Result<()> {
         let resolver = node.resolver().await?;
-        let address = Address::from_slice(
-            Calldata::new_compound([address_string])
-                .calldata(resolver.as_ref(), self.default_resolution_context())
-                .await?
-                .get(12..32)
-                .expect("Can't fail"),
-        );
+        let address = address
+            .resolve_address(resolver.as_ref(), self.default_resolution_context())
+            .await?;
 
         let balance = node.balance_of(address).await?;
 
@@ -616,7 +638,7 @@ impl CaseState {
                 "Balance assertion failed - Expected {} but got {} for {} resolved to {}",
                 expected,
                 actual,
-                address_string,
+                address,
                 address,
             )
         }
@@ -631,17 +653,17 @@ impl CaseState {
         storage_empty_assertion: &StorageEmptyAssertionStep,
         node: &dyn EthereumNode,
     ) -> anyhow::Result<()> {
-        let Some(instance) = storage_empty_assertion
-            .address
-            .strip_suffix(".address")
-            .map(ContractInstance::new)
-        else {
+        let Some(address) = storage_empty_assertion.address.as_resolvable_address() else {
             return Ok(());
         };
+        let Some(instance) = address.strip_suffix(".address").map(ContractInstance::new) else {
+            return Ok(());
+        };
+
         self.get_or_deploy_contract_instance(
             &instance,
             metadata,
-            FunctionCallStep::default_caller(),
+            FunctionCallStep::default_caller_address(),
             None,
             None,
             node,
@@ -654,20 +676,16 @@ impl CaseState {
     pub async fn handle_storage_empty_assertion_execution(
         &mut self,
         StorageEmptyAssertionStep {
-            address: address_string,
+            address,
             is_storage_empty,
             ..
         }: &StorageEmptyAssertionStep,
         node: &dyn EthereumNode,
     ) -> anyhow::Result<()> {
         let resolver = node.resolver().await?;
-        let address = Address::from_slice(
-            Calldata::new_compound([address_string])
-                .calldata(resolver.as_ref(), self.default_resolution_context())
-                .await?
-                .get(12..32)
-                .expect("Can't fail"),
-        );
+        let address = address
+            .resolve_address(resolver.as_ref(), self.default_resolution_context())
+            .await?;
 
         let storage = node.latest_state_proof(address, Default::default()).await?;
         let is_empty = storage.storage_hash == EMPTY_ROOT_HASH;
@@ -681,7 +699,7 @@ impl CaseState {
                 "Storage Empty Assertion failed - Expected {} but got {} for {} resolved to {}",
                 expected,
                 actual,
-                address_string,
+                address,
                 address,
             )
         };
@@ -875,4 +893,5 @@ pub enum StepOutput {
     BalanceAssertion,
     StorageEmptyAssertion,
     Repetition,
+    AccountAllocation,
 }

crates/core/src/main.rs

@@ -26,10 +26,14 @@ use revive_dt_report::{
 };
 use schemars::schema_for;
 use serde_json::{Value, json};
+use tokio::sync::Mutex;
 use tracing::{debug, error, info, info_span, instrument};
 use tracing_subscriber::{EnvFilter, FmtSubscriber};
 
-use revive_dt_common::{iterators::EitherIter, types::Mode};
+use revive_dt_common::{
+    iterators::EitherIter,
+    types::{Mode, PrivateKeyAllocator},
+};
 use revive_dt_compiler::SolidityCompiler;
 use revive_dt_config::{Context, *};
 use revive_dt_core::{
@@ -39,9 +43,9 @@ use revive_dt_core::{
 use revive_dt_format::{
     case::{Case, CaseIdx},
     corpus::Corpus,
-    input::{FunctionCallStep, Step},
     metadata::{ContractPathAndIdent, Metadata, MetadataFile},
     mode::ParsedMode,
+    steps::{FunctionCallStep, Step},
 };
 
 use crate::cached_compiler::CachedCompiler;
@@ -326,8 +330,13 @@ async fn start_driver_task<'a>(
                         .expect("Can't fail");
                 }
 
+                let private_key_allocator = Arc::new(Mutex::new(PrivateKeyAllocator::new(
+                    context.wallet_configuration.highest_private_key_exclusive(),
+                )));
+
                 let reporter = test.reporter.clone();
-                let result = handle_case_driver(&test, cached_compiler).await;
+                let result =
+                    handle_case_driver(&test, cached_compiler, private_key_allocator).await;
 
                 match result {
                     Ok(steps_executed) => reporter
@@ -438,6 +447,7 @@ async fn start_cli_reporting_task(reporter: Reporter) {
 async fn handle_case_driver<'a>(
     test: &Test<'a>,
     cached_compiler: Arc<CachedCompiler<'a>>,
+    private_key_allocator: Arc<Mutex<PrivateKeyAllocator>>,
 ) -> anyhow::Result<usize> {
     let platform_state = stream::iter(test.platforms.iter())
         // Compiling the pre-link contracts.
@@ -511,13 +521,14 @@ async fn handle_case_driver<'a>(
                         .steps
                         .iter()
                         .filter_map(|step| match step {
-                            Step::FunctionCall(input) => Some(input.caller),
+                            Step::FunctionCall(input) => input.caller.as_address().copied(),
                             Step::BalanceAssertion(..) => None,
                             Step::StorageEmptyAssertion(..) => None,
                             Step::Repeat(..) => None,
+                            Step::AllocateAccount(..) => None,
                         })
                         .next()
-                        .unwrap_or(FunctionCallStep::default_caller());
+                        .unwrap_or(FunctionCallStep::default_caller_address());
                     let tx = TransactionBuilder::<Ethereum>::with_deploy_code(
                         TransactionRequest::default().from(deployer_address),
                         code,
@@ -564,6 +575,7 @@ async fn handle_case_driver<'a>(
         .filter_map(
             |(test, platform, node, compiler, reporter, _, deployed_libraries)| {
                 let cached_compiler = cached_compiler.clone();
+                let private_key_allocator = private_key_allocator.clone();
 
                 async move {
                     let compiler_output = cached_compiler
@@ -591,6 +603,7 @@ async fn handle_case_driver<'a>(
                         compiler_output.contracts,
                         deployed_libraries.unwrap_or_default(),
                         reporter.clone(),
+                        private_key_allocator,
                     );
 
                     Some((*node, platform.platform_identifier(), case_state))

crates/format/src/case.rs

@@ -4,8 +4,8 @@ use serde::{Deserialize, Serialize};
 use revive_dt_common::{macros::define_wrapper_type, types::Mode};
 
 use crate::{
-    input::{Expected, Step},
     mode::ParsedMode,
+    steps::{Expected, Step},
 };
 
 #[derive(Debug, Default, Serialize, Deserialize, Clone, Eq, PartialEq, JsonSchema)]

crates/format/src/lib.rs

@@ -2,7 +2,7 @@
 
 pub mod case;
 pub mod corpus;
-pub mod input;
 pub mod metadata;
 pub mod mode;
+pub mod steps;
 pub mod traits;

crates/format/src/steps.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, fmt::Display};
 
 use alloy::{
     eips::BlockNumberOrTag,
@@ -29,12 +29,19 @@ use crate::{metadata::ContractInstance, traits::ResolutionContext};
 pub enum Step {
     /// A function call or an invocation to some function on some smart contract.
     FunctionCall(Box<FunctionCallStep>),
+
     /// A step for performing a balance assertion on some account or contract.
     BalanceAssertion(Box<BalanceAssertionStep>),
+
     /// A step for asserting that the storage of some contract or account is empty.
     StorageEmptyAssertion(Box<StorageEmptyAssertionStep>),
+
     /// A special step for repeating a bunch of steps a certain number of times.
     Repeat(Box<RepeatStep>),
+
+    /// A step type that allows for a new account address to be allocated and to later on be used
+    /// as the caller in another step.
+    AllocateAccount(Box<AllocateAccountStep>),
 }
 
 define_wrapper_type!(
@@ -49,7 +56,7 @@ pub struct FunctionCallStep {
     /// The address of the account performing the call and paying the fees for it.
     #[serde(default = "FunctionCallStep::default_caller")]
     #[schemars(with = "String")]
-    pub caller: Address,
+    pub caller: StepAddress,
 
     /// An optional comment on the step which has no impact on the execution in any way.
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -86,7 +93,7 @@ pub struct FunctionCallStep {
 
 /// This represents a balance assertion step where the framework needs to query the balance of some
 /// account or contract and assert that it's some amount.
-#[derive(Clone, Debug, Default, Serialize, Deserialize, Eq, PartialEq, JsonSchema)]
+#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, JsonSchema)]
 pub struct BalanceAssertionStep {
     /// An optional comment on the balance assertion.
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -98,7 +105,7 @@ pub struct BalanceAssertionStep {
     /// this could be a normal hex address, a variable such as `Test.address`, or perhaps even a
     /// full on variable like `$VARIABLE:Uniswap`. It follows the same resolution rules that are
     /// followed in the calldata.
-    pub address: String,
+    pub address: StepAddress,
 
     /// The amount of balance to assert that the account or contract has. This is a 256 bit string
     /// that's serialized and deserialized into a decimal string.
@@ -108,7 +115,7 @@ pub struct BalanceAssertionStep {
 
 /// This represents an assertion for the storage of some contract or account and whether it's empty
 /// or not.
-#[derive(Clone, Debug, Default, Serialize, Deserialize, Eq, PartialEq, JsonSchema)]
+#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, JsonSchema)]
 pub struct StorageEmptyAssertionStep {
     /// An optional comment on the storage empty assertion.
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -120,7 +127,7 @@ pub struct StorageEmptyAssertionStep {
     /// this could be a normal hex address, a variable such as `Test.address`, or perhaps even a
     /// full on variable like `$VARIABLE:Uniswap`. It follows the same resolution rules that are
     /// followed in the calldata.
-    pub address: String,
+    pub address: StepAddress,
 
     /// A boolean of whether the storage of the address is empty or not.
     pub is_storage_empty: bool,
@@ -130,6 +137,10 @@ pub struct StorageEmptyAssertionStep {
 /// steps to be repeated (on different drivers) a certain number of times.
 #[derive(Clone, Debug, Default, Serialize, Deserialize, Eq, PartialEq, JsonSchema)]
 pub struct RepeatStep {
+    /// An optional comment on the repetition step.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub comment: Option<String>,
+
     /// The number of repetitions that the steps should be repeated for.
     pub repeat: usize,
 
@@ -137,6 +148,19 @@ pub struct RepeatStep {
     pub steps: Vec<Step>,
 }
 
+#[derive(Clone, Debug, Default, Serialize, Deserialize, Eq, PartialEq, JsonSchema)]
+pub struct AllocateAccountStep {
+    /// An optional comment on the account allocation step.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub comment: Option<String>,
+
+    /// An instruction to allocate a new account with the value being the variable name of that
+    /// account. This must start with `$VARIABLE:` and then be followed by the variable name of the
+    /// account.
+    #[serde(rename = "allocate_account")]
+    pub variable_name: String,
+}
+
 /// A set of expectations and assertions to make about the transaction after it ran.
 ///
 /// If this is not specified then the only assertion that will be ran is that the transaction
@@ -177,7 +201,7 @@ pub struct ExpectedOutput {
 pub struct Event {
     /// An optional field of the address of the emitter of the event.
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub address: Option<String>,
+    pub address: Option<StepAddress>,
 
     /// The set of topics to expect the event to have.
     pub topics: Vec<String>,
@@ -310,13 +334,74 @@ pub struct VariableAssignments {
     pub return_data: Vec<String>,
 }
 
+/// An address type that might either be an address literal or a resolvable address.
+#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, JsonSchema)]
+#[schemars(with = "String")]
+#[serde(untagged)]
+pub enum StepAddress {
+    Address(Address),
+    ResolvableAddress(String),
+}
+
+impl Default for StepAddress {
+    fn default() -> Self {
+        Self::Address(Default::default())
+    }
+}
+
+impl Display for StepAddress {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            StepAddress::Address(address) => Display::fmt(address, f),
+            StepAddress::ResolvableAddress(address) => Display::fmt(address, f),
+        }
+    }
+}
+
+impl StepAddress {
+    pub fn as_address(&self) -> Option<&Address> {
+        match self {
+            StepAddress::Address(address) => Some(address),
+            StepAddress::ResolvableAddress(_) => None,
+        }
+    }
+
+    pub fn as_resolvable_address(&self) -> Option<&str> {
+        match self {
+            StepAddress::ResolvableAddress(address) => Some(address),
+            StepAddress::Address(..) => None,
+        }
+    }
+
+    pub async fn resolve_address(
+        &self,
+        resolver: &(impl ResolverApi + ?Sized),
+        context: ResolutionContext<'_>,
+    ) -> anyhow::Result<Address> {
+        match self {
+            StepAddress::Address(address) => Ok(*address),
+            StepAddress::ResolvableAddress(address) => Ok(Address::from_slice(
+                Calldata::new_compound([address])
+                    .calldata(resolver, context)
+                    .await?
+                    .get(12..32)
+                    .expect("Can't fail"),
+            )),
+        }
+    }
+}
+
 impl FunctionCallStep {
-    pub const fn default_caller() -> Address {
+    pub const fn default_caller_address() -> Address {
         Address(FixedBytes(alloy::hex!(
             "0x90F8bf6A479f320ead074411a4B0e7944Ea8c9C1"
         )))
     }
 
+    pub const fn default_caller() -> StepAddress {
+        StepAddress::Address(Self::default_caller_address())
+    }
+
     fn default_instance() -> ContractInstance {
         ContractInstance::new("Test")
     }
@@ -399,7 +484,8 @@ impl FunctionCallStep {
             .encoded_input(resolver, context)
             .await
             .context("Failed to encode input bytes for transaction request")?;
-        let transaction_request = TransactionRequest::default().from(self.caller).value(
+        let caller = self.caller.resolve_address(resolver, context).await?;
+        let transaction_request = TransactionRequest::default().from(caller).value(
             self.value
                 .map(|value| value.into_inner())
                 .unwrap_or_default(),

crates/report/src/common.rs

@@ -4,7 +4,7 @@ use std::{path::PathBuf, sync::Arc};
 
 use revive_dt_common::{define_wrapper_type, types::PlatformIdentifier};
 use revive_dt_compiler::Mode;
-use revive_dt_format::{case::CaseIdx, input::StepIdx};
+use revive_dt_format::{case::CaseIdx, steps::StepIdx};
 use serde::{Deserialize, Serialize};
 
 define_wrapper_type!(

schema.json

@@ -201,6 +201,10 @@
         {
           "description": "A special step for repeating a bunch of steps a certain number of times.",
           "$ref": "#/$defs/RepeatStep"
+        },
+        {
+          "description": "A step type that allows for a new account address to be allocated and to later on be used\nas the caller in another step.",
+          "$ref": "#/$defs/AllocateAccountStep"
         }
       ]
     },
@@ -377,9 +381,13 @@
       "properties": {
         "address": {
           "description": "An optional field of the address of the emitter of the event.",
-          "type": [
-            "string",
-            "null"
+          "anyOf": [
+            {
+              "$ref": "#/$defs/StepAddress"
+            },
+            {
+              "type": "null"
+            }
           ]
         },
         "topics": {
@@ -399,6 +407,10 @@
         "values"
       ]
     },
+    "StepAddress": {
+      "description": "An address type that might either be an address literal or a resolvable address.",
+      "type": "string"
+    },
     "EtherValue": {
       "description": "Defines an Ether value.\n\nThis is an unsigned 256 bit integer that's followed by some denomination which can either be\neth, ether, gwei, or wei.",
       "type": "string"
@@ -431,7 +443,7 @@
         },
         "address": {
           "description": "The address that the balance assertion should be done on.\n\nThis is a string which will be resolved into an address when being processed. Therefore,\nthis could be a normal hex address, a variable such as `Test.address`, or perhaps even a\nfull on variable like `$VARIABLE:Uniswap`. It follows the same resolution rules that are\nfollowed in the calldata.",
-          "type": "string"
+          "$ref": "#/$defs/StepAddress"
         },
         "expected_balance": {
           "description": "The amount of balance to assert that the account or contract has. This is a 256 bit string\nthat's serialized and deserialized into a decimal string.",
@@ -456,7 +468,7 @@
         },
         "address": {
           "description": "The address that the balance assertion should be done on.\n\nThis is a string which will be resolved into an address when being processed. Therefore,\nthis could be a normal hex address, a variable such as `Test.address`, or perhaps even a\nfull on variable like `$VARIABLE:Uniswap`. It follows the same resolution rules that are\nfollowed in the calldata.",
-          "type": "string"
+          "$ref": "#/$defs/StepAddress"
         },
         "is_storage_empty": {
           "description": "A boolean of whether the storage of the address is empty or not.",
@@ -472,6 +484,13 @@
       "description": "This represents a repetition step which is a special step type that allows for a sequence of\nsteps to be repeated (on different drivers) a certain number of times.",
       "type": "object",
       "properties": {
+        "comment": {
+          "description": "An optional comment on the repetition step.",
+          "type": [
+            "string",
+            "null"
+          ]
+        },
         "repeat": {
           "description": "The number of repetitions that the steps should be repeated for.",
           "type": "integer",
@@ -491,6 +510,25 @@
         "steps"
       ]
     },
+    "AllocateAccountStep": {
+      "type": "object",
+      "properties": {
+        "comment": {
+          "description": "An optional comment on the account allocation step.",
+          "type": [
+            "string",
+            "null"
+          ]
+        },
+        "allocate_account": {
+          "description": "An instruction to allocate a new account with the value being the variable name of that\naccount. This must start with `$VARIABLE:` and then be followed by the variable name of the\naccount.",
+          "type": "string"
+        }
+      },
+      "required": [
+        "allocate_account"
+      ]
+    },
     "ContractPathAndIdent": {
       "description": "Represents an identifier used for contracts.\n\nThe type supports serialization from and into the following string format:\n\n```text\n${path}:${contract_ident}\n```",
       "type": "string"
@@ -542,4 +580,4 @@
       ]
     }
   }
-}
+}