test larger contract analysis with pox-4

Author: brady-stacks

clarity/src/vm/costs/analysis.rs

@@ -953,30 +953,81 @@ fn build_listlike_cost_analysis_tree(
         children.push(child_tree);
     }
 
-    let function_name = get_function_name(&exprs[0])?;
-    // Try to lookup the function as a native function first
-    let (expr_node, cost) = if let Some(native_function) =
-        NativeFunctions::lookup_by_name_at_version(function_name.as_str(), clarity_version)
-    {
-        CostExprNode::NativeFunction(native_function);
-        let cost = calculate_function_cost_from_native_function(
-            native_function,
-            children.len() as u64,
-            clarity_version,
-        )?;
-        (CostExprNode::NativeFunction(native_function), cost)
-    } else {
-        // If not a native function, treat as user-defined function and look it up
-        let expr_node = CostExprNode::UserFunction(function_name.clone());
-        let cost = calculate_function_cost(function_name.to_string(), cost_map, clarity_version)?;
-        (expr_node, cost)
+    // Try to get function name from first element
+    let (expr_node, cost, function_name_opt) = match get_function_name(&exprs[0]) {
+        Ok(function_name) => {
+            // Try to lookup the function as a native function first
+            if let Some(native_function) =
+                NativeFunctions::lookup_by_name_at_version(function_name.as_str(), clarity_version)
+            {
+                let cost = calculate_function_cost_from_native_function(
+                    native_function,
+                    children.len() as u64,
+                    clarity_version,
+                )?;
+                (
+                    CostExprNode::NativeFunction(native_function),
+                    cost,
+                    Some(function_name),
+                )
+            } else {
+                // If not a native function, treat as user-defined function and look it up
+                let expr_node = CostExprNode::UserFunction(function_name.clone());
+                let cost =
+                    calculate_function_cost(function_name.to_string(), cost_map, clarity_version)?;
+                (expr_node, cost, Some(function_name))
+            }
+        }
+        Err(_) => {
+            // First element is not an atom - it might be a List that needs to be recursively analyzed
+            match &exprs[0].expr {
+                SymbolicExpressionType::List(_) => {
+                    // Recursively analyze the nested list structure
+                    let (_, nested_tree) =
+                        build_cost_analysis_tree(&exprs[0], user_args, cost_map, clarity_version)?;
+                    // Add the nested tree as a child (its cost will be included when summing children)
+                    children.insert(0, nested_tree);
+                    // The root cost is zero - the actual cost comes from the nested expression
+                    let expr_node = CostExprNode::Atom(ClarityName::from("nested-expression"));
+                    (expr_node, StaticCost::ZERO, None)
+                }
+                SymbolicExpressionType::Atom(name) => {
+                    // It's an atom but not a function name - treat as atom with zero cost
+                    (CostExprNode::Atom(name.clone()), StaticCost::ZERO, None)
+                }
+                SymbolicExpressionType::AtomValue(value) => {
+                    // It's an atom value - calculate its cost
+                    let cost = calculate_value_cost(value)?;
+                    (CostExprNode::AtomValue(value.clone()), cost, None)
+                }
+                SymbolicExpressionType::TraitReference(trait_name, _trait_definition) => (
+                    CostExprNode::TraitReference(trait_name.clone()),
+                    StaticCost::ZERO,
+                    None,
+                ),
+                SymbolicExpressionType::Field(field_identifier) => (
+                    CostExprNode::FieldIdentifier(field_identifier.clone()),
+                    StaticCost::ZERO,
+                    None,
+                ),
+                SymbolicExpressionType::LiteralValue(value) => {
+                    let cost = calculate_value_cost(value)?;
+                    // TODO not sure if LiteralValue is needed in the CostExprNode types
+                    (CostExprNode::AtomValue(value.clone()), cost, None)
+                }
+            }
+        }
     };
 
     // Handle special cases for string arguments to functions that include their processing cost
-    if FUNCTIONS_WITH_ZERO_STRING_ARG_COST.contains(&function_name.as_str()) {
-        for child in &mut children {
-            if let CostExprNode::AtomValue(Value::Sequence(SequenceData::String(_))) = &child.expr {
-                child.cost = StaticCost::ZERO;
+    if let Some(function_name) = &function_name_opt {
+        if FUNCTIONS_WITH_ZERO_STRING_ARG_COST.contains(&function_name.as_str()) {
+            for child in &mut children {
+                if let CostExprNode::AtomValue(Value::Sequence(SequenceData::String(_))) =
+                    &child.expr
+                {
+                    child.cost = StaticCost::ZERO;
+                }
             }
         }
     }

clarity/src/vm/tests/analysis.rs

@@ -1,5 +1,6 @@
 // TODO: This needs work to get the dynamic vs static testing working
 use std::collections::HashMap;
+use std::path::Path;
 
 use rstest::rstest;
 use stacks_common::types::StacksEpochId;
@@ -314,3 +315,55 @@ fn execute_contract_function_and_get_cost(
         runtime: final_cost.runtime - initial_cost.runtime,
     }
 }
+
+#[test]
+fn test_pox_4_costs() {
+    let workspace_root = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap();
+    let pox_4_path = workspace_root
+        .join("contrib")
+        .join("boot-contracts-unit-tests")
+        .join("boot_contracts")
+        .join("pox-4.clar");
+    let contract_source = std::fs::read_to_string(&pox_4_path)
+        .unwrap_or_else(|e| panic!("Failed to read pox-4.clar file at {:?}: {}", pox_4_path, e));
+
+    let contract_id = QualifiedContractIdentifier::transient();
+    let epoch = StacksEpochId::Epoch32;
+    let clarity_version = ClarityVersion::Clarity3;
+
+    let ast = crate::vm::ast::build_ast(
+        &contract_id,
+        &contract_source,
+        &mut (),
+        clarity_version,
+        epoch,
+    )
+    .expect("Failed to build AST from pox-4.clar");
+
+    let cost_map = static_cost_from_ast(&ast, &clarity_version)
+        .expect("Failed to perform static cost analysis on pox-4.clar");
+
+    // Check some functions in the cost map
+    let key_functions = vec![
+        "stack-stx",
+        "delegate-stx",
+        "get-stacker-info",
+        "current-pox-reward-cycle",
+        "stack-aggregation-commit",
+        "stack-increase",
+        "stack-extend",
+    ];
+
+    for function_name in key_functions {
+        assert!(
+            cost_map.contains_key(function_name),
+            "Expected function '{}' to be present in cost map",
+            function_name
+        );
+
+        let (_cost, _trait_count) = cost_map.get(function_name).expect(&format!(
+            "Failed to get cost for function '{}'",
+            function_name
+        ));
+    }
+}