2-pass for dependent function costs

Author: brady-stacks

clarity/src/vm/costs/analysis.rs

@@ -207,7 +207,7 @@ fn make_ast(
 /// somewhat of a passthrough since we don't have to build the whole context we can jsut return the cost of the single expression
 fn static_cost_native(
     source: &str,
-    cost_map: &HashMap<String, StaticCost>,
+    cost_map: &HashMap<String, Option<StaticCost>>,
     clarity_version: &ClarityVersion,
 ) -> Result<StaticCost, String> {
     let epoch = StacksEpochId::latest(); // XXX this should be matched with the clarity version
@@ -234,26 +234,53 @@ pub fn static_cost(
     }
     let exprs = &ast.expressions;
     let user_args = UserArgumentsContext::new();
-    let mut costs = HashMap::new();
+    let mut costs: HashMap<String, Option<StaticCost>> = HashMap::new();
+
+    // First pass registers all function definitions
+    for expr in exprs {
+        if let Some(function_name) = extract_function_name(expr) {
+            costs.insert(function_name, None);
+        }
+    }
+
+    // Second pass computes costs
     for expr in exprs {
-        let (_, cost_analysis_tree) =
-            build_cost_analysis_tree(expr, &user_args, &costs, clarity_version)?;
-
-        let summing_cost = calculate_total_cost_with_branching(&cost_analysis_tree);
-        costs.insert(
-            expr.match_atom()
-                .map(|name| name.to_string())
-                .unwrap_or_default(),
-            summing_cost.into(),
-        );
-    }
-    Ok(costs)
+        if let Some(function_name) = extract_function_name(expr) {
+            let (_, cost_analysis_tree) =
+                build_cost_analysis_tree(expr, &user_args, &costs, clarity_version)?;
+
+            let summing_cost = calculate_total_cost_with_branching(&cost_analysis_tree);
+            costs.insert(function_name, Some(summing_cost.into()));
+        }
+    }
+
+    Ok(costs
+        .into_iter()
+        .filter_map(|(name, cost)| cost.map(|c| (name, c)))
+        .collect())
+}
+
+/// Extract function name from a symbolic expression
+fn extract_function_name(expr: &SymbolicExpression) -> Option<String> {
+    if let Some(list) = expr.match_list() {
+        if let Some(first_atom) = list.first().and_then(|first| first.match_atom()) {
+            if is_function_definition(first_atom.as_str()) {
+                if let Some(signature) = list.get(1).and_then(|sig| sig.match_list()) {
+                    return signature
+                        .first()
+                        .and_then(|name| name.match_atom())
+                        .map(|name| name.to_string());
+                }
+            }
+        }
+    }
+    None
 }
 
 pub fn build_cost_analysis_tree(
     expr: &SymbolicExpression,
     user_args: &UserArgumentsContext,
-    cost_map: &HashMap<String, StaticCost>,
+    cost_map: &HashMap<String, Option<StaticCost>>,
     clarity_version: &ClarityVersion,
 ) -> Result<(Option<String>, CostAnalysisNode), String> {
     match &expr.expr {
@@ -339,7 +366,7 @@ fn parse_atom_expression(
 fn build_function_definition_cost_analysis_tree(
     list: &[SymbolicExpression],
     _user_args: &UserArgumentsContext,
-    cost_map: &HashMap<String, StaticCost>,
+    cost_map: &HashMap<String, Option<StaticCost>>,
     clarity_version: &ClarityVersion,
 ) -> Result<(String, CostAnalysisNode), String> {
     let define_type = list[0]
@@ -417,7 +444,7 @@ fn get_function_name(expr: &SymbolicExpression) -> Result<ClarityName, String> {
 fn build_listlike_cost_analysis_tree(
     exprs: &[SymbolicExpression],
     user_args: &UserArgumentsContext,
-    cost_map: &HashMap<String, StaticCost>,
+    cost_map: &HashMap<String, Option<StaticCost>>,
     clarity_version: &ClarityVersion,
 ) -> Result<CostAnalysisNode, String> {
     let mut children = Vec::new();
@@ -442,8 +469,9 @@ fn build_listlike_cost_analysis_tree(
         (CostExprNode::NativeFunction(native_function), cost)
     } else {
         // If not a native function, treat as user-defined function and look it up
+        println!("in user-defined function");
         let expr_node = CostExprNode::UserFunction(function_name.clone());
-        let cost = calculate_function_cost(function_name.to_string(), cost_map)?;
+        let cost = calculate_function_cost(function_name.to_string(), cost_map, clarity_version)?;
         (expr_node, cost)
     };
 
@@ -459,14 +487,31 @@ fn build_listlike_cost_analysis_tree(
     Ok(CostAnalysisNode::new(expr_node, cost, children))
 }
 
-// this is a bit tricky, we need to ensure the previously defined function is
-// within the cost_map already or we need to find it and compute the cost first
+// Calculate function cost with lazy evaluation support
 fn calculate_function_cost(
     function_name: String,
-    cost_map: &HashMap<String, StaticCost>,
+    cost_map: &HashMap<String, Option<StaticCost>>,
+    _clarity_version: &ClarityVersion,
 ) -> Result<StaticCost, String> {
-    let cost = cost_map.get(&function_name).unwrap_or(&StaticCost::ZERO);
-    Ok(cost.clone())
+    match cost_map.get(&function_name) {
+        Some(Some(cost)) => {
+            // Cost already computed
+            Ok(cost.clone())
+        }
+        Some(None) => {
+            // Function exists but cost not yet computed - this indicates a circular dependency
+            // For now, return zero cost to avoid infinite recursion
+            println!(
+                "Circular dependency detected for function: {}",
+                function_name
+            );
+            Ok(StaticCost::ZERO)
+        }
+        None => {
+            // Function not found
+            Ok(StaticCost::ZERO)
+        }
+    }
 }
 /// This function is no longer needed - we now use NativeFunctions::lookup_by_name_at_version
 /// directly in build_listlike_cost_analysis_tree
@@ -760,7 +805,7 @@ mod tests {
         source: &str,
         clarity_version: &ClarityVersion,
     ) -> Result<StaticCost, String> {
-        let cost_map = HashMap::new();
+        let cost_map: HashMap<String, Option<StaticCost>> = HashMap::new();
         static_cost_native(source, &cost_map, clarity_version)
     }
 

clarity/src/vm/tests/analysis.rs

@@ -181,18 +181,18 @@ fn test_complex_trait_implementation_costs(
 
 #[test]
 fn test_build_cost_analysis_tree_function_definition() {
-    let source = r#"(define-public (somefunc (a uint))
+    let src = r#"(define-public (somefunc (a uint))
   (ok (+ a 1))
 )"#;
 
     let contract_id = QualifiedContractIdentifier::transient();
     let ast = ast::parse(
         &contract_id,
-        source,
+        src,
         ClarityVersion::Clarity3,
         StacksEpochId::Epoch32,
     )
-    .expect("Failed to parse source code");
+    .expect("Failed to parse");
 
     let expr = &ast[0];
     let user_args = UserArgumentsContext::new();
@@ -214,3 +214,25 @@ fn test_build_cost_analysis_tree_function_definition() {
         }
     }
 }
+
+#[test]
+fn test_dependent_function_calls() {
+    let src = r#"(define-public (add-one (a uint))
+  (begin
+    (print "somefunc")
+    (somefunc a)
+  )
+)
+(define-private (somefunc (a uint))
+  (ok (+ a 1))
+)"#;
+
+    let contract_id = QualifiedContractIdentifier::transient();
+    let function_map = static_cost(src, &ClarityVersion::Clarity3).unwrap();
+
+    let add_one_cost = function_map.get("add-one").unwrap();
+    let somefunc_cost = function_map.get("somefunc").unwrap();
+
+    assert!(add_one_cost.min.runtime >= somefunc_cost.min.runtime);
+    assert!(add_one_cost.max.runtime >= somefunc_cost.max.runtime);
+}