fix value increment

Author: FrancoGiachetta

src/bin/cairo-native-run.rs

@@ -135,14 +135,41 @@ fn main() -> anyhow::Result<()> {
                     unsafe { *counter_id = 0 };
                 }
             }
+            
+            #[cfg(feature = "with-libfunc-counter")]
+            let libfuncs_amount = sierra_program.libfunc_declarations.len();
 
             Box::new(move |function_id, args, gas, syscall_handler| {
-                executor.invoke_dynamic_with_syscall_handler(
+                let result = executor.invoke_dynamic_with_syscall_handler(
                     function_id,
                     args,
                     gas,
                     syscall_handler,
-                )
+                );
+
+                #[cfg(feature = "with-libfunc-counter")]
+                unsafe {
+                    use cairo_native::metadata::libfunc_counter::{
+                        libfunc_counter_runtime, LibfuncCounterBinding,
+                    };
+
+                    let counter_id_ptr = executor
+                        .find_symbol_ptr(LibfuncCounterBinding::CounterId.symbol())
+                        .unwrap()
+                        .cast::<u64>();
+                    let counters_array_ptr_ptr = executor
+                        .find_symbol_ptr(LibfuncCounterBinding::CounterArray.symbol())
+                        .unwrap()
+                        .cast::<*mut u32>();
+
+                    libfunc_counter_runtime::store_counters_array(
+                        counter_id_ptr,
+                        counters_array_ptr_ptr,
+                        libfuncs_amount,
+                    );
+                }
+
+                result
             })
         }
         RunMode::Jit => {
@@ -152,6 +179,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 };
@@ -173,6 +201,7 @@ 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())
                 {
@@ -181,13 +210,40 @@ fn main() -> anyhow::Result<()> {
                 }
             }
 
+            #[cfg(feature = "with-libfunc-counter")]
+            let libfuncs_amount = sierra_program.libfunc_declarations.len();
+
             Box::new(move |function_id, args, gas, syscall_handler| {
-                executor.invoke_dynamic_with_syscall_handler(
+                let result = executor.invoke_dynamic_with_syscall_handler(
                     function_id,
                     args,
                     gas,
                     syscall_handler,
-                )
+                );
+
+                #[cfg(feature = "with-libfunc-counter")]
+                unsafe {
+                    use cairo_native::metadata::libfunc_counter::{
+                        libfunc_counter_runtime, LibfuncCounterBinding,
+                    };
+
+                    let counter_id_ptr = executor
+                        .find_symbol_ptr(LibfuncCounterBinding::CounterId.symbol())
+                        .unwrap()
+                        .cast::<u64>();
+                    let counters_array_ptr_ptr = executor
+                        .find_symbol_ptr(LibfuncCounterBinding::CounterArray.symbol())
+                        .unwrap()
+                        .cast::<*mut u32>();
+
+                    libfunc_counter_runtime::store_counters_array(
+                        counter_id_ptr,
+                        counters_array_ptr_ptr,
+                        libfuncs_amount,
+                    );
+                }
+
+                result
             })
         }
     };

src/compiler.rs

@@ -1064,10 +1064,6 @@ 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);
@@ -1449,8 +1445,6 @@ 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));
@@ -1483,20 +1477,6 @@ 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,9 +63,6 @@ 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,9 +333,6 @@ 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,9 +74,6 @@ 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/libfunc_counter.rs

@@ -20,13 +20,12 @@
 //! contract. Since a contract can call other contracts, we need a way of restoring the counter after every execution.
 //!
 //! See `cairo-native-run` for an example on how to do it.
-use std::{collections::HashSet, os::raw::c_void, ptr};
+use std::collections::HashSet;
 
 use melior::{
     dialect::{llvm, memref, ods},
     ir::{
         attribute::{FlatSymbolRefAttribute, StringAttribute, TypeAttribute},
-        operation::OperationBuilder,
         r#type::{IntegerType, MemRefType},
         Attribute, Block, BlockLike, Location, Module, Region, Value,
     },
@@ -41,27 +40,15 @@ use crate::{
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
 pub enum LibfuncCounterBinding {
-    StoreArrayCounter,
     CounterId,
-    ArrayCounter,
+    CounterArray,
 }
 
 impl LibfuncCounterBinding {
     pub const fn symbol(self) -> &'static str {
         match self {
-            LibfuncCounterBinding::StoreArrayCounter => "cairo_native__store_array_counter",
             LibfuncCounterBinding::CounterId => "cairo_native__counter_id",
-            LibfuncCounterBinding::ArrayCounter => "cairo_native__array_counter",
-        }
-    }
-
-    const fn function_ptr(self) -> *const () {
-        match self {
-            LibfuncCounterBinding::StoreArrayCounter => {
-                libfunc_counter_runtime::store_array_counter as *const ()
-            }
-            LibfuncCounterBinding::CounterId => ptr::null(),
-            LibfuncCounterBinding::ArrayCounter => ptr::null(),
+            LibfuncCounterBinding::CounterArray => "cairo_native__counter_array",
         }
     }
 }
@@ -155,68 +142,24 @@ impl LibfuncCounterMeta {
         block.append_op_result(memref::load(libfunc_counter_id_ptr, &[], location))
     }
 
-    /// Indexes the array of counters and increments the counter relative
-    /// to the given libfunc index
-    pub fn store_array_counter(
-        &mut self,
-        context: &Context,
-        module: &Module,
-        block: &Block<'_>,
-        location: Location,
-        libfunc_amount: u32,
-    ) -> Result<()> {
-        let counter_id = self.build_counter_id(context, module, block, location)?;
-        let function_ptr = self.build_function(
-            context,
-            module,
-            block,
-            location,
-            LibfuncCounterBinding::StoreArrayCounter,
-        )?;
-        let lifuncs_amount = block.const_int(context, location, libfunc_amount, 32)?;
-        // by this time, the array counter should be initialized
-        let array_counter_ptr_ptr = block.append_op_result(
-            ods::llvm::mlir_addressof(
-                context,
-                llvm::r#type::pointer(context, 0),
-                FlatSymbolRefAttribute::new(context, LibfuncCounterBinding::ArrayCounter.symbol()),
-                location,
-            )
-            .into(),
-        )?;
-        let array_counter_ptr = block.load(
-            context,
-            location,
-            array_counter_ptr_ptr,
-            llvm::r#type::pointer(context, 0),
-        )?;
-
-        block.append_operation(
-            OperationBuilder::new("llvm.call", location)
-                .add_operands(&[function_ptr])
-                .add_operands(&[counter_id, array_counter_ptr, lifuncs_amount])
-                .build()?,
-        );
-
-        Ok(())
-    }
-
     /// Build the array of counters
-    fn get_array_counter<'c, 'a>(
+    fn build_array_counter<'c, 'a>(
         &mut self,
         context: &'c Context,
         module: &Module,
         block: &'a Block<'c>,
         location: Location<'c>,
         libfunc_amount: u32,
     ) -> Result<Value<'c, 'a>> {
-        if self.active_map.insert(LibfuncCounterBinding::ArrayCounter) {
+        if self.active_map.insert(LibfuncCounterBinding::CounterArray) {
+            self.build_counter_id(context, module, block, location)?;
+
             module.body().append_operation(
                 ods::llvm::mlir_global(
                     context,
                     Region::new(),
                     TypeAttribute::new(llvm::r#type::pointer(context, 0)),
-                    StringAttribute::new(context, LibfuncCounterBinding::ArrayCounter.symbol()),
+                    StringAttribute::new(context, LibfuncCounterBinding::CounterArray.symbol()),
                     Attribute::parse(context, "#llvm.linkage<weak>")
                         .ok_or(Error::ParseAttributeError)?,
                     location,
@@ -240,7 +183,7 @@ impl LibfuncCounterMeta {
                     llvm::r#type::pointer(context, 0),
                     FlatSymbolRefAttribute::new(
                         context,
-                        LibfuncCounterBinding::ArrayCounter.symbol(),
+                        LibfuncCounterBinding::CounterArray.symbol(),
                     ),
                     location,
                 )
@@ -264,13 +207,13 @@ impl LibfuncCounterMeta {
             ods::llvm::mlir_addressof(
                 context,
                 llvm::r#type::pointer(context, 0),
-                FlatSymbolRefAttribute::new(context, LibfuncCounterBinding::ArrayCounter.symbol()),
+                FlatSymbolRefAttribute::new(context, LibfuncCounterBinding::CounterArray.symbol()),
                 location,
             )
             .into(),
         )?;
 
-        // // return the pointer to array counter
+        // return the pointer to array counter
         block.load(
             context,
             location,
@@ -289,10 +232,11 @@ impl LibfuncCounterMeta {
         libfuncs_amount: u32,
     ) -> Result<()> {
         let u32_ty = IntegerType::new(context, 32).into();
-        let k1 = block.const_int(context, location, 0, 32)?;
+        let k1 = block.const_int(context, location, 1, 32)?;
 
         let array_counter_ptr =
-            self.get_array_counter(context, module, block, location, libfuncs_amount)?;
+            self.build_array_counter(context, module, block, location, libfuncs_amount)?;
+
         let value_counter_ptr = block.gep(
             context,
             location,
@@ -310,24 +254,13 @@ impl LibfuncCounterMeta {
     }
 }
 
-pub fn setup_runtime(find_symbol_ptr: impl Fn(&str) -> Option<*mut c_void>) {
-    let bindings = &[LibfuncCounterBinding::StoreArrayCounter];
-
-    for binding in bindings {
-        if let Some(global) = find_symbol_ptr(binding.symbol()) {
-            let global = global.cast::<*const ()>();
-            unsafe { *global = binding.function_ptr() };
-        }
-    }
-}
-
 pub mod libfunc_counter_runtime {
+    use core::slice;
     use std::{
         collections::HashMap,
         sync::{LazyLock, Mutex},
     };
 
-    use itertools::Itertools;
     use melior::{
         ir::{Block, Location, Module},
         Context,
@@ -364,16 +297,18 @@ pub mod libfunc_counter_runtime {
         )
     }
 
-    pub unsafe extern "C" fn store_array_counter(
-        counter_id: u64,
-        array_counter: *const u32,
-        libfuncs_amount: u32,
+    pub unsafe fn store_counters_array(
+        counter_id_ptr: *mut u64,
+        array_ptr_ptr: *mut *mut u32,
+        libfuncs_amount: usize,
     ) {
-        let mut libfunc_counter = LIBFUNC_COUNTER.lock().unwrap();
-        let vec = (0..libfuncs_amount)
-            .map(|i| *array_counter.add(i as usize))
-            .collect_vec();
+        let counters_vec = slice::from_raw_parts(*array_ptr_ptr, libfuncs_amount).to_vec();
+
+        LIBFUNC_COUNTER
+            .lock()
+            .unwrap()
+            .insert(*counter_id_ptr, counters_vec);
 
-        libfunc_counter.insert(counter_id, vec);
+        libc::free(*array_ptr_ptr.cast::<*mut libc::c_void>());
     }
 }