use Environment for static_cost

brady-stacks · brady-stacks · commit fc009feddcf3 · 2025-11-07T11:42:35.000-08:00
diff --git a/clarity/src/vm/costs/analysis.rs b/clarity/src/vm/costs/analysis.rs
@@ -6,6 +6,7 @@ use clarity_types::types::{CharType, SequenceData, TraitIdentifier};
 use stacks_common::types::StacksEpochId;
 
 use crate::vm::ast::build_ast;
+use crate::vm::contexts::Environment;
 use crate::vm::costs::cost_functions::{linear, CostValues};
 use crate::vm::costs::costs_3::Costs3;
 use crate::vm::costs::ExecutionCost;
@@ -221,18 +222,17 @@ fn static_cost_native(
     let summing_cost = calculate_total_cost_with_branching(&cost_analysis_tree);
     Ok(summing_cost.into())
 }
-/// Parse Clarity source code and calculate its static execution cost for the specified function
-pub fn static_cost(
-    source: &str,
+
+pub fn static_cost_from_ast(
+    contract_ast: &crate::vm::ast::ContractAST,
     clarity_version: &ClarityVersion,
 ) -> Result<HashMap<String, StaticCost>, String> {
-    let epoch = StacksEpochId::latest(); // XXX this should be matched with the clarity version
-    let ast = make_ast(source, epoch, clarity_version)?;
+    let exprs = &contract_ast.expressions;
 
-    if ast.expressions.is_empty() {
-        return Err("No expressions found".to_string());
+    if exprs.is_empty() {
+        return Err("No expressions found in contract AST".to_string());
     }
-    let exprs = &ast.expressions;
+
     let user_args = UserArgumentsContext::new();
     let mut costs: HashMap<String, Option<StaticCost>> = HashMap::new();
 
@@ -260,6 +260,40 @@ pub fn static_cost(
         .collect())
 }
 
+/// Calculate static execution cost for functions using Environment context
+/// This replaces the old source-string based approach with Environment integration
+pub fn static_cost(
+    env: &mut Environment,
+    contract_identifier: &QualifiedContractIdentifier,
+) -> Result<HashMap<String, StaticCost>, String> {
+    // Get the contract source from the environment's database
+    let contract_source = env
+        .global_context
+        .database
+        .get_contract_src(contract_identifier)
+        .ok_or_else(|| "Contract source not found in database".to_string())?;
+
+    // Get the contract's clarity version from the environment
+    let contract = env
+        .global_context
+        .database
+        .get_contract(contract_identifier)
+        .map_err(|e| format!("Failed to get contract: {:?}", e))?;
+
+    let clarity_version = contract.contract_context.get_clarity_version();
+
+    let epoch = env.global_context.epoch_id;
+    let ast = make_ast(&contract_source, epoch, clarity_version)?;
+
+    static_cost_from_ast(&ast, clarity_version)
+}
+
+// pub fn static_cost_tree(
+//     source: &str,
+//     clarity_version: &ClarityVersion,
+// ) -> Result<HashMap<String, CostAnalysisNode>, String> {
+// }
+
 /// Extract function name from a symbolic expression
 fn extract_function_name(expr: &SymbolicExpression) -> Option<String> {
     if let Some(list) = expr.match_list() {
@@ -469,7 +503,6 @@ 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, clarity_version)?;
         (expr_node, cost)
@@ -499,6 +532,7 @@ fn calculate_function_cost(
             Ok(cost.clone())
         }
         Some(None) => {
+            // Should be impossible but alas..
             // Function exists but cost not yet computed - this indicates a circular dependency
             // For now, return zero cost to avoid infinite recursion
             println!(
@@ -809,6 +843,15 @@ mod tests {
         static_cost_native(source, &cost_map, clarity_version)
     }
 
+    fn static_cost_test(
+        source: &str,
+        clarity_version: &ClarityVersion,
+    ) -> Result<HashMap<String, StaticCost>, String> {
+        let epoch = StacksEpochId::latest();
+        let ast = make_ast(source, epoch, clarity_version)?;
+        static_cost_from_ast(&ast, clarity_version)
+    }
+
     fn build_test_ast(src: &str) -> crate::vm::ast::ContractAST {
         let contract_identifier = QualifiedContractIdentifier::transient();
         let mut cost_tracker = ();
@@ -1051,4 +1094,41 @@ mod tests {
             assert_eq!(arg_type.as_str(), "uint");
         }
     }
+
+    #[test]
+    fn test_static_cost_simple_addition() {
+        let source = "(define-public (add (a uint) (b uint)) (+ a b))";
+        let ast_cost = static_cost_test(source, &ClarityVersion::Clarity3).unwrap();
+
+        // Should have one function
+        assert_eq!(ast_cost.len(), 1);
+        assert!(ast_cost.contains_key("add"));
+
+        // Check that the cost is reasonable (non-zero for addition)
+        let add_cost = ast_cost.get("add").unwrap();
+        assert!(add_cost.min.runtime > 0);
+        assert!(add_cost.max.runtime > 0);
+    }
+
+    #[test]
+    fn test_static_cost_multiple_functions() {
+        let source = r#"
+            (define-public (func1 (x uint)) (+ x 1))
+            (define-private (func2 (y uint)) (* y 2))
+        "#;
+        let ast_cost = static_cost_test(source, &ClarityVersion::Clarity3).unwrap();
+
+        // Should have 2 functions
+        assert_eq!(ast_cost.len(), 2);
+
+        // Check that both functions are present
+        assert!(ast_cost.contains_key("func1"));
+        assert!(ast_cost.contains_key("func2"));
+
+        // Check that costs are reasonable
+        let func1_cost = ast_cost.get("func1").unwrap();
+        let func2_cost = ast_cost.get("func2").unwrap();
+        assert!(func1_cost.min.runtime > 0);
+        assert!(func2_cost.min.runtime > 0);
+    }
 }
diff --git a/clarity/src/vm/tests/analysis.rs b/clarity/src/vm/tests/analysis.rs
@@ -1,26 +1,12 @@
-// Copyright (C) 2013-2020 Blockstack PBC, a public benefit corporation
-// Copyright (C) 2020 Stacks Open Internet Foundation
-//
-// This program is free software: you can redistribute it and/or modify
-// it under the terms of the GNU General Public License as published by
-// the Free Software Foundation, either version 3 of the License, or
-// (at your option) any later version.
-//
-// This program is distributed in the hope that it will be useful,
-// but WITHOUT ANY WARRANTY; without even the implied warranty of
-// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-// GNU General Public License for more details.
-//
-// You should have received a copy of the GNU General Public License
-// along with this program.  If not, see <http://www.gnu.org/licenses/>.
-
 use std::collections::HashMap;
 
 use rstest::rstest;
 use stacks_common::types::StacksEpochId;
 
 use crate::vm::contexts::OwnedEnvironment;
-use crate::vm::costs::analysis::{build_cost_analysis_tree, static_cost, UserArgumentsContext};
+use crate::vm::costs::analysis::{
+    build_cost_analysis_tree, static_cost_from_ast, UserArgumentsContext,
+};
 use crate::vm::costs::ExecutionCost;
 use crate::vm::tests::{tl_env_factory, TopLevelMemoryEnvironmentGenerator};
 use crate::vm::types::{PrincipalData, QualifiedContractIdentifier};
@@ -38,17 +24,23 @@ fn test_simple_trait_implementation_costs(
     #[case] epoch: StacksEpochId,
     mut tl_env_factory: TopLevelMemoryEnvironmentGenerator,
 ) {
-    // Simple trait implementation - very brief function that basically does nothing
     let simple_impl = r#"(impl-trait .mytrait.mytrait)
         (define-public (somefunc (a uint) (b uint))
           (ok (+ a b))
         )"#;
 
-    // Set up environment with cost tracking - use regular environment but try to get actual costs
     let mut owned_env = tl_env_factory.get_env(epoch);
 
-    // Get static cost analysis
-    let static_cost = static_cost(simple_impl, &version).unwrap();
+    let epoch = StacksEpochId::Epoch21;
+    let ast = crate::vm::ast::build_ast(
+        &QualifiedContractIdentifier::transient(),
+        simple_impl,
+        &mut (),
+        version,
+        epoch,
+    )
+    .unwrap();
+    let static_cost = static_cost_from_ast(&ast, &version).unwrap();
     // Deploy and execute the contract to get dynamic costs
     let contract_id = QualifiedContractIdentifier::local("simple-impl").unwrap();
     owned_env
@@ -71,71 +63,13 @@ fn test_simple_trait_implementation_costs(
     assert!(dynamic_cost.runtime <= cost.max.runtime);
 }
 
-/// Helper function to execute a contract function and return the execution cost
-fn execute_contract_function_and_get_cost(
-    env: &mut OwnedEnvironment,
-    contract_id: &QualifiedContractIdentifier,
-    function_name: &str,
-    args: &[u64],
-    version: ClarityVersion,
-) -> ExecutionCost {
-    // Start with a fresh cost tracker
-    let initial_cost = env.get_cost_total();
-
-    // Create a dummy sender
-    let sender = PrincipalData::parse_qualified_contract_principal(
-        "ST1PQHQKV0RJXZFY1DGX8MNSNYVE3VGZJSRTPGZGM.sender",
-    )
-    .unwrap();
-
-    // Build function call string
-    let arg_str = args
-        .iter()
-        .map(|a| format!("u{}", a))
-        .collect::<Vec<_>>()
-        .join(" ");
-    let function_call = format!("({} {})", function_name, arg_str);
-
-    // Parse the function call into a symbolic expression
-    let ast = crate::vm::ast::parse(
-        &QualifiedContractIdentifier::transient(),
-        &function_call,
-        version,
-        StacksEpochId::Epoch21,
-    )
-    .expect("Failed to parse function call");
-
-    if !ast.is_empty() {
-        let _result = env.execute_transaction(
-            sender,
-            None,
-            contract_id.clone(),
-            &function_call,
-            &ast[0..1],
-        );
-    }
-
-    // Get the cost after execution
-    let final_cost = env.get_cost_total();
-
-    // Return the difference
-    ExecutionCost {
-        write_length: final_cost.write_length - initial_cost.write_length,
-        write_count: final_cost.write_count - initial_cost.write_count,
-        read_length: final_cost.read_length - initial_cost.read_length,
-        read_count: final_cost.read_count - initial_cost.read_count,
-        runtime: final_cost.runtime - initial_cost.runtime,
-    }
-}
-
 #[rstest]
 #[case::clarity2(ClarityVersion::Clarity2, StacksEpochId::Epoch21)]
 fn test_complex_trait_implementation_costs(
     #[case] version: ClarityVersion,
     #[case] epoch: StacksEpochId,
     mut tl_env_factory: TopLevelMemoryEnvironmentGenerator,
 ) {
-    // Complex trait implementation with expensive operations but no external calls
     let complex_impl = r#"(define-public (somefunc (a uint) (b uint))
     (begin
         ;; do something expensive
@@ -152,7 +86,16 @@ fn test_complex_trait_implementation_costs(
 
     let mut owned_env = tl_env_factory.get_env(epoch);
 
-    let static_cost_result = static_cost(complex_impl, &version);
+    let epoch = StacksEpochId::Epoch21;
+    let ast = crate::vm::ast::build_ast(
+        &QualifiedContractIdentifier::transient(),
+        complex_impl,
+        &mut (),
+        version,
+        epoch,
+    )
+    .unwrap();
+    let static_cost_result = static_cost_from_ast(&ast, &version);
     match static_cost_result {
         Ok(static_cost) => {
             let contract_id = QualifiedContractIdentifier::local("complex-impl").unwrap();
@@ -228,11 +171,74 @@ fn test_dependent_function_calls() {
 )"#;
 
     let contract_id = QualifiedContractIdentifier::transient();
-    let function_map = static_cost(src, &ClarityVersion::Clarity3).unwrap();
+    let epoch = StacksEpochId::Epoch32;
+    let ast = crate::vm::ast::build_ast(
+        &QualifiedContractIdentifier::transient(),
+        src,
+        &mut (),
+        ClarityVersion::Clarity3,
+        epoch,
+    )
+    .unwrap();
+    let function_map = static_cost_from_ast(&ast, &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);
 }
+
+/// Helper function to execute a contract function and return the execution cost
+fn execute_contract_function_and_get_cost(
+    env: &mut OwnedEnvironment,
+    contract_id: &QualifiedContractIdentifier,
+    function_name: &str,
+    args: &[u64],
+    version: ClarityVersion,
+) -> ExecutionCost {
+    // Start with a fresh cost tracker
+    let initial_cost = env.get_cost_total();
+
+    // Create a dummy sender
+    let sender = PrincipalData::parse_qualified_contract_principal(
+        "ST1PQHQKV0RJXZFY1DGX8MNSNYVE3VGZJSRTPGZGM.sender",
+    )
+    .unwrap();
+
+    let arg_str = args
+        .iter()
+        .map(|a| format!("u{}", a))
+        .collect::<Vec<_>>()
+        .join(" ");
+    let function_call = format!("({} {})", function_name, arg_str);
+
+    let ast = crate::vm::ast::parse(
+        &QualifiedContractIdentifier::transient(),
+        &function_call,
+        version,
+        StacksEpochId::Epoch21,
+    )
+    .expect("Failed to parse function call");
+
+    if !ast.is_empty() {
+        let _result = env.execute_transaction(
+            sender,
+            None,
+            contract_id.clone(),
+            &function_call,
+            &ast[0..1],
+        );
+    }
+
+    // Get the cost after execution
+    let final_cost = env.get_cost_total();
+
+    ExecutionCost {
+        write_length: final_cost.write_length - initial_cost.write_length,
+        write_count: final_cost.write_count - initial_cost.write_count,
+        read_length: final_cost.read_length - initial_cost.read_length,
+        read_count: final_cost.read_count - initial_cost.read_count,
+        runtime: final_cost.runtime - initial_cost.runtime,
+    }
+}
