Implement Libfunc Counter

Cargo.toml

@@ -61,6 +61,7 @@ with-cheatcode = []
 with-debug-utils = []
 with-mem-tracing = []
 with-libfunc-profiling = []
+with-libfunc-counter = []
 with-segfault-catcher = []
 with-trace-dump = ["dep:sierra-emu"]
 

src/bin/cairo-native-run.rs

@@ -6,6 +6,8 @@ use cairo_lang_runner::short_string::as_cairo_short_string;
 #[cfg(feature = "with-libfunc-profiling")]
 use cairo_lang_sierra::ids::ConcreteLibfuncId;
 use cairo_lang_sierra_to_casm::metadata::MetadataComputationConfig;
+#[cfg(feature = "with-libfunc-counter")]
+use cairo_native::metadata::libfunc_counter::libfunc_counter_runtime::CountersArrayGuard;
 #[cfg(feature = "with-libfunc-profiling")]
 use cairo_native::metadata::profiler::LibfuncProfileData;
 use cairo_native::{
@@ -57,6 +59,11 @@ struct Args {
     /// The output path for the libfunc profilling results
     profiler_output: Option<PathBuf>,
 
+    #[cfg(feature = "with-libfunc-counter")]
+    #[arg(long)]
+    /// The output path for the execution trace
+    libfunc_counter_output: Option<PathBuf>,
+
     #[cfg(feature = "with-trace-dump")]
     #[arg(long)]
     /// The output path for the execution trace
@@ -106,6 +113,9 @@ fn main() -> anyhow::Result<()> {
         .compile(&sierra_program, false, Some(Default::default()), None)
         .unwrap();
 
+    #[cfg(feature = "with-libfunc-counter")]
+    let libfuncs_amount = sierra_program.clone().libfunc_declarations.len();
+
     let native_executor: Box<dyn Fn(_, _, _, &mut StubSyscallHandler) -> _> = match args.run_mode {
         RunMode::Aot => {
             let executor =
@@ -120,13 +130,48 @@ fn main() -> anyhow::Result<()> {
                 }
             }
 
+            #[cfg(feature = "with-libfunc-counter")]
+            {
+                use cairo_native::metadata::libfunc_counter::LibfuncCounterBinding;
+                if let Some(counter_id) =
+                    executor.find_symbol_ptr(LibfuncCounterBinding::CounterId.symbol())
+                {
+                    let counter_id = counter_id.cast::<u64>();
+                    unsafe { *counter_id = 0 };
+                }
+            }
+
             Box::new(move |function_id, args, gas, syscall_handler| {
-                executor.invoke_dynamic_with_syscall_handler(
+                #[cfg(feature = "with-libfunc-counter")]
+                let array_counter_guard = CountersArrayGuard::init(libfuncs_amount);
+
+                let result = executor.invoke_dynamic_with_syscall_handler(
                     function_id,
                     args,
                     gas,
                     syscall_handler,
-                )
+                );
+
+                #[cfg(feature = "with-libfunc-counter")]
+                {
+                    use cairo_native::metadata::libfunc_counter::libfunc_counter_runtime;
+                    use cairo_native::metadata::libfunc_counter::LibfuncCounterBinding;
+
+                    let counter_id_ptr = executor
+                        .find_symbol_ptr(LibfuncCounterBinding::CounterId.symbol())
+                        .expect("");
+
+                    unsafe {
+                        libfunc_counter_runtime::store_and_free_counters_array(
+                            counter_id_ptr as *mut u64,
+                            libfuncs_amount,
+                        );
+                    }
+
+                    drop(array_counter_guard);
+                };
+
+                result
             })
         }
         RunMode::Jit => {
@@ -136,6 +181,7 @@ fn main() -> anyhow::Result<()> {
             #[cfg(feature = "with-trace-dump")]
             {
                 use cairo_native::metadata::trace_dump::TraceBinding;
+
                 if let Some(trace_id) = executor.find_symbol_ptr(TraceBinding::TraceId.symbol()) {
                     let trace_id = trace_id.cast::<u64>();
                     unsafe { *trace_id = 0 };
@@ -154,13 +200,49 @@ fn main() -> anyhow::Result<()> {
                 }
             }
 
+            #[cfg(feature = "with-libfunc-counter")]
+            {
+                use cairo_native::metadata::libfunc_counter::LibfuncCounterBinding;
+
+                if let Some(counter_id) =
+                    executor.find_symbol_ptr(LibfuncCounterBinding::CounterId.symbol())
+                {
+                    let counter_id = counter_id.cast::<u64>();
+                    unsafe { *counter_id = 0 };
+                }
+            }
+
             Box::new(move |function_id, args, gas, syscall_handler| {
-                executor.invoke_dynamic_with_syscall_handler(
+                #[cfg(feature = "with-libfunc-counter")]
+                let array_counter_guard = CountersArrayGuard::init(libfuncs_amount);
+
+                let result = executor.invoke_dynamic_with_syscall_handler(
                     function_id,
                     args,
                     gas,
                     syscall_handler,
-                )
+                );
+
+                #[cfg(feature = "with-libfunc-counter")]
+                {
+                    use cairo_native::metadata::libfunc_counter::libfunc_counter_runtime;
+                    use cairo_native::metadata::libfunc_counter::LibfuncCounterBinding;
+
+                    let counter_id_ptr = executor
+                        .find_symbol_ptr(LibfuncCounterBinding::CounterId.symbol())
+                        .expect("");
+
+                    unsafe {
+                        libfunc_counter_runtime::store_and_free_counters_array(
+                            counter_id_ptr as *mut u64,
+                            libfuncs_amount,
+                        );
+                    }
+
+                    drop(array_counter_guard);
+                };
+
+                result
             })
         }
     };
@@ -281,6 +363,35 @@ fn main() -> anyhow::Result<()> {
         }
     }
 
+    #[cfg(feature = "with-libfunc-counter")]
+    if let Some(libfunc_counter_output) = args.libfunc_counter_output {
+        use std::collections::HashMap;
+
+        let counters =
+            cairo_native::metadata::libfunc_counter::libfunc_counter_runtime::LIBFUNC_COUNTER
+                .lock()
+                .unwrap();
+        assert_eq!(counters.len(), 1);
+
+        let libfunc_counter = counters.values().next().unwrap();
+
+        let libfunc_counts = libfunc_counter
+            .iter()
+            .enumerate()
+            .map(|(i, count)| {
+                let libfunc = &sierra_program.libfunc_declarations[i];
+                let debug_name = libfunc.id.debug_name.clone().unwrap().to_string();
+
+                (debug_name, *count)
+            })
+            .collect::<HashMap<String, u32>>();
+        serde_json::to_writer_pretty(
+            std::fs::File::create(libfunc_counter_output).unwrap(),
+            &libfunc_counts,
+        )
+        .unwrap();
+    }
+
     #[cfg(feature = "with-trace-dump")]
     if let Some(trace_output) = args.trace_output {
         let traces = cairo_native::metadata::trace_dump::trace_dump_runtime::TRACE_DUMP

src/compiler.rs

@@ -60,6 +60,8 @@ use crate::{
     utils::{generate_function_name, walk_ir::walk_mlir_block, BlockExt},
 };
 use bumpalo::Bump;
+#[cfg(feature = "with-libfunc-counter")]
+use cairo_lang_sierra::ids::ConcreteLibfuncId;
 use cairo_lang_sierra::{
     edit_state,
     extensions::{
@@ -151,6 +153,14 @@ pub fn compile(
     let n_libfuncs = program.libfunc_declarations.len() + 1;
     let sierra_stmt_start_offset = num_types + n_libfuncs + 1;
 
+    #[cfg(feature = "with-libfunc-counter")]
+    let libfunc_indexes = program
+        .libfunc_declarations
+        .iter()
+        .enumerate()
+        .map(|(idx, libf)| (libf.id.clone(), idx))
+        .collect::<HashMap<ConcreteLibfuncId, usize>>();
+
     for function in &program.funcs {
         tracing::info!("Compiling function `{}`.", function.id);
         compile_func(
@@ -159,6 +169,8 @@ pub fn compile(
             registry,
             function,
             &program.statements,
+            #[cfg(feature = "with-libfunc-counter")]
+            &libfunc_indexes,
             metadata,
             di_compile_unit_id,
             sierra_stmt_start_offset,
@@ -186,6 +198,7 @@ fn compile_func(
     registry: &ProgramRegistry<CoreType, CoreLibfunc>,
     function: &Function,
     statements: &[Statement],
+    #[cfg(feature = "with-libfunc-counter")] libfunc_indexes: &HashMap<ConcreteLibfuncId, usize>,
     metadata: &mut MetadataStorage,
     di_compile_unit_id: Attribute,
     sierra_stmt_start_offset: usize,
@@ -640,6 +653,23 @@ fn compile_func(
                         },
                     };
 
+                    #[cfg(feature = "with-libfunc-counter")]
+                    {
+                        // Can't fail. If we had a key from the invocation statement that is not included in this hashmap,
+                        // that would mean that there's an error in the sierra program since we got to invoke a libfunc that has
+                        // not been declared.
+                        let libfunc_idx = libfunc_indexes.get(&invocation.libfunc_id).unwrap();
+
+                        crate::metadata::libfunc_counter::libfunc_counter_runtime::count_libfunc(
+                            context,
+                            module,
+                            block,
+                            location,
+                            metadata,
+                            *libfunc_idx,
+                        )?;
+                    }
+
                     libfunc.build(
                         context,
                         registry,

src/executor/aot.rs

@@ -63,6 +63,9 @@ impl AotNativeExecutor {
         #[cfg(feature = "with-libfunc-profiling")]
         crate::metadata::profiler::setup_runtime(|name| executor.find_symbol_ptr(name));
 
+        #[cfg(feature = "with-libfunc-counter")]
+        crate::metadata::libfunc_counter::setup_runtime(|name| executor.find_symbol_ptr(name));
+
         executor
     }
 

src/executor/contract.rs

@@ -337,6 +337,9 @@ impl AotContractExecutor {
         #[cfg(feature = "with-libfunc-profiling")]
         crate::metadata::profiler::setup_runtime(|name| executor.find_symbol_ptr(name));
 
+        #[cfg(feature = "with-libfunc-counter")]
+        crate::metadata::libfunc_counter::setup_runtime(|name| executor.find_symbol_ptr(name));
+
         Ok(Some(executor))
     }
 

src/executor/jit.rs

@@ -74,6 +74,9 @@ impl<'m> JitNativeExecutor<'m> {
         #[cfg(feature = "with-libfunc-profiling")]
         crate::metadata::profiler::setup_runtime(|name| executor.find_symbol_ptr(name));
 
+        #[cfg(feature = "with-libfunc-counter")]
+        crate::metadata::libfunc_counter::setup_runtime(|name| executor.find_symbol_ptr(name));
+
         Ok(executor)
     }
 

src/metadata.rs

@@ -21,6 +21,7 @@ pub mod dup_overrides;
 pub mod enum_snapshot_variants;
 pub mod felt252_dict;
 pub mod gas;
+pub mod libfunc_counter;
 pub mod profiler;
 pub mod realloc_bindings;
 pub mod runtime_bindings;