Test cfg for gas_budget::compute_unit_limit (anza-xyz#21)

Tests can now add compute_unit_limit as a config.

Followups
- Remove `DEFAULT_EXECUTION_BOUND_SOLANA_STDLIB_TEST` added in anza-xyz#18
- Parse GAS_BUDGET_HEAP_SIZE
- Parse GAS_BUDGET_MAX_CALL_DEPTH

Fixes: anza-xyz#20

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<BTreeSet<AttributePosition>> =
             Lazy::new(|| BTreeSet::from([AttributePosition::Function]));
+        static GAS_BUDGET_POSITIONS: Lazy<BTreeSet<AttributePosition>> =
+            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 {

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<MoveValue>,
     pub expected_failure: Option<ExpectedFailure>,
+    pub gas_budget: Option<GasBudgetParams>,
 }
 
 #[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,

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<ModuleIdent, UniqueMap<ConstantName, (Loc, Option<u64>)>>,
@@ -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<GasBudgetParams> = 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<GasBudgetParams> {
+    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<String, (Loc, Attribute)> = 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::<Vec<_>>();
+            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,

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)
     }