Skip to main content

huff_neo_test_runner/
spec.rs

1use crate::errors::{DecoratorError, DecoratorErrorKind, RunnerError};
2use crate::outcome::{Expectations, ExpectedOutcome};
3use alloy_primitives::{Address, B256, U256, hex, keccak256};
4use huff_neo_codegen::Codegen;
5use huff_neo_utils::ast::huff::{
6    BuiltinFunctionArg, BuiltinFunctionCall, BuiltinFunctionKind, Decorator, DecoratorArg, DecoratorFlag, DecoratorValue,
7};
8use huff_neo_utils::ast::span::AstSpan;
9use huff_neo_utils::prelude::Contract;
10use std::collections::HashMap;
11use std::ops::RangeInclusive;
12
13/// The setup a test's decorators ask for before the test call
14#[derive(Debug, Clone, Default, PartialEq, Eq)]
15pub struct TestSpec {
16    /// The call value of the test transaction
17    pub value: U256,
18    /// The calldata of the test transaction
19    pub calldata: Vec<u8>,
20    /// Storage slots written before the test call, in source order
21    pub storage: Vec<StorageSlot>,
22    /// Balances set before the test call
23    pub balances: Vec<AccountValue<U256>>,
24    /// Code placed at addresses before the test call
25    pub code: Vec<AccountValue<Vec<u8>>>,
26    /// Nonces set before the test call
27    pub nonces: Vec<AccountValue<u64>>,
28    /// Transaction and block fields of the test call that differ from the runner's
29    pub env: EnvOverrides,
30    /// How the test call must end, and its gas bounds
31    pub expectations: Expectations,
32}
33
34/// Transaction and block fields a test's decorators override; `None` keeps the runner's value
35#[derive(Debug, Clone, Default, PartialEq, Eq)]
36pub struct EnvOverrides {
37    /// Sender of the test transaction, both `msg.sender` and `tx.origin`
38    pub caller: Option<Address>,
39    /// Gas limit of the test transaction
40    pub gas_limit: Option<u64>,
41    /// Gas price of the test transaction
42    pub gas_price: Option<u128>,
43    /// Block number
44    pub block_number: Option<U256>,
45    /// Block timestamp
46    pub timestamp: Option<U256>,
47    /// Block base fee
48    pub basefee: Option<u64>,
49    /// Blob base fee
50    pub blob_basefee: Option<u128>,
51    /// Block beneficiary
52    pub coinbase: Option<Address>,
53    /// Chain id
54    pub chain_id: Option<u64>,
55    /// Block `prevrandao`
56    pub prevrandao: Option<B256>,
57    /// Beacon slot number read by `slotnum`
58    pub slot_number: Option<u64>,
59}
60
61/// A storage slot of an account with its value, set before the test or expected after it
62#[derive(Debug, Clone, PartialEq, Eq)]
63pub struct StorageSlot {
64    /// The account that owns the slot
65    pub address: Address,
66    /// The storage slot
67    pub slot: U256,
68    /// The value of the slot
69    pub value: U256,
70}
71
72/// A property of an account, such as its balance, set before the test or expected after it
73#[derive(Debug, Clone, PartialEq, Eq)]
74pub struct AccountValue<T> {
75    /// The account
76    pub address: Address,
77    /// The value of the property
78    pub value: T,
79}
80
81/// Addresses that decorator arguments can refer to by name
82#[derive(Debug, Clone, Copy)]
83pub struct SpecContext {
84    /// The address of the deployed test, referred to as `this`
85    pub this: Address,
86    /// The runner's sender of the test transaction; a `caller` flag replaces it as `caller`
87    pub caller: Address,
88}
89
90/// The value an `expect_balance` or `expect_storage` flag expects after the test
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub enum ExpectedValue {
93    /// Exactly this value
94    Exact(U256),
95    /// A value from `min` to `max`, both inclusive and each optional
96    Range { min: Option<U256>, max: Option<U256> },
97}
98
99impl ExpectedValue {
100    /// Whether the value meets the expectation
101    pub fn matches(&self, value: U256) -> bool {
102        match *self {
103            Self::Exact(expected) => value == expected,
104            Self::Range { min, max } => min.is_none_or(|min| value >= min) && max.is_none_or(|max| value <= max),
105        }
106    }
107
108    /// The expectation in words, with each value written by `show`, e.g. "at least 5 and at most 9"
109    pub fn describe(&self, show: impl Fn(U256) -> String) -> String {
110        match *self {
111            Self::Exact(expected) => show(expected),
112            Self::Range { min: Some(min), max: Some(max) } => format!("at least {} and at most {}", show(min), show(max)),
113            Self::Range { min: Some(min), max: None } => format!("at least {}", show(min)),
114            Self::Range { min: None, max: Some(max) } => format!("at most {}", show(max)),
115            Self::Range { min: None, max: None } => String::from("any value"),
116        }
117    }
118}
119
120/// A storage slot of an account and the value it must hold after the test
121#[derive(Debug, Clone, PartialEq, Eq)]
122pub struct ExpectedSlot {
123    /// The account that owns the slot
124    pub address: Address,
125    /// The storage slot
126    pub slot: U256,
127    /// The value the slot must hold
128    pub expected: ExpectedValue,
129}
130
131/// The decorator flags `hnc test` understands
132#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
133enum DecoratorKind {
134    /// `calldata(bytes...)`: the calldata of the test transaction
135    Calldata,
136    /// `value(word)`: the call value of the test transaction
137    Value,
138    /// `store(slot, value, at = address)`: a storage slot set before the test runs
139    Store,
140    /// `deal(address, wei)`: a balance set before the test runs
141    Deal,
142    /// `etch(address, bytes...)`: code placed at an address before the test runs
143    Etch,
144    /// `nonce(address, n)`: a nonce set before the test runs
145    Nonce,
146    /// `caller(address)`: the sender of the test, both `caller` and `origin`
147    Caller,
148    /// `coinbase(address)`: the block beneficiary
149    Coinbase,
150    /// `gas_limit(n)`: the gas limit of the test transaction
151    GasLimit,
152    /// `gas_price(n)`: the gas price of the test transaction
153    GasPrice,
154    /// `block_number(n)`: the block number
155    BlockNumber,
156    /// `timestamp(n)`: the block timestamp
157    Timestamp,
158    /// `basefee(n)`: the block base fee
159    Basefee,
160    /// `blob_basefee(n)`: the blob base fee
161    BlobBasefee,
162    /// `chain_id(n)`: the chain id
163    ChainId,
164    /// `prevrandao(word)`: the block `prevrandao`
165    Prevrandao,
166    /// `slot_number(n)`: the beacon slot number read by `slotnum`
167    SlotNumber,
168    /// A flag that says how the test must end
169    Outcome(OutcomeFlag),
170    /// `expect_gas(min = n, max = m)`: bounds of the gas the test uses
171    ExpectGas,
172    /// `expect_storage(slot, value, at = address)` or `expect_storage(slot, min = n, max = m)`: a storage slot after the test
173    ExpectStorage,
174    /// `expect_balance(address, wei)` or `expect_balance(address, min = n, max = m)`: a balance after the test
175    ExpectBalance,
176    /// `skip` or `skip("reason")`: the test is not run
177    Skip,
178}
179
180/// The flags that say how the test must end; a test can have only one of them
181#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
182enum OutcomeFlag {
183    /// `expect_revert` or `expect_revert(bytes...)`: the test reverts, with exactly this data when given
184    Revert,
185    /// `expect_partial_revert(bytes...)`: the test reverts with data that starts with this
186    PartialRevert,
187    /// `expect_halt` or `expect_halt("Reason")`: the test halts exceptionally, for this reason when given
188    Halt,
189    /// `expect_return(bytes...)`: the test returns exactly this data
190    Return,
191}
192
193impl DecoratorKind {
194    /// Every flag with its name in the source
195    const NAMES: [(Self, &'static str); 25] = [
196        (Self::Calldata, "calldata"),
197        (Self::Value, "value"),
198        (Self::Store, "store"),
199        (Self::Deal, "deal"),
200        (Self::Etch, "etch"),
201        (Self::Nonce, "nonce"),
202        (Self::Caller, "caller"),
203        (Self::Coinbase, "coinbase"),
204        (Self::GasLimit, "gas_limit"),
205        (Self::GasPrice, "gas_price"),
206        (Self::BlockNumber, "block_number"),
207        (Self::Timestamp, "timestamp"),
208        (Self::Basefee, "basefee"),
209        (Self::BlobBasefee, "blob_basefee"),
210        (Self::ChainId, "chain_id"),
211        (Self::Prevrandao, "prevrandao"),
212        (Self::SlotNumber, "slot_number"),
213        (Self::Outcome(OutcomeFlag::Revert), "expect_revert"),
214        (Self::Outcome(OutcomeFlag::PartialRevert), "expect_partial_revert"),
215        (Self::Outcome(OutcomeFlag::Halt), "expect_halt"),
216        (Self::Outcome(OutcomeFlag::Return), "expect_return"),
217        (Self::ExpectGas, "expect_gas"),
218        (Self::ExpectStorage, "expect_storage"),
219        (Self::ExpectBalance, "expect_balance"),
220        (Self::Skip, "skip"),
221    ];
222
223    /// The flag with this name, or `None` for a name `hnc test` does not know
224    fn from_name(name: &str) -> Option<Self> {
225        Self::NAMES.iter().find(|(_, n)| *n == name).map(|(kind, _)| *kind)
226    }
227
228    /// Whether the flag may be given more than once in a test
229    fn is_repeatable(self) -> bool {
230        matches!(self, Self::Store | Self::Deal | Self::Etch | Self::Nonce | Self::ExpectStorage | Self::ExpectBalance)
231    }
232
233    /// Whether the flag checks state after the test, which a revert or halt rolls back
234    fn checks_post_state(self) -> bool {
235        matches!(self, Self::ExpectStorage | Self::ExpectBalance)
236    }
237}
238
239impl TestSpec {
240    /// Checks a test's decorator flags and evaluates their arguments into the setup the runner applies
241    pub fn from_decorator(decorator: Option<&Decorator>, contract: &Contract, ctx: SpecContext) -> Result<Self, RunnerError> {
242        let Some(decorator) = decorator else {
243            return Ok(Self::default());
244        };
245
246        let mut spec = Self::default();
247        let mut eval = Evaluator { contract, ctx };
248        // `caller` in other values means the test's sender, wherever the `caller` flag appears
249        if let Some(flag) = decorator.flags.iter().find(|f| DecoratorKind::from_name(&f.name) == Some(DecoratorKind::Caller))
250            && let [arg] = flag.args.as_slice()
251            && arg.key.is_none()
252        {
253            eval.ctx.caller = eval.address(flag, arg)?;
254        }
255        let mut single_use: HashMap<DecoratorKind, &DecoratorFlag> = HashMap::new();
256        let mut stored_slots: HashMap<(Address, U256), &DecoratorFlag> = HashMap::new();
257        let mut expected_slots: HashMap<(Address, U256), &DecoratorFlag> = HashMap::new();
258        // Accounts set or checked per flag, so each flag may name an address only once
259        let mut accounts: HashMap<(DecoratorKind, Address), &DecoratorFlag> = HashMap::new();
260        let mut outcome_flag: Option<(&DecoratorFlag, OutcomeFlag)> = None;
261        let mut post_state_flag: Option<&DecoratorFlag> = None;
262
263        for flag in &decorator.flags {
264            let Some(kind) = DecoratorKind::from_name(&flag.name) else {
265                return Err(decorator_error(
266                    flag,
267                    DecoratorErrorKind::UnknownFlag,
268                    format!("Unknown decorator flag: {}", flag.name),
269                    flag.span.clone(),
270                ));
271            };
272            if !kind.is_repeatable() {
273                check_single_use(&mut single_use, kind, flag)?;
274            }
275            if kind.checks_post_state() {
276                post_state_flag.get_or_insert(flag);
277            }
278
279            match kind {
280                DecoratorKind::Calldata => {
281                    let args = FlagArgs::new(
282                        flag,
283                        1..=usize::MAX,
284                        &[],
285                        "`calldata` takes the calldata bytes, e.g. calldata(__FUNC_SIG(transfer), 0x01)",
286                    )?;
287                    spec.calldata = eval.bytes(flag, &args.positional)?;
288                }
289                DecoratorKind::Value => {
290                    let arg = single_arg(flag, "a single call value, e.g. value(0x01)")?;
291                    spec.value = eval.word(flag, arg)?;
292                }
293                DecoratorKind::Store => {
294                    let args = FlagArgs::new(
295                        flag,
296                        2..=2,
297                        &["at"],
298                        "`store` takes a slot, a value and an optional `at = address`, e.g. store([SLOT], 0x01, at = 0x1000)",
299                    )?;
300                    let slot = eval.word(flag, args.positional[0])?;
301                    let value = eval.word(flag, args.positional[1])?;
302                    let address = eval.at_address(flag, &args)?;
303
304                    if let Some(first) = stored_slots.insert((address, slot), flag) {
305                        return Err(decorator_error(
306                            flag,
307                            DecoratorErrorKind::Duplicate,
308                            format!("Storage slot {slot:#x} at {address} is already set by an earlier `store`"),
309                            joined_span(first, flag),
310                        ));
311                    }
312                    spec.storage.push(StorageSlot { address, slot, value });
313                }
314                DecoratorKind::Deal => {
315                    let usage = "`deal` takes an address and a balance in wei, e.g. deal(caller, 1000)";
316                    let args = FlagArgs::new(flag, 2..=2, &[], usage)?;
317                    let address = eval.address(flag, args.positional[0])?;
318                    check_account_once(&mut accounts, kind, address, flag)?;
319                    let value = eval.word(flag, args.positional[1])?;
320                    spec.balances.push(AccountValue { address, value });
321                }
322                DecoratorKind::Etch => {
323                    let usage = "`etch` takes an address and the code bytes, e.g. etch(0x1000, \"0x5f5ff3\")";
324                    let args = FlagArgs::new(flag, 2..=usize::MAX, &[], usage)?;
325                    let address = eval.address(flag, args.positional[0])?;
326                    let replaced = if address == eval.ctx.this {
327                        Some("the test itself")
328                    } else if address == eval.ctx.caller {
329                        Some("the sender of the test, which may not have code")
330                    } else {
331                        None
332                    };
333                    if let Some(replaced) = replaced {
334                        return Err(decorator_error(
335                            flag,
336                            DecoratorErrorKind::InvalidValue,
337                            format!("`etch` cannot replace the code of {replaced}"),
338                            args.positional[0].span.clone(),
339                        ));
340                    }
341                    check_account_once(&mut accounts, kind, address, flag)?;
342                    let code = eval.bytes(flag, &args.positional[1..])?;
343                    spec.code.push(AccountValue { address, value: code });
344                }
345                DecoratorKind::Nonce => {
346                    let usage = "`nonce` takes an address and a nonce, e.g. nonce(caller, 5)";
347                    let args = FlagArgs::new(flag, 2..=2, &[], usage)?;
348                    let address = eval.address(flag, args.positional[0])?;
349                    check_account_once(&mut accounts, kind, address, flag)?;
350                    let value = eval.uint(flag, args.positional[1])?;
351                    spec.nonces.push(AccountValue { address, value });
352                }
353                DecoratorKind::ExpectStorage => {
354                    let usage = "`expect_storage` takes a slot, then a value or bounds `min =` and/or `max =`, and an optional \
355                                 `at = address`, e.g. expect_storage([SLOT], 0x01) or expect_storage([SLOT], max = 10, at = 0x1000)";
356                    let args = FlagArgs::new(flag, 1..=2, &["at", "min", "max"], usage)?;
357                    let slot = eval.word(flag, args.positional[0])?;
358                    let expected = eval.expected_value(flag, &args, usage)?;
359                    let address = eval.at_address(flag, &args)?;
360
361                    if let Some(first) = expected_slots.insert((address, slot), flag) {
362                        return Err(decorator_error(
363                            flag,
364                            DecoratorErrorKind::Duplicate,
365                            format!("Storage slot {slot:#x} at {address} is already checked by an earlier `expect_storage`"),
366                            joined_span(first, flag),
367                        ));
368                    }
369                    spec.expectations.storage.push(ExpectedSlot { address, slot, expected });
370                }
371                DecoratorKind::ExpectBalance => {
372                    let usage = "`expect_balance` takes an address, then a balance in wei or bounds `min =` and/or `max =`, \
373                                 e.g. expect_balance(this, 1000) or expect_balance(caller, min = 1, max = 1000)";
374                    let args = FlagArgs::new(flag, 1..=2, &["min", "max"], usage)?;
375                    let address = eval.address(flag, args.positional[0])?;
376                    check_account_once(&mut accounts, kind, address, flag)?;
377                    let expected = eval.expected_value(flag, &args, usage)?;
378                    spec.expectations.balances.push(AccountValue { address, value: expected });
379                }
380                DecoratorKind::Caller => {
381                    let arg = single_arg(flag, "an address, e.g. caller(0xbeef)")?;
382                    spec.env.caller = Some(eval.address(flag, arg)?);
383                }
384                DecoratorKind::Coinbase => {
385                    let arg = single_arg(flag, "an address, e.g. coinbase(0xbeef)")?;
386                    spec.env.coinbase = Some(eval.address(flag, arg)?);
387                }
388                DecoratorKind::GasLimit => {
389                    let arg = single_arg(flag, "a gas amount, e.g. gas_limit(100000)")?;
390                    spec.env.gas_limit = Some(eval.uint(flag, arg)?);
391                }
392                DecoratorKind::GasPrice => {
393                    let arg = single_arg(flag, "a price in wei, e.g. gas_price(10)")?;
394                    spec.env.gas_price = Some(eval.uint(flag, arg)?);
395                }
396                DecoratorKind::BlockNumber => {
397                    let arg = single_arg(flag, "a number, e.g. block_number(100)")?;
398                    spec.env.block_number = Some(eval.word(flag, arg)?);
399                }
400                DecoratorKind::Timestamp => {
401                    let arg = single_arg(flag, "a unix time, e.g. timestamp(1700000000)")?;
402                    spec.env.timestamp = Some(eval.word(flag, arg)?);
403                }
404                DecoratorKind::Basefee => {
405                    let arg = single_arg(flag, "a fee in wei, e.g. basefee(7)")?;
406                    spec.env.basefee = Some(eval.uint(flag, arg)?);
407                }
408                DecoratorKind::BlobBasefee => {
409                    let arg = single_arg(flag, "a fee in wei, e.g. blob_basefee(1)")?;
410                    spec.env.blob_basefee = Some(eval.uint(flag, arg)?);
411                }
412                DecoratorKind::ChainId => {
413                    let arg = single_arg(flag, "a chain id, e.g. chain_id(1)")?;
414                    spec.env.chain_id = Some(eval.uint(flag, arg)?);
415                }
416                DecoratorKind::Prevrandao => {
417                    let arg = single_arg(flag, "a 32-byte value, e.g. prevrandao(0x01)")?;
418                    spec.env.prevrandao = Some(eval.word(flag, arg)?.into());
419                }
420                DecoratorKind::SlotNumber => {
421                    let arg = single_arg(flag, "a slot, e.g. slot_number(100)")?;
422                    spec.env.slot_number = Some(eval.uint(flag, arg)?);
423                }
424                DecoratorKind::Outcome(outcome) => {
425                    if let Some((first, _)) = outcome_flag {
426                        return Err(decorator_error(
427                            flag,
428                            DecoratorErrorKind::Conflict,
429                            format!("`{}` conflicts with `{}`: a test can expect only one way to end", flag.name, first.name),
430                            joined_span(first, flag),
431                        ));
432                    }
433                    outcome_flag = Some((flag, outcome));
434                    spec.expectations.outcome = eval.expected_outcome(flag, outcome)?;
435                    spec.expectations.outcome_flag = Some(flag.name.clone());
436                }
437                // The runner skips the test before it gets here, see `skipped`
438                DecoratorKind::Skip => {}
439                DecoratorKind::ExpectGas => {
440                    let usage = "`expect_gas` takes `min =` and/or `max =`, e.g. expect_gas(max = 30000)";
441                    let args = FlagArgs::new(flag, 0..=0, &["min", "max"], usage)?;
442                    if args.named.is_empty() {
443                        return Err(decorator_error(flag, DecoratorErrorKind::InvalidArguments, usage, flag.span.clone()));
444                    }
445                    let min = args.named.get("min").map(|arg| eval.uint::<u64>(flag, arg)).transpose()?;
446                    let max = args.named.get("max").map(|arg| eval.uint::<u64>(flag, arg)).transpose()?;
447                    if let (Some(min), Some(max)) = (min, max)
448                        && min > max
449                    {
450                        return Err(decorator_error(
451                            flag,
452                            DecoratorErrorKind::InvalidValue,
453                            format!("`min` ({min}) is larger than `max` ({max})"),
454                            flag.span.clone(),
455                        ));
456                    }
457                    spec.expectations.min_gas = min;
458                    spec.expectations.max_gas = max;
459                }
460            }
461        }
462
463        // A revert or halt rolls the state back, so there is nothing left to check
464        if let (Some((outcome, kind)), Some(post_state)) = (outcome_flag, post_state_flag)
465            && kind != OutcomeFlag::Return
466        {
467            return Err(decorator_error(
468                post_state,
469                DecoratorErrorKind::Conflict,
470                format!("`{}` conflicts with `{}`: a reverted or halted test leaves no state to check", post_state.name, outcome.name),
471                joined_span(outcome, post_state),
472            ));
473        }
474
475        Ok(spec)
476    }
477}
478
479/// The arguments of a flag, checked against the positional count and keys it accepts
480struct FlagArgs<'a> {
481    positional: Vec<&'a DecoratorArg>,
482    named: HashMap<&'a str, &'a DecoratorArg>,
483}
484
485impl<'a> FlagArgs<'a> {
486    fn new(flag: &'a DecoratorFlag, positional: RangeInclusive<usize>, keys: &[&str], usage: &str) -> Result<Self, RunnerError> {
487        let invalid = |hint: String| decorator_error(flag, DecoratorErrorKind::InvalidArguments, hint, flag.span.clone());
488
489        let mut args = FlagArgs { positional: Vec::new(), named: HashMap::new() };
490        for arg in &flag.args {
491            match &arg.key {
492                None => args.positional.push(arg),
493                Some(key) if !keys.contains(&key.as_str()) => return Err(invalid(format!("Unknown argument `{key}`. {usage}"))),
494                Some(key) => {
495                    if args.named.insert(key, arg).is_some() {
496                        return Err(invalid(format!("Argument `{key}` is given twice. {usage}")));
497                    }
498                }
499            }
500        }
501
502        if !positional.contains(&args.positional.len()) {
503            return Err(invalid(usage.to_string()));
504        }
505        Ok(args)
506    }
507}
508
509/// Evaluates decorator argument values against the contract and the test's addresses
510struct Evaluator<'a> {
511    contract: &'a Contract,
512    ctx: SpecContext,
513}
514
515impl Evaluator<'_> {
516    /// The 32-byte word an argument stands for, the same value Huff would push for it
517    fn word(&self, flag: &DecoratorFlag, arg: &DecoratorArg) -> Result<U256, RunnerError> {
518        let invalid = |hint: String| decorator_error(flag, DecoratorErrorKind::InvalidValue, hint, arg.span.clone());
519
520        match &arg.value {
521            DecoratorValue::Hex(digits) if digits.len() > 64 => Err(invalid(String::from("A hex literal word is at most 32 bytes"))),
522            DecoratorValue::Hex(digits) => Ok(U256::from_str_radix(digits, 16).expect("the lexer only produces hex digits")),
523            DecoratorValue::Int(int) => Ok(U256::from(*int)),
524            DecoratorValue::Bool(b) => Ok(U256::from(*b as u8)),
525            DecoratorValue::Constant(name) => Codegen::gen_constant_value(self.contract, name, &arg.span)?
526                .map(|v| U256::from_be_bytes(v.value.0))
527                .ok_or_else(|| invalid(format!("Constant [{name}] is __NOOP and has no value"))),
528            DecoratorValue::Builtin(bf) if bf.kind == BuiltinFunctionKind::Error => {
529                let (selector, defined) = self.error_selector(flag, arg, bf)?;
530                // A defined error's selector is pushed left-aligned, ready to be stored as revert data
531                let mut word = [0u8; 32];
532                if defined {
533                    word[..4].copy_from_slice(&selector)
534                } else {
535                    word[28..].copy_from_slice(&selector)
536                }
537                Ok(U256::from_be_bytes(word))
538            }
539            DecoratorValue::Builtin(bf) => {
540                let value = Codegen::gen_builtin_bytecode(self.contract, bf, arg.span.clone())?;
541                Ok(U256::from_be_bytes(value.value.0))
542            }
543            DecoratorValue::Ident(ident) => match ident.as_str() {
544                "this" => Ok(U256::from_be_slice(self.ctx.this.as_slice())),
545                "caller" => Ok(U256::from_be_slice(self.ctx.caller.as_slice())),
546                _ => Err(invalid(format!("Unknown identifier `{ident}`; expected `this` or `caller`"))),
547            },
548            DecoratorValue::Str(_) => {
549                Err(invalid(String::from("Expected a hex or decimal literal, boolean, [CONSTANT], builtin call, `this` or `caller`")))
550            }
551        }
552    }
553
554    /// The account in a flag's `at =` argument, or the test itself without one
555    fn at_address(&self, flag: &DecoratorFlag, args: &FlagArgs) -> Result<Address, RunnerError> {
556        match args.named.get("at") {
557            Some(at) => self.address(flag, at),
558            None => Ok(self.ctx.this),
559        }
560    }
561
562    /// The exact value in a flag's second argument, or the inclusive range of its `min =` and `max =` arguments
563    fn expected_value(&self, flag: &DecoratorFlag, args: &FlagArgs, usage: &str) -> Result<ExpectedValue, RunnerError> {
564        let (min, max) = (args.named.get("min"), args.named.get("max"));
565        match (args.positional.get(1), min, max) {
566            (Some(exact), None, None) => Ok(ExpectedValue::Exact(self.word(flag, exact)?)),
567            (None, min, max) if min.is_some() || max.is_some() => {
568                let min = min.map(|arg| self.word(flag, arg)).transpose()?;
569                let max = max.map(|arg| self.word(flag, arg)).transpose()?;
570                if let (Some(min), Some(max)) = (min, max)
571                    && min > max
572                {
573                    return Err(decorator_error(
574                        flag,
575                        DecoratorErrorKind::InvalidValue,
576                        format!("`min` ({min:#x}) is larger than `max` ({max:#x})"),
577                        flag.span.clone(),
578                    ));
579                }
580                Ok(ExpectedValue::Range { min, max })
581            }
582            _ => Err(decorator_error(flag, DecoratorErrorKind::InvalidArguments, usage, flag.span.clone())),
583        }
584    }
585
586    /// How an outcome flag (`expect_revert`, `expect_partial_revert`, `expect_halt`, `expect_return`) expects the test to end
587    fn expected_outcome(&self, flag: &DecoratorFlag, outcome: OutcomeFlag) -> Result<ExpectedOutcome, RunnerError> {
588        let outcome = match outcome {
589            OutcomeFlag::Revert => {
590                let usage = "`expect_revert` takes nothing, or the exact revert data, e.g. expect_revert(__ERROR(Unauthorized))";
591                let args = FlagArgs::new(flag, 0..=usize::MAX, &[], usage)?;
592                ExpectedOutcome::Revert(self.optional_bytes(flag, &args.positional)?)
593            }
594            OutcomeFlag::PartialRevert => {
595                let usage = "`expect_partial_revert` takes the start of the revert data, e.g. expect_partial_revert(__ERROR(Unauthorized))";
596                let args = FlagArgs::new(flag, 1..=usize::MAX, &[], usage)?;
597                ExpectedOutcome::PartialRevert(self.bytes(flag, &args.positional)?)
598            }
599            OutcomeFlag::Halt => {
600                let usage = "`expect_halt` takes nothing, or a halt reason, e.g. expect_halt(\"OutOfGas\")";
601                let args = FlagArgs::new(flag, 0..=1, &[], usage)?;
602                let reason = match args.positional.first() {
603                    None => None,
604                    Some(DecoratorArg { value: DecoratorValue::Str(reason), .. }) => Some(reason.clone()),
605                    Some(arg) => return Err(decorator_error(flag, DecoratorErrorKind::InvalidValue, usage, arg.span.clone())),
606                };
607                ExpectedOutcome::Halt(reason)
608            }
609            OutcomeFlag::Return => {
610                let usage = "`expect_return` takes the exact return data, e.g. expect_return(0x01)";
611                let args = FlagArgs::new(flag, 0..=usize::MAX, &[], usage)?;
612                ExpectedOutcome::Return(Some(self.bytes(flag, &args.positional)?))
613            }
614        };
615        Ok(outcome)
616    }
617
618    /// The bytes of a byte sequence, or `None` when there are no arguments
619    fn optional_bytes(&self, flag: &DecoratorFlag, args: &[&DecoratorArg]) -> Result<Option<Vec<u8>>, RunnerError> {
620        if args.is_empty() { Ok(None) } else { self.bytes(flag, args).map(Some) }
621    }
622
623    /// An argument that must fit in a smaller unsigned integer, such as a `u64` gas limit
624    fn uint<T: TryFrom<U256>>(&self, flag: &DecoratorFlag, arg: &DecoratorArg) -> Result<T, RunnerError> {
625        let word = self.word(flag, arg)?;
626        T::try_from(word).map_err(|_| {
627            decorator_error(flag, DecoratorErrorKind::InvalidValue, format!("{word:#x} is too large for `{}`", flag.name), arg.span.clone())
628        })
629    }
630
631    /// An argument that must be an address
632    fn address(&self, flag: &DecoratorFlag, arg: &DecoratorArg) -> Result<Address, RunnerError> {
633        let word = self.word(flag, arg)?;
634        if word.bit_len() > Address::len_bytes() * 8 {
635            return Err(decorator_error(flag, DecoratorErrorKind::InvalidValue, format!("{word:#x} is not an address"), arg.span.clone()));
636        }
637        Ok(Address::from_word(word.into()))
638    }
639
640    /// The bytes a sequence of arguments stands for: selectors add 4 bytes, hex keeps its written
641    /// width, and every other value adds a 32-byte word
642    fn bytes(&self, flag: &DecoratorFlag, args: &[&DecoratorArg]) -> Result<Vec<u8>, RunnerError> {
643        let mut bytes = Vec::new();
644        for arg in args {
645            match &arg.value {
646                DecoratorValue::Str(s) => {
647                    let decoded = hex::decode(s.strip_prefix("0x").unwrap_or(s)).map_err(|_| {
648                        decorator_error(flag, DecoratorErrorKind::InvalidValue, format!("\"{s}\" is not a hex string"), arg.span.clone())
649                    })?;
650                    bytes.extend(decoded);
651                }
652                DecoratorValue::Hex(digits) => {
653                    // An odd number of digits is padded with a leading zero nibble
654                    let decoded = if digits.len() % 2 == 1 { hex::decode(format!("0{digits}")) } else { hex::decode(digits) };
655                    bytes.extend(decoded.expect("the lexer only produces hex digits"));
656                }
657                DecoratorValue::Builtin(bf) if bf.kind == BuiltinFunctionKind::FunctionSignature => {
658                    let word = self.word(flag, arg)?.to_be_bytes::<32>();
659                    bytes.extend(&word[28..]);
660                }
661                DecoratorValue::Builtin(bf) if bf.kind == BuiltinFunctionKind::Error => {
662                    bytes.extend(self.error_selector(flag, arg, bf)?.0);
663                }
664                _ => bytes.extend(self.word(flag, arg)?.to_be_bytes::<32>()),
665            }
666        }
667        Ok(bytes)
668    }
669
670    /// The selector of `__ERROR(..)`, and whether it names an error defined in the contract
671    fn error_selector(&self, flag: &DecoratorFlag, arg: &DecoratorArg, bf: &BuiltinFunctionCall) -> Result<([u8; 4], bool), RunnerError> {
672        let name = match bf.args.as_slice() {
673            [BuiltinFunctionArg::StringLiteral(s, _) | BuiltinFunctionArg::Identifier(s, _)] => s,
674            _ => {
675                return Err(decorator_error(
676                    flag,
677                    DecoratorErrorKind::InvalidValue,
678                    "__ERROR takes an error name or signature",
679                    arg.span.clone(),
680                ));
681            }
682        };
683        Ok(match self.contract.errors.iter().find(|e| &e.name == name) {
684            Some(error) => (error.selector, true),
685            None => (keccak256(name)[..4].try_into().expect("a selector is 4 bytes"), false),
686        })
687    }
688}
689
690/// The reason a `skip` flag gives, if any
691fn skip_reason(flag: &DecoratorFlag) -> Result<Option<String>, RunnerError> {
692    let usage = "`skip` takes nothing, or a reason, e.g. skip(\"needs a mainnet fork\")";
693    let args = FlagArgs::new(flag, 0..=1, &[], usage)?;
694    match args.positional.first() {
695        None => Ok(None),
696        Some(DecoratorArg { value: DecoratorValue::Str(reason), .. }) => Ok(Some(reason.clone())),
697        Some(arg) => Err(decorator_error(flag, DecoratorErrorKind::InvalidValue, usage, arg.span.clone())),
698    }
699}
700
701/// Whether a test's decorator skips it, with the reason when given. The runner checks this before
702/// it compiles the test, so a skipped test may be unfinished.
703pub fn skipped(decorator: Option<&Decorator>) -> Result<Option<Option<String>>, RunnerError> {
704    decorator
705        .and_then(|decorator| decorator.flags.iter().find(|flag| DecoratorKind::from_name(&flag.name) == Some(DecoratorKind::Skip)))
706        .map(skip_reason)
707        .transpose()
708}
709
710/// The only argument of a flag that takes one value
711fn single_arg<'a>(flag: &'a DecoratorFlag, takes: &str) -> Result<&'a DecoratorArg, RunnerError> {
712    let args = FlagArgs::new(flag, 1..=1, &[], &format!("`{}` takes {takes}", flag.name))?;
713    Ok(args.positional[0])
714}
715
716/// Fails when a flag that may appear only once was already given
717fn check_single_use<'a>(
718    seen: &mut HashMap<DecoratorKind, &'a DecoratorFlag>,
719    kind: DecoratorKind,
720    flag: &'a DecoratorFlag,
721) -> Result<(), RunnerError> {
722    match seen.insert(kind, flag) {
723        Some(first) => Err(decorator_error(
724            flag,
725            DecoratorErrorKind::Duplicate,
726            format!("`{}` may only be given once", flag.name),
727            joined_span(first, flag),
728        )),
729        None => Ok(()),
730    }
731}
732
733/// Fails when a flag that sets or checks an account was already given for the same address
734fn check_account_once<'a>(
735    seen: &mut HashMap<(DecoratorKind, Address), &'a DecoratorFlag>,
736    kind: DecoratorKind,
737    address: Address,
738    flag: &'a DecoratorFlag,
739) -> Result<(), RunnerError> {
740    match seen.insert((kind, address), flag) {
741        Some(first) => Err(decorator_error(
742            flag,
743            DecoratorErrorKind::Duplicate,
744            format!("`{}` is already given for {address}", flag.name),
745            joined_span(first, flag),
746        )),
747        None => Ok(()),
748    }
749}
750
751/// The source of two flags, so a duplicate points at both places
752fn joined_span(first: &DecoratorFlag, second: &DecoratorFlag) -> AstSpan {
753    AstSpan(first.span.0.iter().chain(&second.span.0).cloned().collect())
754}
755
756/// An error pointing at a decorator flag the runner cannot apply
757fn decorator_error(flag: &DecoratorFlag, kind: DecoratorErrorKind, hint: impl Into<String>, span: AstSpan) -> RunnerError {
758    RunnerError::InvalidDecorator(DecoratorError { flag: flag.name.clone(), kind, hint: hint.into(), span })
759}
760
761#[cfg(test)]
762#[test]
763fn test_decorator_names_map_to_their_kind() {
764    let mut names = std::collections::HashSet::new();
765    for (kind, name) in DecoratorKind::NAMES {
766        assert!(names.insert(name), "`{name}` is listed twice");
767        assert_eq!(DecoratorKind::from_name(name), Some(kind), "{name}");
768    }
769    assert_eq!(DecoratorKind::from_name("origin"), None);
770}