add run_with_libfunc_profile + AotWithProgram pairing

binaries/starknet-native-compile/Cargo.toml

@@ -7,6 +7,13 @@ edition.workspace = true
 license.workspace = true
 repository.workspace = true
 
+[features]
+# Builds an instrumented compiler whose emitted `.so` carries the libfunc-profiling
+# symbols that `cairo_native::AotContractExecutor::run_with_libfunc_profile` looks up.
+# Install with `cargo install ... --features with-libfunc-profiling` to profile native
+# execution; a binary built without it produces `.so`s with no profiler symbol.
+with-libfunc-profiling = ["cairo-native/with-libfunc-profiling"]
+
 [dependencies]
 anyhow.workspace = true
 cairo-lang-sierra.workspace = true

src/executor.rs

@@ -5,6 +5,8 @@
 
 #[cfg(feature = "sierra-emu")]
 pub use self::emu_contract_executor::EmuContractExecutor;
+#[cfg(feature = "with-libfunc-profiling")]
+pub use self::libfunc_profile::AotWithProgram;
 pub use self::{aot::AotNativeExecutor, contract::AotContractExecutor, jit::JitNativeExecutor};
 use crate::{
     arch::{AbiArgument, ValueWithInfoWrapper},
@@ -23,6 +25,11 @@ use crate::{
     values::Value,
 };
 use bumpalo::Bump;
+// Profiling and sierra-emu consumers (e.g. blockifier) only ever hold the Sierra
+// program as a shared handle. Export the `Arc` alias so they can name it without
+// taking a direct `cairo-lang-sierra` dependency.
+#[cfg(any(feature = "sierra-emu", feature = "with-libfunc-profiling"))]
+pub type ArcProgram = std::sync::Arc<cairo_lang_sierra::program::Program>;
 use cairo_lang_sierra::{
     extensions::{
         circuit::CircuitTypeConcrete,
@@ -44,6 +51,8 @@ mod contract;
 #[cfg(feature = "sierra-emu")]
 mod emu_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/libfunc_profile.rs

@@ -0,0 +1,245 @@
+//! Profiling-instrumented run wrapper around [`AotContractExecutor::run`].
+//!
+//! Available under the `with-libfunc-profiling` feature (gated at the `mod`
+//! declaration in `src/executor.rs`).
+
+use std::cell::Cell;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+
+use cairo_lang_sierra::program::Program;
+use starknet_types_core::felt::Felt;
+
+use crate::error::{Error, Result};
+use crate::execution_result::ContractExecutionResult;
+use crate::executor::{AotContractExecutor, ArcProgram};
+use crate::metadata::profiler::{Profile, ProfilerBinding, ProfilerImpl, LIBFUNC_PROFILE};
+use crate::starknet::StarknetSyscallHandler;
+use crate::utils::BuiltinCosts;
+
+/// An AOT executor paired with the Sierra program it was compiled from. The program is
+/// needed to resolve libfunc profiling samples against the program's declarations, so
+/// pairing the two spares callers from keeping their own copy around.
+#[derive(Debug)]
+pub struct AotWithProgram {
+    pub executor: AotContractExecutor,
+    pub program: ArcProgram,
+}
+
+impl AotWithProgram {
+    /// Run the contract entry point identified by `selector`, discarding the profile.
+    ///
+    /// Mirrors [`AotContractExecutor::run`] so the executor types are interchangeable at
+    /// the call site. The run is still routed through the profiling bookkeeping: the
+    /// executor's shared library is compiled with profiling instrumentation, and running
+    /// it without a live profile slot makes every statement log a missing-profiler
+    /// error. Use [`Self::run_with_profile`] to capture the profile instead.
+    pub fn run(
+        &self,
+        selector: Felt,
+        args: &[Felt],
+        gas: u64,
+        builtin_costs: Option<BuiltinCosts>,
+        syscall_handler: impl StarknetSyscallHandler,
+    ) -> Result<ContractExecutionResult> {
+        self.executor.run_with_libfunc_profile(
+            &self.program,
+            selector,
+            args,
+            gas,
+            builtin_costs,
+            syscall_handler,
+            |_profile| {},
+        )
+    }
+
+    /// Like [`Self::run`] but hands the captured libfunc profile -- together with the
+    /// program this executor was paired with -- to `on_profile` after the call returns
+    /// successfully. The program is included so callers don't have to keep their own
+    /// copy around just to resolve libfunc samples.
+    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, ArcProgram),
+    {
+        let program_for_cb = Arc::clone(&self.program);
+        self.executor.run_with_libfunc_profile(
+            &self.program,
+            selector,
+            args,
+            gas,
+            builtin_costs,
+            syscall_handler,
+            move |profile| on_profile(profile, program_for_cb),
+        )
+    }
+}
+
+/// Process-wide lock that serializes *top-level* profiled runs across threads. The
+/// profiler hot-swaps a process-global symbol (`cairo_native__profiler__profile_id`);
+/// a concurrent thread would race on that write and on the [`LIBFUNC_PROFILE`] slot
+/// bookkeeping. It is acquired only by the outermost profiled frame on each thread
+/// (see [`PROFILE_DEPTH`]); nested same-thread calls -- a profiled contract invoking
+/// another contract -- re-enter without re-locking, which would otherwise deadlock
+/// this non-reentrant mutex.
+static PROFILE_LOCK: Mutex<()> = Mutex::new(());
+
+thread_local! {
+    /// Nesting depth of profiled runs on the current thread. Only the outermost frame
+    /// (depth 0) takes [`PROFILE_LOCK`]; deeper frames rely on the lock the outer frame
+    /// already holds. The per-call `old_trace_id` save/restore keeps the global trace-id
+    /// symbol correct across nesting without any additional locking.
+    static PROFILE_DEPTH: Cell<usize> = const { Cell::new(0) };
+}
+
+impl AotContractExecutor {
+    /// Run the entrypoint with libfunc-level profiling instrumentation.
+    ///
+    /// Wraps [`AotContractExecutor::run`] with the bookkeeping the
+    /// `with-libfunc-profiling` runtime needs:
+    ///
+    /// 1. Acquires [`PROFILE_LOCK`] so concurrent profile calls serialize on the
+    ///    global trace-id symbol. The lock is recovered if poisoned.
+    /// 2. Looks up the executor's `cairo_native__profiler__profile_id` symbol. If
+    ///    absent (the .so was compiled without profiling instrumentation) the call
+    ///    returns an error before touching any global state.
+    /// 3. Allocates a unique trace ID and inserts an empty `ProfilerImpl` slot in
+    ///    [`LIBFUNC_PROFILE`]; points the profile-id symbol at the new ID, saving
+    ///    the previous value.
+    /// 4. Calls `run`. Per-statement samples accumulate in the slot via the runtime
+    ///    `push_stmt` callback.
+    /// 5. Drains the slot. On success (and only on success) hands the resulting
+    ///    [`Profile`] to `on_profile`; on failure the callback is not invoked
+    ///    (partial profiles aren't meaningful).
+    /// 6. A [`ProfilerGuard`] restores the previous trace ID and clears the slot on
+    ///    both the 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.
+    #[allow(clippy::too_many_arguments)]
+    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),
+    {
+        // Acquire the cross-thread lock only at the outermost profiled frame on this
+        // thread. A profiled contract that calls another contract re-enters this
+        // function on the same thread; re-locking the non-reentrant `PROFILE_LOCK`
+        // there would self-deadlock, so nested frames inherit the outer frame's lock.
+        // Recover from a poisoned lock -- it only gates access to the global trace-id
+        // symbol, on which we hold no data invariants.
+        let _profile_lock = PROFILE_DEPTH
+            .with(|depth| depth.get() == 0)
+            .then(|| PROFILE_LOCK.lock().unwrap_or_else(|e| e.into_inner()));
+        PROFILE_DEPTH.with(|depth| depth.set(depth.get() + 1));
+        let _depth_guard = ProfileDepthGuard;
+
+        // Look up the profile-id symbol before touching any global state. If the
+        // executor wasn't compiled with libfunc-profiling instrumentation, the
+        // symbol is absent -- return a typed error rather than panicking.
+        let trace_id_ptr = self
+            .find_symbol_ptr(ProfilerBinding::ProfileId.symbol())
+            .ok_or_else(|| {
+                Error::UnexpectedValue(format!(
+                    "AOT executor missing libfunc-profiling symbol `{}`; \
+                     was the program compiled with libfunc-profiling enabled?",
+                    ProfilerBinding::ProfileId.symbol()
+                ))
+            })?
+            .cast::<u64>();
+
+        static COUNTER: AtomicU64 = AtomicU64::new(0);
+        let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
+
+        LIBFUNC_PROFILE
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .insert(counter, ProfilerImpl::new());
+
+        // SAFETY: the pointer targets a memref-global emitted into the executor's
+        // shared library; the executor outlives the call. `PROFILE_LOCK` serializes
+        // us against any other writer, and the JIT/AOT code reads through the same
+        // address. Reads/writes are aligned `u64`s.
+        let old_trace_id = unsafe { *trace_id_ptr };
+        unsafe {
+            *trace_id_ptr = counter;
+        }
+
+        let _guard = ProfilerGuard {
+            trace_id_ptr,
+            old_trace_id,
+            counter,
+        };
+
+        let result = self.run(selector, args, gas, builtin_costs, syscall_handler);
+
+        // Drain the slot. `ProfilerGuard::drop` would also remove it; doing it here
+        // means we hold the lock for the shortest time and can hand the profile to
+        // the callback. Tolerate a poisoned mutex (we'd lose the profile, not state).
+        let drained = LIBFUNC_PROFILE
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .remove(&counter);
+
+        // Only call the user's callback when `run` succeeded -- a partial profile
+        // captured against an aborted execution wouldn't be meaningful.
+        if let (Some(profiler), Ok(_)) = (drained, &result) {
+            on_profile(profiler.get_profile(program));
+        }
+
+        result
+    }
+}
+
+/// RAII cleanup for the profiler globals. Restores `*trace_id_ptr` on success or
+/// unwind. The [`LIBFUNC_PROFILE`] slot at `counter` is normally drained on the
+/// success path; this guard removes it if it's still occupied (panic case).
+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. `PROFILE_LOCK` is held
+        // by the enclosing scope (still in flight while we drop) so no other thread
+        // races us.
+        unsafe {
+            *self.trace_id_ptr = self.old_trace_id;
+        }
+        // Tolerate a poisoned mutex silently -- Drop must not panic. Slot leak on
+        // poison is intentional and matches the behavior of other Drop impls in
+        // this crate; the alternative (panic in Drop) is worse.
+        if let Ok(mut profile) = LIBFUNC_PROFILE.lock() {
+            profile.remove(&self.counter);
+        }
+    }
+}
+
+/// Decrements [`PROFILE_DEPTH`] on every exit path (including unwind) so the
+/// outermost-frame lock bookkeeping stays correct even if the profiled run panics.
+struct ProfileDepthGuard;
+
+impl Drop for ProfileDepthGuard {
+    fn drop(&mut self) {
+        PROFILE_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1)));
+    }
+}

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)]
@@ -362,7 +362,11 @@ impl ProfilerImpl {
 
     // Push a profiler frame
     pub extern "C" fn push_stmt(profile_id: u64, statement_idx: u64, tick_delta: u64) {
-        let mut profiler = LIBFUNC_PROFILE.lock().unwrap();
+        // Invoked directly from compiled/JIT'd Cairo code across the C ABI: a panic
+        // here (e.g. `.unwrap()` on a poisoned mutex) would unwind into machine code
+        // that cannot handle it. Recover from poison instead of panicking, matching
+        // the lock handling in `run_with_libfunc_profile`.
+        let mut profiler = LIBFUNC_PROFILE.lock().unwrap_or_else(|e| e.into_inner());
 
         let Some(profiler) = profiler.get_mut(&profile_id) else {
             eprintln!("Could not find libfunc profiler!");