diff --git a/external-crates/move/crates/move-compiler/src/unit_test/plan_builder.rs b/external-crates/move/crates/move-compiler/src/unit_test/plan_builder.rs index 32d2bfe2c7414..0ea12594d6743 100644 --- a/external-crates/move/crates/move-compiler/src/unit_test/plan_builder.rs +++ b/external-crates/move/crates/move-compiler/src/unit_test/plan_builder.rs @@ -532,9 +532,10 @@ fn parse_gas_budget_attribute(context: &mut Context, gas_budget_attr: Option<&E: Some((k, attr_opt)) }) .collect::>(); - if gas_budget_kind_vec.len() != 1 { + assert!(attrs.is_empty(), "Invalid #[gas_budget(...)] attribute"); + if gas_budget_kind_vec.len() > 3 { let invalid_attr_msg = format!( - "Invalid #[gas_budget(...)] attribute, expected 1 failure kind but found {}. Expected one of: {}", + "Invalid #[gas_budget(...)] attribute, found {} attributes. Expected only: {}", gas_budget_kind_vec.len(), TestingAttribute::gas_budget_cases().to_vec().join(", ") ); @@ -543,26 +544,27 @@ fn parse_gas_budget_attribute(context: &mut Context, gas_budget_attr: Option<&E: .add_diag(diag!(Attributes::InvalidValue, (*aloc, invalid_attr_msg))); return None; } - let (gas_budget_kind, (attr_loc, attr)) = gas_budget_kind_vec.pop().unwrap(); - match gas_budget_kind.as_str() { - TestingAttribute::GAS_BUDGET_COMPUTE_UNIT_LIMIT => { - let (value_name_loc, attr_value) = get_assigned_attribute( - context, - TestingAttribute::GAS_BUDGET_COMPUTE_UNIT_LIMIT, - attr_loc, - attr, - )?; - let (_, _, u) = - convert_constant_value_u64_constant_or_value( - context, - value_name_loc, - &attr_value, - )?; - // TODO: Do some sanity check that u shouldn't be larger than a max value. - assigned_gas_budget.compute_budget = u; - } - _ => unreachable!(), - }; + while !gas_budget_kind_vec.is_empty() { + let (gas_budget_kind, (attr_loc, attr)) = gas_budget_kind_vec.pop().unwrap(); + match gas_budget_kind.as_str() { + TestingAttribute::GAS_BUDGET_COMPUTE_UNIT_LIMIT => { + let u = get_assigned_attribute_as_u64(context, TestingAttribute::GAS_BUDGET_COMPUTE_UNIT_LIMIT, attr_loc, attr)?; + // TODO: Do some sanity check that u shouldn't be larger than a max value. + assigned_gas_budget.compute_budget = u; + } + TestingAttribute::GAS_BUDGET_HEAP_SIZE => { + let u = get_assigned_attribute_as_u64(context, TestingAttribute::GAS_BUDGET_HEAP_SIZE, attr_loc, attr)?; + // TODO: Do some sanity check that u shouldn't be larger than a max value. + assigned_gas_budget.heap_size = u; + } + TestingAttribute::GAS_BUDGET_MAX_CALL_DEPTH => { + let u = get_assigned_attribute_as_u64(context, TestingAttribute::GAS_BUDGET_MAX_CALL_DEPTH, attr_loc, attr)?; + // TODO: Do some sanity check that u shouldn't be larger than a max value. + assigned_gas_budget.max_call_depth = u; + } + _ => return None, + }; + } return Some(assigned_gas_budget); } } @@ -622,6 +624,26 @@ fn get_assigned_attribute( } } +fn get_assigned_attribute_as_u64( + context: &mut Context, + kind: &str, + attr_loc: Loc, + attr: Attribute, +) -> Option { + let (value_name_loc, attr_value) = get_assigned_attribute( + context, + kind, + attr_loc, + attr, + )?; + let (_, _, u) = convert_constant_value_u64_constant_or_value( + context, + value_name_loc, + &attr_value, + )?; + return Some(u); +} + fn convert_location(context: &mut Context, attr_loc: Loc, attr: Attribute) -> Option { use E::AttributeValue_ as EAV; let (loc, value) = diff --git a/external-crates/move/crates/move-stdlib/tests/bit_vector_tests.move b/external-crates/move/crates/move-stdlib/tests/bit_vector_tests.move index e3e089144ae1a..5e1c4297fa1e6 100644 --- a/external-crates/move/crates/move-stdlib/tests/bit_vector_tests.move +++ b/external-crates/move/crates/move-stdlib/tests/bit_vector_tests.move @@ -53,7 +53,7 @@ module std::bit_vector_tests { } #[test] - #[gas_budget(compute_unit_limit=10000000)] + #[gas_budget(compute_unit_limit=10000000, heap_size=1000, max_call_depth=10)] fun test_set_bit_and_index_basic() { test_bitvector_set_unset_of_size(8) }