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#[derive(Debug, Clone, Default, PartialEq, Eq)]
15pub struct TestSpec {
16 pub value: U256,
18 pub calldata: Vec<u8>,
20 pub storage: Vec<StorageSlot>,
22 pub balances: Vec<AccountValue<U256>>,
24 pub code: Vec<AccountValue<Vec<u8>>>,
26 pub nonces: Vec<AccountValue<u64>>,
28 pub env: EnvOverrides,
30 pub expectations: Expectations,
32}
33
34#[derive(Debug, Clone, Default, PartialEq, Eq)]
36pub struct EnvOverrides {
37 pub caller: Option<Address>,
39 pub gas_limit: Option<u64>,
41 pub gas_price: Option<u128>,
43 pub block_number: Option<U256>,
45 pub timestamp: Option<U256>,
47 pub basefee: Option<u64>,
49 pub blob_basefee: Option<u128>,
51 pub coinbase: Option<Address>,
53 pub chain_id: Option<u64>,
55 pub prevrandao: Option<B256>,
57 pub slot_number: Option<u64>,
59}
60
61#[derive(Debug, Clone, PartialEq, Eq)]
63pub struct StorageSlot {
64 pub address: Address,
66 pub slot: U256,
68 pub value: U256,
70}
71
72#[derive(Debug, Clone, PartialEq, Eq)]
74pub struct AccountValue<T> {
75 pub address: Address,
77 pub value: T,
79}
80
81#[derive(Debug, Clone, Copy)]
83pub struct SpecContext {
84 pub this: Address,
86 pub caller: Address,
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub enum ExpectedValue {
93 Exact(U256),
95 Range { min: Option<U256>, max: Option<U256> },
97}
98
99impl ExpectedValue {
100 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 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#[derive(Debug, Clone, PartialEq, Eq)]
122pub struct ExpectedSlot {
123 pub address: Address,
125 pub slot: U256,
127 pub expected: ExpectedValue,
129}
130
131#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
133enum DecoratorKind {
134 Calldata,
136 Value,
138 Store,
140 Deal,
142 Etch,
144 Nonce,
146 Caller,
148 Coinbase,
150 GasLimit,
152 GasPrice,
154 BlockNumber,
156 Timestamp,
158 Basefee,
160 BlobBasefee,
162 ChainId,
164 Prevrandao,
166 SlotNumber,
168 Outcome(OutcomeFlag),
170 ExpectGas,
172 ExpectStorage,
174 ExpectBalance,
176 Skip,
178}
179
180#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
182enum OutcomeFlag {
183 Revert,
185 PartialRevert,
187 Halt,
189 Return,
191}
192
193impl DecoratorKind {
194 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 fn from_name(name: &str) -> Option<Self> {
225 Self::NAMES.iter().find(|(_, n)| *n == name).map(|(kind, _)| *kind)
226 }
227
228 fn is_repeatable(self) -> bool {
230 matches!(self, Self::Store | Self::Deal | Self::Etch | Self::Nonce | Self::ExpectStorage | Self::ExpectBalance)
231 }
232
233 fn checks_post_state(self) -> bool {
235 matches!(self, Self::ExpectStorage | Self::ExpectBalance)
236 }
237}
238
239impl TestSpec {
240 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 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 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 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 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
479struct 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
509struct Evaluator<'a> {
511 contract: &'a Contract,
512 ctx: SpecContext,
513}
514
515impl Evaluator<'_> {
516 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 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 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 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 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 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 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 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 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 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 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
690fn 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
701pub 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
710fn 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
716fn 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
733fn 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
751fn joined_span(first: &DecoratorFlag, second: &DecoratorFlag) -> AstSpan {
753 AstSpan(first.span.0.iter().chain(&second.span.0).cloned().collect())
754}
755
756fn 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}