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

@@ -57,6 +57,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
@@ -120,6 +125,17 @@ 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(
                     function_id,
@@ -154,6 +170,17 @@ 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(
                     function_id,
@@ -281,6 +308,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,22 @@ fn compile_func(
                         },
                     };
 
+                    #[cfg(feature = "with-libfunc-counter")]
+                    {
+                        // Can't fail since that would we got a libfunc which is not defined in the program
+                        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_indexes.len() as u32,
+                        )?;
+                    }
+
                     libfunc.build(
                         context,
                         registry,
@@ -1033,6 +1062,10 @@ fn compile_func(
             sierra_stmt_start_offset + function.entry_point.0,
             0,
         ),
+        #[cfg(feature = "with-libfunc-counter")]
+        libfunc_indexes,
+        #[cfg(feature = "with-libfunc-counter")]
+        metadata,
     )?;
 
     tracing::debug!("Done generating function {}.", function.id);
@@ -1414,6 +1447,8 @@ fn generate_entry_point_wrapper<'c>(
     arg_types: &[(Type<'c>, Location<'c>)],
     ret_types: &[Type<'c>],
     location: Location<'c>,
+    #[cfg(feature = "with-libfunc-counter")] libfunc_indexes: &HashMap<ConcreteLibfuncId, usize>,
+    #[cfg(feature = "with-libfunc-counter")] metadata: &mut MetadataStorage,
 ) -> Result<(), Error> {
     let region = Region::new();
     let block = region.append_block(Block::new(arg_types));
@@ -1446,6 +1481,20 @@ fn generate_entry_point_wrapper<'c>(
         returns.push(block.extract_value(context, location, result, *ty, i)?);
     }
 
+    #[cfg(feature = "with-libfunc-counter")]
+    {
+        use crate::metadata::libfunc_counter::LibfuncCounterMeta;
+
+        let libfunc_counter = metadata.get_mut::<LibfuncCounterMeta>().unwrap();
+        libfunc_counter.store_array_counter(
+            context,
+            module,
+            &block,
+            location,
+            libfunc_indexes.len() as u32,
+        )?;
+    }
+
     block.append_operation(func::r#return(&returns, location));
 
     module.body().append_operation(func::func(

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

@@ -333,6 +333,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;