diff --git a/external-crates/move/crates/move-compiler/src/shared/known_attributes.rs b/external-crates/move/crates/move-compiler/src/shared/known_attributes.rs index 559e029a9696c..e1638d15b5849 100644 --- a/external-crates/move/crates/move-compiler/src/shared/known_attributes.rs +++ b/external-crates/move/crates/move-compiler/src/shared/known_attributes.rs @@ -35,6 +35,8 @@ pub enum TestingAttribute { Test, // This test is expected to fail ExpectedFailure, + // Solana gas budget config + GasBudget, } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] @@ -86,6 +88,7 @@ impl KnownAttribute { TestingAttribute::TEST => TestingAttribute::Test.into(), TestingAttribute::TEST_ONLY => TestingAttribute::TestOnly.into(), TestingAttribute::EXPECTED_FAILURE => TestingAttribute::ExpectedFailure.into(), + TestingAttribute::GAS_BUDGET => TestingAttribute::GasBudget.into(), VerificationAttribute::VERIFY_ONLY => VerificationAttribute::VerifyOnly.into(), NativeAttribute::BYTECODE_INSTRUCTION => NativeAttribute::BytecodeInstruction.into(), DiagnosticAttribute::ALLOW => DiagnosticAttribute::Allow.into(), @@ -133,12 +136,17 @@ impl TestingAttribute { pub const MAJOR_STATUS_NAME: &'static str = "major_status"; pub const MINOR_STATUS_NAME: &'static str = "minor_status"; pub const ERROR_LOCATION: &'static str = "location"; + pub const GAS_BUDGET: &'static str = "gas_budget"; + pub const GAS_BUDGET_COMPUTE_UNIT_LIMIT: &'static str = "compute_unit_limit"; + pub const GAS_BUDGET_HEAP_SIZE: &'static str = "heap_size"; + pub const GAS_BUDGET_MAX_CALL_DEPTH: &'static str = "max_call_depth"; pub const fn name(&self) -> &str { match self { Self::Test => Self::TEST, Self::TestOnly => Self::TEST_ONLY, Self::ExpectedFailure => Self::EXPECTED_FAILURE, + Self::GasBudget => Self::GAS_BUDGET, } } @@ -158,10 +166,13 @@ impl TestingAttribute { Lazy::new(|| BTreeSet::from([AttributePosition::Function])); static EXPECTED_FAILURE_POSITIONS: Lazy> = Lazy::new(|| BTreeSet::from([AttributePosition::Function])); + static GAS_BUDGET_POSITIONS: Lazy> = + Lazy::new(|| BTreeSet::from([AttributePosition::Function])); match self { TestingAttribute::TestOnly => &TEST_ONLY_POSITIONS, TestingAttribute::Test => &TEST_POSITIONS, TestingAttribute::ExpectedFailure => &EXPECTED_FAILURE_POSITIONS, + TestingAttribute::GasBudget => &GAS_BUDGET_POSITIONS, } } @@ -174,6 +185,13 @@ impl TestingAttribute { Self::MAJOR_STATUS_NAME, ] } + pub fn gas_budget_cases() -> &'static [&'static str] { + &[ + Self::GAS_BUDGET_COMPUTE_UNIT_LIMIT, + Self::GAS_BUDGET_HEAP_SIZE, + Self::GAS_BUDGET_MAX_CALL_DEPTH, + ] + } } impl VerificationAttribute { diff --git a/external-crates/move/crates/move-compiler/src/unit_test/mod.rs b/external-crates/move/crates/move-compiler/src/unit_test/mod.rs index 3506f2a9ce571..93e521b4e03e9 100644 --- a/external-crates/move/crates/move-compiler/src/unit_test/mod.rs +++ b/external-crates/move/crates/move-compiler/src/unit_test/mod.rs @@ -34,6 +34,7 @@ pub struct TestCase { pub test_name: TestName, pub arguments: Vec, pub expected_failure: Option, + pub gas_budget: Option, } #[derive(Debug, Clone)] @@ -46,6 +47,13 @@ pub enum ExpectedFailure { ExpectedWithError(ExpectedMoveError), } +#[derive(Debug, Clone)] +pub struct GasBudgetParams { + pub compute_budget:u64, + pub heap_size:u64, + pub max_call_depth:u64, +} + #[derive(Debug, Clone, Ord, PartialOrd, PartialEq, Eq)] pub struct ExpectedMoveError( pub StatusCode, 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 99f670f5ac4df..32d2bfe2c7414 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 @@ -24,6 +24,8 @@ use move_ir_types::location::Loc; use move_symbol_pool::Symbol; use std::collections::BTreeMap; +use super::GasBudgetParams; + struct Context<'env> { env: &'env mut CompilationEnv, constants: UniqueMap)>>, @@ -124,7 +126,7 @@ fn build_test_info<'func>( let test_attribute_opt = get_attrs(TestingAttribute::Test); let abort_attribute_opt = get_attrs(TestingAttribute::ExpectedFailure); let test_only_attribute_opt = get_attrs(TestingAttribute::TestOnly); - + let gas_budget_attr = get_attrs(TestingAttribute::GasBudget); let test_attribute = match test_attribute_opt { None => { // expected failures cannot be annotated on non-#[test] functions @@ -182,11 +184,13 @@ fn build_test_info<'func>( None => None, Some(abort_attribute) => parse_failure_attribute(context, abort_attribute), }; + let gas_budget: Option = parse_gas_budget_attribute(context, gas_budget_attr); Some(TestCase { test_name: fn_name.to_string(), arguments, expected_failure, + gas_budget, }) } @@ -487,6 +491,83 @@ fn parse_failure_attribute( } } +fn parse_gas_budget_attribute(context: &mut Context, gas_budget_attr: Option<&E::Attribute>) -> Option { + let mut assigned_gas_budget = GasBudgetParams { + compute_budget:1, + heap_size: 1, + max_call_depth: 1, + }; + if gas_budget_attr.is_none(){ + return Some(assigned_gas_budget); + } + let sp!(aloc, expected_attr) = match gas_budget_attr { + Some(gb) => gb, + None => { + unreachable!() + } + }; + + use E::Attribute_ as EA; + match expected_attr { + EA::Name(_) => { + return Some(assigned_gas_budget); + } + EA::Assigned(_, _) => { + return Some(assigned_gas_budget); + } + EA::Parameterized(sp!(_, nm), attrs) => { + assert!( + nm.as_str() == TestingAttribute::GasBudget.name(), + "ICE: expected failure attribute {nm} must have the right name" + ); + let mut attrs: BTreeMap = attrs + .key_cloned_iter() + .map(|(sp!(kloc, k_), v)| (k_.to_string(), (kloc, v.clone()))) + .collect(); + let mut gas_budget_kind_vec = TestingAttribute::gas_budget_cases() + .iter() + .filter_map(|k| { + let k = k.to_string(); + let attr_opt = attrs.remove(&k)?; + Some((k, attr_opt)) + }) + .collect::>(); + if gas_budget_kind_vec.len() != 1 { + let invalid_attr_msg = format!( + "Invalid #[gas_budget(...)] attribute, expected 1 failure kind but found {}. Expected one of: {}", + gas_budget_kind_vec.len(), + TestingAttribute::gas_budget_cases().to_vec().join(", ") + ); + context + .env + .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!(), + }; + return Some(assigned_gas_budget); + } + } +} + fn check_attribute_unassigned( context: &mut Context, kind: &str, 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 2dfa2e047471a..e3e089144ae1a 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,6 +53,7 @@ module std::bit_vector_tests { } #[test] + #[gas_budget(compute_unit_limit=10000000)] fun test_set_bit_and_index_basic() { test_bitvector_set_unset_of_size(8) }