pass epoch along for cost eval

Author: brady-stacks

clarity/src/vm/ast/static_cost/mod.rs

@@ -2,6 +2,7 @@ mod trait_counter;
 use std::collections::HashMap;
 
 use clarity_types::types::{CharType, SequenceData};
+use stacks_common::types::StacksEpochId;
 pub use trait_counter::{
     TraitCount, TraitCountCollector, TraitCountContext, TraitCountPropagator, TraitCountVisitor,
 };
@@ -97,10 +98,12 @@ pub(crate) fn calculate_value_cost(value: &Value) -> Result<StaticCost, String>
 pub(crate) fn calculate_function_cost_from_native_function(
     native_function: NativeFunctions,
     arg_count: u64,
-    clarity_version: &ClarityVersion,
+    epoch: StacksEpochId,
 ) -> Result<StaticCost, String> {
+    // Derive clarity_version from epoch for lookup_reserved_functions
+    let clarity_version = ClarityVersion::default_for_epoch(epoch);
     let cost_function =
-        match lookup_reserved_functions(native_function.to_string().as_str(), clarity_version) {
+        match lookup_reserved_functions(native_function.to_string().as_str(), &clarity_version) {
             Some(CallableType::NativeFunction(_, _, cost_fn)) => cost_fn,
             Some(CallableType::NativeFunction205(_, _, cost_fn, _)) => cost_fn,
             Some(CallableType::SpecialFunction(_, _)) => return Ok(StaticCost::ZERO),
@@ -110,11 +113,22 @@ pub(crate) fn calculate_function_cost_from_native_function(
             }
         };
 
-    let cost = match clarity_version {
-        ClarityVersion::Clarity1 => cost_function.eval::<Costs1>(arg_count),
-        ClarityVersion::Clarity2 => cost_function.eval::<Costs2>(arg_count),
-        ClarityVersion::Clarity3 => cost_function.eval::<Costs3>(arg_count),
-        ClarityVersion::Clarity4 => cost_function.eval::<Costs4>(arg_count),
+    let cost = match epoch {
+        StacksEpochId::Epoch20 => cost_function.eval::<Costs1>(arg_count),
+        StacksEpochId::Epoch2_05 => cost_function.eval::<Costs2>(arg_count),
+        StacksEpochId::Epoch21
+        | StacksEpochId::Epoch22
+        | StacksEpochId::Epoch23
+        | StacksEpochId::Epoch24
+        | StacksEpochId::Epoch25
+        | StacksEpochId::Epoch30
+        | StacksEpochId::Epoch31
+        | StacksEpochId::Epoch32 => cost_function.eval::<Costs3>(arg_count),
+        StacksEpochId::Epoch33 => cost_function.eval::<Costs4>(arg_count),
+        StacksEpochId::Epoch10 => {
+            // fallback to costs 1 since epoch 1 doesn't have direct cost mapping
+            cost_function.eval::<Costs1>(arg_count)
+        }
     }
     .map_err(|e| format!("Cost calculation error: {:?}", e))?;
     Ok(StaticCost {

clarity/src/vm/costs/analysis.rs

@@ -234,7 +234,7 @@ pub fn static_cost(
     let epoch = env.global_context.epoch_id;
     let ast = make_ast(&contract_source, epoch, clarity_version)?;
 
-    static_cost_tree_from_ast(&ast, clarity_version)
+    static_cost_tree_from_ast(&ast, clarity_version, epoch)
 }
 
 /// same idea as `static_cost` but returns the root of the cost analysis tree for each function
@@ -265,14 +265,15 @@ pub fn static_cost_tree(
     let epoch = env.global_context.epoch_id;
     let ast = make_ast(&contract_source, epoch, clarity_version)?;
 
-    static_cost_tree_from_ast(&ast, clarity_version)
+    static_cost_tree_from_ast(&ast, clarity_version, epoch)
 }
 
 pub fn static_cost_from_ast(
     contract_ast: &crate::vm::ast::ContractAST,
     clarity_version: &ClarityVersion,
+    epoch: StacksEpochId,
 ) -> Result<HashMap<String, (StaticCost, Option<TraitCount>)>, String> {
-    let cost_trees_with_traits = static_cost_tree_from_ast(contract_ast, clarity_version)?;
+    let cost_trees_with_traits = static_cost_tree_from_ast(contract_ast, clarity_version, epoch)?;
 
     // Extract trait_count from the first entry (all entries have the same trait_count)
     let trait_count = cost_trees_with_traits
@@ -298,6 +299,7 @@ pub fn static_cost_from_ast(
 pub(crate) fn static_cost_tree_from_ast(
     ast: &crate::vm::ast::ContractAST,
     clarity_version: &ClarityVersion,
+    epoch: StacksEpochId,
 ) -> Result<HashMap<String, (CostAnalysisNode, Option<TraitCount>)>, String> {
     let exprs = &ast.expressions;
     let user_args = UserArgumentsContext::new();
@@ -313,7 +315,7 @@ pub(crate) fn static_cost_tree_from_ast(
     for expr in exprs {
         if let Some(function_name) = extract_function_name(expr) {
             let (_, cost_analysis_tree) =
-                build_cost_analysis_tree(expr, &user_args, &costs_map, clarity_version)?;
+                build_cost_analysis_tree(expr, &user_args, &costs_map, clarity_version, epoch)?;
             costs.insert(function_name, Some(cost_analysis_tree));
         }
     }
@@ -353,6 +355,7 @@ pub fn build_cost_analysis_tree(
     user_args: &UserArgumentsContext,
     cost_map: &HashMap<String, Option<StaticCost>>,
     clarity_version: &ClarityVersion,
+    epoch: StacksEpochId,
 ) -> Result<(Option<String>, CostAnalysisNode), String> {
     match &expr.expr {
         SymbolicExpressionType::List(list) => {
@@ -364,6 +367,7 @@ pub fn build_cost_analysis_tree(
                             user_args,
                             cost_map,
                             clarity_version,
+                            epoch,
                         )?;
                     Ok((Some(returned_function_name), cost_analysis_tree))
                 } else {
@@ -372,12 +376,18 @@ pub fn build_cost_analysis_tree(
                         user_args,
                         cost_map,
                         clarity_version,
+                        epoch,
                     )?;
                     Ok((None, cost_analysis_tree))
                 }
             } else {
-                let cost_analysis_tree =
-                    build_listlike_cost_analysis_tree(list, user_args, cost_map, clarity_version)?;
+                let cost_analysis_tree = build_listlike_cost_analysis_tree(
+                    list,
+                    user_args,
+                    cost_map,
+                    clarity_version,
+                    epoch,
+                )?;
                 Ok((None, cost_analysis_tree))
             }
         }
@@ -439,6 +449,7 @@ fn build_function_definition_cost_analysis_tree(
     _user_args: &UserArgumentsContext,
     cost_map: &HashMap<String, Option<StaticCost>>,
     clarity_version: &ClarityVersion,
+    epoch: StacksEpochId,
 ) -> Result<(String, CostAnalysisNode), String> {
     let define_type = list[0]
         .match_atom()
@@ -476,7 +487,7 @@ fn build_function_definition_cost_analysis_tree(
 
     // Process the function body with the function's user arguments context
     let (_, body_tree) =
-        build_cost_analysis_tree(body, &function_user_args, cost_map, clarity_version)?;
+        build_cost_analysis_tree(body, &function_user_args, cost_map, clarity_version, epoch)?;
     children.push(body_tree);
 
     // Get the function name from the signature
@@ -508,20 +519,22 @@ fn build_listlike_cost_analysis_tree(
     user_args: &UserArgumentsContext,
     cost_map: &HashMap<String, Option<StaticCost>>,
     clarity_version: &ClarityVersion,
+    epoch: StacksEpochId,
 ) -> Result<CostAnalysisNode, String> {
     let mut children = Vec::new();
 
     // Build children for all exprs
     for expr in exprs[1..].iter() {
-        let (_, child_tree) = build_cost_analysis_tree(expr, user_args, cost_map, clarity_version)?;
+        let (_, child_tree) =
+            build_cost_analysis_tree(expr, user_args, cost_map, clarity_version, epoch)?;
         children.push(child_tree);
     }
 
     let (expr_node, cost) = 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)?;
+                build_cost_analysis_tree(&exprs[0], user_args, cost_map, clarity_version, epoch)?;
             // 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
@@ -537,7 +550,7 @@ fn build_listlike_cost_analysis_tree(
                 let cost = calculate_function_cost_from_native_function(
                     native_function,
                     children.len() as u64,
-                    clarity_version,
+                    epoch,
                 )?;
                 (CostExprNode::NativeFunction(native_function), cost)
             } else {
@@ -608,7 +621,7 @@ mod tests {
         let user_args = UserArgumentsContext::new();
         let expr = &exprs[0];
         let (_, cost_analysis_tree) =
-            build_cost_analysis_tree(&expr, &user_args, &cost_map, clarity_version)?;
+            build_cost_analysis_tree(&expr, &user_args, &cost_map, clarity_version, epoch)?;
 
         let summing_cost = calculate_total_cost_with_branching(&cost_analysis_tree);
         Ok(summing_cost.into())
@@ -620,7 +633,7 @@ mod tests {
     ) -> Result<HashMap<String, StaticCost>, String> {
         let epoch = StacksEpochId::latest();
         let ast = make_ast(source, epoch, clarity_version)?;
-        let costs = static_cost_from_ast(&ast, clarity_version)?;
+        let costs = static_cost_from_ast(&ast, clarity_version, epoch)?;
         Ok(costs
             .into_iter()
             .map(|(name, (cost, _trait_count))| (name, cost))
@@ -713,9 +726,15 @@ mod tests {
         let expr = &ast.expressions[0];
         let user_args = UserArgumentsContext::new();
         let cost_map = HashMap::new(); // Empty cost map for tests
-        let (_, cost_tree) =
-            build_cost_analysis_tree(expr, &user_args, &cost_map, &ClarityVersion::Clarity3)
-                .unwrap();
+        let epoch = StacksEpochId::Epoch32;
+        let (_, cost_tree) = build_cost_analysis_tree(
+            expr,
+            &user_args,
+            &cost_map,
+            &ClarityVersion::Clarity3,
+            epoch,
+        )
+        .unwrap();
 
         // Root should be an If node
         assert!(matches!(
@@ -760,9 +779,15 @@ mod tests {
         let expr = &ast.expressions[0];
         let user_args = UserArgumentsContext::new();
         let cost_map = HashMap::new(); // Empty cost map for tests
-        let (_, cost_tree) =
-            build_cost_analysis_tree(expr, &user_args, &cost_map, &ClarityVersion::Clarity3)
-                .unwrap();
+        let epoch = StacksEpochId::Epoch32;
+        let (_, cost_tree) = build_cost_analysis_tree(
+            expr,
+            &user_args,
+            &cost_map,
+            &ClarityVersion::Clarity3,
+            epoch,
+        )
+        .unwrap();
 
         assert!(matches!(
             cost_tree.expr,
@@ -793,9 +818,15 @@ mod tests {
         let expr = &ast.expressions[0];
         let user_args = UserArgumentsContext::new();
         let cost_map = HashMap::new(); // Empty cost map for tests
-        let (_, cost_tree) =
-            build_cost_analysis_tree(expr, &user_args, &cost_map, &ClarityVersion::Clarity3)
-                .unwrap();
+        let epoch = StacksEpochId::Epoch32;
+        let (_, cost_tree) = build_cost_analysis_tree(
+            expr,
+            &user_args,
+            &cost_map,
+            &ClarityVersion::Clarity3,
+            epoch,
+        )
+        .unwrap();
 
         assert!(matches!(
             cost_tree.expr,
@@ -816,9 +847,15 @@ mod tests {
         let expr = &ast.expressions[0];
         let user_args = UserArgumentsContext::new();
         let cost_map = HashMap::new(); // Empty cost map for tests
-        let (_, cost_tree) =
-            build_cost_analysis_tree(expr, &user_args, &cost_map, &ClarityVersion::Clarity3)
-                .unwrap();
+        let epoch = StacksEpochId::Epoch32;
+        let (_, cost_tree) = build_cost_analysis_tree(
+            expr,
+            &user_args,
+            &cost_map,
+            &ClarityVersion::Clarity3,
+            epoch,
+        )
+        .unwrap();
 
         assert_eq!(cost_tree.children.len(), 3);
 

clarity/src/vm/tests/analysis.rs

@@ -33,7 +33,8 @@ fn test_build_cost_analysis_tree_function_definition() {
     let cost_map = HashMap::new();
 
     let clarity_version = ClarityVersion::Clarity3;
-    let result = build_cost_analysis_tree(expr, &user_args, &cost_map, &clarity_version);
+    let epoch = StacksEpochId::Epoch32;
+    let result = build_cost_analysis_tree(expr, &user_args, &cost_map, &clarity_version, epoch);
 
     match result {
         Ok((function_name, node)) => {
@@ -71,7 +72,7 @@ fn test_dependent_function_calls() {
         epoch,
     )
     .unwrap();
-    let function_map = static_cost_from_ast(&ast, &ClarityVersion::Clarity3).unwrap();
+    let function_map = static_cost_from_ast(&ast, &ClarityVersion::Clarity3, epoch).unwrap();
 
     let (add_one_cost, _) = function_map.get("add-one").unwrap();
     let (somefunc_cost, _) = function_map.get("somefunc").unwrap();
@@ -103,7 +104,8 @@ fn test_get_trait_count_direct() {
     )
     .unwrap();
 
-    let costs = static_cost_tree_from_ast(&ast, &ClarityVersion::Clarity3).unwrap();
+    let costs =
+        static_cost_tree_from_ast(&ast, &ClarityVersion::Clarity3, StacksEpochId::Epoch32).unwrap();
 
     // Extract trait_count from the result (all entries have the same trait_count)
     let trait_count = costs
@@ -133,15 +135,11 @@ fn test_trait_counting() {
 (define-private (send (trait <trait-name>) (addr principal)) (trait addr))
 "#;
     let contract_id = QualifiedContractIdentifier::local("trait-counting").unwrap();
-    let ast = crate::vm::ast::build_ast(
-        &contract_id,
-        src,
-        &mut (),
-        ClarityVersion::Clarity3,
-        StacksEpochId::Epoch32,
-    )
-    .unwrap();
-    let static_cost = static_cost_from_ast(&ast, &ClarityVersion::Clarity3)
+    let epoch = StacksEpochId::Epoch32;
+    let ast =
+        crate::vm::ast::build_ast(&contract_id, src, &mut (), ClarityVersion::Clarity3, epoch)
+            .unwrap();
+    let static_cost = static_cost_from_ast(&ast, &ClarityVersion::Clarity3, epoch)
         .unwrap()
         .clone();
     let send_trait_count_map = static_cost.get("send").unwrap().1.clone().unwrap();
@@ -230,7 +228,7 @@ fn test_pox_4_costs() {
     )
     .expect("Failed to build AST from pox-4.clar");
 
-    let cost_map = static_cost_from_ast(&ast, &clarity_version)
+    let cost_map = static_cost_from_ast(&ast, &clarity_version, epoch)
         .expect("Failed to get static cost analysis for pox-4.clar");
 
     // Check some functions in the cost map