add run_with_libfunc_profile + AotWithProgram variant for ContractExecutor

src/executor.rs

@@ -11,6 +11,8 @@ pub use self::{
 };
 #[cfg(feature = "sierra-emu")]
 pub use self::contract_executor::EmuContractInfo;
+#[cfg(feature = "with-libfunc-profiling")]
+pub use self::contract_executor::AotWithProgram;
 use crate::{
     arch::{AbiArgument, ValueWithInfoWrapper},
     error::{panic::ToNativeAssertError, Error},
@@ -48,6 +50,8 @@ mod aot;
 mod contract;
 mod contract_executor;
 mod jit;
+#[cfg(feature = "with-libfunc-profiling")]
+mod libfunc_profile;
 
 #[cfg(target_arch = "aarch64")]
 global_asm!(include_str!("arch/aarch64.s"));

src/executor/contract_executor.rs

@@ -3,21 +3,26 @@
 //!
 //! The `Emu` variant is gated on the `sierra-emu` feature and uses
 //! [`crate::sierra_emu_bridge::SierraEmuSyscallBridge`] to thread a cairo-native syscall
-//! handler through the sierra-emu VM.
+//! handler through the sierra-emu VM. The `AotWithProgram` variant is gated on
+//! `with-libfunc-profiling` and bundles an [`AotContractExecutor`] with the Sierra
+//! `Program` it was built from so [`ContractExecutor::run_with_profile`] can resolve
+//! libfunc IDs after collecting samples.
 
-#[cfg(feature = "sierra-emu")]
+#[cfg(any(feature = "sierra-emu", feature = "with-libfunc-profiling"))]
 use cairo_lang_sierra::program::Program;
 #[cfg(feature = "sierra-emu")]
 use cairo_lang_starknet_classes::compiler_version::VersionId;
 #[cfg(feature = "sierra-emu")]
 use cairo_lang_starknet_classes::contract_class::ContractEntryPoints;
 use starknet_types_core::felt::Felt;
-#[cfg(feature = "sierra-emu")]
+#[cfg(any(feature = "sierra-emu", feature = "with-libfunc-profiling"))]
 use std::sync::Arc;
 
 use crate::error::Result;
 use crate::execution_result::ContractExecutionResult;
 use crate::executor::AotContractExecutor;
+#[cfg(feature = "with-libfunc-profiling")]
+use crate::metadata::profiler::Profile;
 #[cfg(feature = "sierra-emu")]
 use crate::sierra_emu_bridge::SierraEmuSyscallBridge;
 use crate::starknet::StarknetSyscallHandler;
@@ -33,6 +38,8 @@ pub enum ContractExecutor {
     Aot(AotContractExecutor),
     #[cfg(feature = "sierra-emu")]
     Emu(EmuContractInfo),
+    #[cfg(feature = "with-libfunc-profiling")]
+    AotWithProgram(AotWithProgram),
 }
 
 /// Inputs required to construct a `sierra_emu::VirtualMachine` for the `Emu` variant.
@@ -44,6 +51,16 @@ pub struct EmuContractInfo {
     pub sierra_version: VersionId,
 }
 
+/// AOT executor paired with the Sierra program it was built from. Required by
+/// [`ContractExecutor::run_with_profile`] so libfunc samples can be resolved against
+/// the program's declarations.
+#[cfg(feature = "with-libfunc-profiling")]
+#[derive(Debug)]
+pub struct AotWithProgram {
+    pub executor: AotContractExecutor,
+    pub program: Arc<Program>,
+}
+
 impl From<AotContractExecutor> for ContractExecutor {
     fn from(value: AotContractExecutor) -> Self {
         Self::Aot(value)
@@ -57,6 +74,13 @@ impl From<EmuContractInfo> for ContractExecutor {
     }
 }
 
+#[cfg(feature = "with-libfunc-profiling")]
+impl From<AotWithProgram> for ContractExecutor {
+    fn from(value: AotWithProgram) -> Self {
+        Self::AotWithProgram(value)
+    }
+}
+
 impl ContractExecutor {
     /// Run the contract entry point identified by `selector`.
     ///
@@ -99,6 +123,51 @@ impl ContractExecutor {
                     builtin_stats: Default::default(),
                 })
             }
+            #[cfg(feature = "with-libfunc-profiling")]
+            ContractExecutor::AotWithProgram(AotWithProgram { executor, program }) => executor
+                .run_with_libfunc_profile(
+                    program,
+                    selector,
+                    args,
+                    gas,
+                    builtin_costs,
+                    syscall_handler,
+                    // Profile is collected and dropped on this path. Use
+                    // `run_with_profile` to capture it.
+                    |_profile| {},
+                ),
+        }
+    }
+
+    /// Like [`Self::run`] but, for the `AotWithProgram` variant, hands the captured
+    /// libfunc profile to `on_profile` after the call returns successfully. For other
+    /// variants this is identical to `run` and `on_profile` is never invoked.
+    #[cfg(feature = "with-libfunc-profiling")]
+    pub fn run_with_profile<H, F>(
+        &self,
+        selector: Felt,
+        args: &[Felt],
+        gas: u64,
+        builtin_costs: Option<BuiltinCosts>,
+        syscall_handler: H,
+        on_profile: F,
+    ) -> Result<ContractExecutionResult>
+    where
+        H: StarknetSyscallHandler,
+        F: FnOnce(Profile),
+    {
+        match self {
+            ContractExecutor::AotWithProgram(AotWithProgram { executor, program }) => executor
+                .run_with_libfunc_profile(
+                    program,
+                    selector,
+                    args,
+                    gas,
+                    builtin_costs,
+                    syscall_handler,
+                    on_profile,
+                ),
+            _ => self.run(selector, args, gas, builtin_costs, syscall_handler),
         }
     }
 }

src/executor/libfunc_profile.rs

@@ -0,0 +1,103 @@
+//! Profiling-instrumented run wrapper around [`AotContractExecutor::run`].
+//!
+//! Available under the `with-libfunc-profiling` feature.
+
+#![cfg(feature = "with-libfunc-profiling")]
+
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+
+use cairo_lang_sierra::program::Program;
+use starknet_types_core::felt::Felt;
+
+use crate::error::Result;
+use crate::execution_result::ContractExecutionResult;
+use crate::executor::AotContractExecutor;
+use crate::metadata::profiler::{Profile, ProfilerBinding, ProfilerImpl, LIBFUNC_PROFILE};
+use crate::starknet::StarknetSyscallHandler;
+use crate::utils::BuiltinCosts;
+
+impl AotContractExecutor {
+    /// Run the entrypoint with libfunc-level profiling instrumentation.
+    ///
+    /// Wraps [`AotContractExecutor::run`] with the bookkeeping the
+    /// `with-libfunc-profiling` runtime needs:
+    ///
+    /// 1. Allocates a unique trace ID and inserts an empty `ProfilerImpl` slot in
+    ///    [`LIBFUNC_PROFILE`].
+    /// 2. Points the executor's `cairo_native__profiler__profile_id` symbol at the new
+    ///    trace ID, saving the previous value.
+    /// 3. Calls `run`. Per-statement samples accumulate in the slot via the runtime
+    ///    `push_stmt` callback.
+    /// 4. Drains the slot, calls [`ProfilerImpl::get_profile`] with `program`, and hands
+    ///    the resulting [`Profile`] to `on_profile`.
+    /// 5. A [`ProfilerGuard`] restores the previous trace ID — and removes the slot if
+    ///    the success path didn't — on both success and unwind paths.
+    ///
+    /// `program` must be the Sierra program this executor was compiled from; it's used
+    /// by `get_profile` to map runtime libfunc IDs back to declarations.
+    ///
+    /// Profiling is intended to run single-threaded; concurrent calls would race on the
+    /// global `trace_id` symbol.
+    pub fn run_with_libfunc_profile<H, F>(
+        &self,
+        program: &Arc<Program>,
+        selector: Felt,
+        args: &[Felt],
+        gas: u64,
+        builtin_costs: Option<BuiltinCosts>,
+        syscall_handler: H,
+        on_profile: F,
+    ) -> Result<ContractExecutionResult>
+    where
+        H: StarknetSyscallHandler,
+        F: FnOnce(Profile),
+    {
+        static COUNTER: AtomicU64 = AtomicU64::new(0);
+        let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
+
+        LIBFUNC_PROFILE.lock().unwrap().insert(counter, ProfilerImpl::new());
+
+        // The pointer targets a global symbol in the executor's shared library; it lives
+        // for the executor's lifetime. Single-threaded profiling means no concurrent writer.
+        let trace_id_ptr =
+            self.find_symbol_ptr(ProfilerBinding::ProfileId.symbol()).unwrap().cast::<u64>();
+        // SAFETY: see above. Read/write to a non-null, properly-aligned `*mut u64`.
+        let old_trace_id = unsafe { *trace_id_ptr };
+        unsafe {
+            *trace_id_ptr = counter;
+        }
+
+        // Restore on the success path AND on unwind. On success the caller drains the
+        // slot below; the guard's `remove` is then a no-op.
+        let _guard = ProfilerGuard { trace_id_ptr, old_trace_id, counter };
+
+        let result = self.run(selector, args, gas, builtin_costs, syscall_handler);
+
+        let profiler = LIBFUNC_PROFILE.lock().unwrap().remove(&counter).unwrap();
+        on_profile(profiler.get_profile(program));
+
+        result
+    }
+}
+
+/// RAII cleanup for the profiler globals. Restores `*trace_id_ptr` and drops the
+/// `LIBFUNC_PROFILE` slot at `counter` if it's still occupied.
+struct ProfilerGuard {
+    trace_id_ptr: *mut u64,
+    old_trace_id: u64,
+    counter: u64,
+}
+
+impl Drop for ProfilerGuard {
+    fn drop(&mut self) {
+        // SAFETY: same provenance as the construction site; single-threaded use.
+        unsafe {
+            *self.trace_id_ptr = self.old_trace_id;
+        }
+        // Tolerate a poisoned mutex silently — Drop must not panic.
+        if let Ok(mut profile) = LIBFUNC_PROFILE.lock() {
+            profile.remove(&self.counter);
+        }
+    }
+}

src/metadata/profiler.rs

@@ -330,7 +330,7 @@ impl ProfilerMeta {
 /// Represents the entire profile of the execution.
 ///
 /// It maps the libfunc ID to a libfunc profile.
-type Profile = HashMap<ConcreteLibfuncId, LibfuncProfileData>;
+pub type Profile = HashMap<ConcreteLibfuncId, LibfuncProfileData>;
 
 /// Represents the profile data for a particular libfunc.
 #[derive(Default)]