make static_cost functions return trait counts also

Author: brady-stacks

clarity/src/vm/costs/analysis.rs

@@ -207,11 +207,11 @@ fn make_ast(
 
 /// STatic execution cost for functions within Environment
 /// returns the top level cost for specific functions
-/// {function_name: cost}
+/// {some-function-name: (CostAnalysisNode, Some({some-function-name: (1,1)}))}
 pub fn static_cost(
     env: &mut Environment,
     contract_identifier: &QualifiedContractIdentifier,
-) -> Result<HashMap<String, StaticCost>, String> {
+) -> Result<HashMap<String, (CostAnalysisNode, Option<TraitCount>)>, String> {
     let contract_source = env
         .global_context
         .database
@@ -234,19 +234,15 @@ pub fn static_cost(
     let epoch = env.global_context.epoch_id;
     let ast = make_ast(&contract_source, epoch, clarity_version)?;
 
-    let costs = static_cost_from_ast(&ast, clarity_version)?;
-    Ok(costs
-        .into_iter()
-        .map(|(name, (cost, _trait_count))| (name, cost))
-        .collect())
+    static_cost_tree_from_ast(&ast, clarity_version)
 }
 
 /// same idea as `static_cost` but returns the root of the cost analysis tree for each function
 /// Useful if you need to analyze specific nodes in the cost tree
 pub fn static_cost_tree(
     env: &mut Environment,
     contract_identifier: &QualifiedContractIdentifier,
-) -> Result<HashMap<String, CostAnalysisNode>, String> {
+) -> Result<HashMap<String, (CostAnalysisNode, Option<TraitCount>)>, String> {
     let contract_source = env
         .global_context
         .database
@@ -276,16 +272,23 @@ pub fn static_cost_from_ast(
     contract_ast: &crate::vm::ast::ContractAST,
     clarity_version: &ClarityVersion,
 ) -> Result<HashMap<String, (StaticCost, Option<TraitCount>)>, String> {
-    let cost_trees = static_cost_tree_from_ast(contract_ast, clarity_version)?;
+    let cost_trees_with_traits = static_cost_tree_from_ast(contract_ast, clarity_version)?;
 
-    let trait_count = get_trait_count(&cost_trees);
-    let costs: HashMap<String, StaticCost> = cost_trees
+    // Extract trait_count from the first entry (all entries have the same trait_count)
+    let trait_count = cost_trees_with_traits
+        .values()
+        .next()
+        .and_then(|(_, trait_count)| trait_count.clone());
+
+    // Convert CostAnalysisNode to StaticCost
+    let costs: HashMap<String, StaticCost> = cost_trees_with_traits
         .into_iter()
-        .map(|(name, cost_analysis_node)| {
+        .map(|(name, (cost_analysis_node, _))| {
             let summing_cost = calculate_total_cost_with_branching(&cost_analysis_node);
             (name, summing_cost.into())
         })
         .collect();
+
     Ok(costs
         .into_iter()
         .map(|(name, cost)| (name, (cost, trait_count.clone())))
@@ -295,26 +298,39 @@ pub fn static_cost_from_ast(
 pub(crate) fn static_cost_tree_from_ast(
     ast: &crate::vm::ast::ContractAST,
     clarity_version: &ClarityVersion,
-) -> Result<HashMap<String, CostAnalysisNode>, String> {
+) -> Result<HashMap<String, (CostAnalysisNode, Option<TraitCount>)>, String> {
     let exprs = &ast.expressions;
     let user_args = UserArgumentsContext::new();
     let costs_map: HashMap<String, Option<StaticCost>> = HashMap::new();
     let mut costs: HashMap<String, Option<CostAnalysisNode>> = HashMap::new();
+    // first pass extracts the function names
     for expr in exprs {
         if let Some(function_name) = extract_function_name(expr) {
             costs.insert(function_name, None);
         }
     }
+    // second pass computes the cost
     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)?;
             costs.insert(function_name, Some(cost_analysis_tree));
         }
     }
-    Ok(costs
+
+    // Build the final map with cost analysis nodes
+    let cost_trees: HashMap<String, CostAnalysisNode> = costs
         .into_iter()
         .filter_map(|(name, cost)| cost.map(|c| (name, c)))
+        .collect();
+
+    // Compute trait_count while creating the root CostAnalysisNode
+    let trait_count = get_trait_count(&cost_trees);
+
+    // Return each node with its trait_count
+    Ok(cost_trees
+        .into_iter()
+        .map(|(name, node)| (name, (node, trait_count.clone())))
         .collect())
 }
 

clarity/src/vm/tests/analysis.rs

@@ -7,8 +7,7 @@ use stacks_common::types::StacksEpochId;
 
 use crate::vm::contexts::OwnedEnvironment;
 use crate::vm::costs::analysis::{
-    build_cost_analysis_tree, get_trait_count, static_cost_from_ast, static_cost_tree_from_ast,
-    UserArgumentsContext,
+    build_cost_analysis_tree, static_cost_from_ast, static_cost_tree_from_ast, UserArgumentsContext,
 };
 use crate::vm::costs::ExecutionCost;
 use crate::vm::types::{PrincipalData, QualifiedContractIdentifier};
@@ -106,7 +105,11 @@ fn test_get_trait_count_direct() {
 
     let costs = static_cost_tree_from_ast(&ast, &ClarityVersion::Clarity3).unwrap();
 
-    let trait_count = get_trait_count(&costs);
+    // Extract trait_count from the result (all entries have the same trait_count)
+    let trait_count = costs
+        .values()
+        .next()
+        .and_then(|(_, trait_count)| trait_count.clone());
 
     let expected = {
         let mut map = HashMap::new();