Add tracing-texray streaming instrumentation for Aiur proving (#401)

Wraps `AiurSystem::prove` in `aiur/execute` and `aiur/witness` spans and
exposes `TracingTexray.init` for Lean callers. Once installed, each
tracked span emits `[texray] <name>: <dur>  ── RAM Δ +X peak Y` on
stderr as it closes (Linux RAM stats), and the full texray graph prints
when `aiur/prove` exits.

Author: samuelburnham

Cargo.lock

@@ -13,7 +13,7 @@ itertools = "0.14.0"
 indexmap = { version = "2", features = ["rayon"] }
 lean-ffi = { git = "https://github.com/argumentcomputer/lean-ffi", rev = "cc98ebf67bf453ac3827cb767f78b13ea674dd6a" }
 mimalloc = { version = "0.1", default-features = false }
-multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "a8a15ea6aa2890f9f60f32a6e0e5e66afc1535ff" }
+multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "9ecab51d553445c0cc7b571af00a76b8a83a6f8c" }
 num-bigint = "0.4.6"
 rayon = "1"
 rustc-hash = "2"
@@ -27,8 +27,9 @@ iroh = { version = "0.97", optional = true }
 iroh-base = { version = "0.97", optional = true }
 n0-error = { version = "0.1", optional = true }
 getrandom = { version = "0.3", optional = true }
-tracing = { version = "0.1", optional = true }
-tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true }
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+tracing-texray = { git = "https://github.com/argumentcomputer/tracing-texray", rev = "8ce04e3422cd48e68ef47fab95dba7d06b8c368c" }
 bincode = { version = "2.0.1", optional = true }
 serde = { version = "1.0.219", features = ["derive"], optional = true }
 
@@ -41,7 +42,7 @@ quickcheck_macros = "1.0.0"
 default = []
 parallel = ["multi-stark/parallel"]
 test-ffi = []
-net = ["bytes", "tokio", "iroh", "iroh-base", "n0-error", "getrandom", "tracing", "tracing-subscriber", "bincode", "serde" ]
+net = ["bytes", "tokio", "iroh", "iroh-base", "n0-error", "getrandom", "bincode", "serde" ]
 
 [profile.dev]
 panic = "abort"

Cargo.toml

@@ -15,3 +15,4 @@ public import Ix.Claim
 public import Ix.Commit
 public import Ix.Benchmark.Bench
 public import Ix.Aiur
+public import Ix.TracingTexray

Ix.lean

@@ -0,0 +1,44 @@
+/-
+  tracing-texray + streaming Lean API.
+
+  The Rust prover emits `tracing` spans for each major Aiur proving stage
+  (`aiur/execute`, `aiur/witness`, `aiur/prove`) and each STARK proving
+  stage (`stark/...`). After `init` installs the subscriber:
+    * a one-line `[texray] <span>: <dur>  ── RAM Δ +X peak Y` is streamed
+      to stderr as each tracked span closes (when `streaming = true`);
+    * the full texray graph (with RAM block, Linux-only) prints when
+      `aiur/prove` exits.
+-/
+
+module
+
+public section
+
+namespace TracingTexray
+
+/-- Configuration for the tracing-texray subscriber. -/
+structure Settings where
+  /-- Comma-separated span-name prefixes to keep
+      (default `"aiur/,stark/"`; empty string disables filtering). -/
+  namePrefixes : String := "aiur/,stark/"
+  /-- Sample VmRSS/VmHWM per span. Linux-only; zeros elsewhere. -/
+  trackRam : Bool := true
+  /-- Emit per-span `[texray] <name>: <dur>  ── RAM Δ +X peak Y` lines
+      on stderr as each span closes. The texray graph prints either way
+      when the examined root span exits. -/
+  streaming : Bool := true
+
+@[extern "rs_texray_init"]
+private opaque initWith
+  (namePrefixes : @& String)
+  (trackRam : Bool)
+  (streaming : Bool) : IO Unit
+
+/-- Install the global tracing-texray subscriber with the given settings.
+    Idempotent: subsequent calls are no-ops once installed. -/
+def init (s : Settings := {}) : IO Unit :=
+  initWith s.namePrefixes s.trackRam s.streaming
+
+end TracingTexray
+
+end

Ix/TracingTexray.lean

@@ -109,21 +109,28 @@ impl AiurSystem {
     AiurSystem { system, key, toplevel }
   }
 
+  #[tracing::instrument(level = "info", skip_all, name = "aiur/prove")]
   pub fn prove(
     &self,
     fri_parameters: FriParameters,
     fun_idx: FunIdx,
     input: &[G],
     io_buffer: &mut IOBuffer,
   ) -> (Vec<G>, Proof) {
+    tracing_texray::examine_current();
+
+    // Execute the Aiur bytecode.
+    let _g = tracing::info_span!("aiur/execute").entered();
     // Execute the Aiur bytecode. The prover assumes inputs are valid; any
     // execution error here is a programmer bug, so we unwrap.
     let (query_record, output) = self
       .toplevel
       .execute(fun_idx, input.to_vec(), io_buffer)
       .expect("Aiur execution failed during prove");
+    drop(_g);
 
     // Build the `SystemWitness`
+    let _g = tracing::info_span!("aiur/witness").entered();
     let functions =
       (0..self.toplevel.functions.len()).into_par_iter().filter_map(|idx| {
         if self.toplevel.functions[idx].constrained {
@@ -155,6 +162,7 @@ impl AiurSystem {
     drop(query_record); // Early drop to free memory.
     let (traces, lookups) = witness_data.into_iter().unzip();
     let witness = SystemWitness { traces, lookups };
+    drop(_g);
 
     // Construct the claim.
     let mut claim = vec![function_channel(), G::from_usize(fun_idx)];

src/aiur/synthesis.rs

@@ -14,6 +14,7 @@ pub mod byte_array;
 pub mod iroh;
 pub mod keccak;
 pub mod lean_env;
+pub mod texray;
 pub mod unsigned;
 
 // Modular FFI structure

src/ffi.rs

@@ -0,0 +1,58 @@
+//! FFI shim for installing the `tracing-texray` subscriber from Lean.
+//!
+//! With `streaming = true`, each tracked span (e.g. `aiur/execute`,
+//! `aiur/witness`, `aiur/prove`) emits a one-line `[texray] <name>: <dur>
+//!   ── RAM Δ +X peak Y` to stderr as it closes, and the full texray
+//! graph prints when the examined root span exits.
+
+use lean_ffi::object::{LeanBorrowed, LeanIOResult, LeanOwned, LeanString};
+use tracing_subscriber::{
+  Layer, Registry, filter::FilterFn, layer::SubscriberExt,
+  util::SubscriberInitExt,
+};
+
+/// Install the `tracing-texray` subscriber as the global default.
+///
+/// Idempotent: subsequent calls are no-ops once the global subscriber has
+/// been set.
+///
+/// Parameters:
+/// - `name_prefixes`: comma-separated list of span-name prefixes to render
+///   (e.g. `"aiur/,stark/"`). Empty string disables filtering and renders
+///   everything, including upstream library spans.
+/// - `track_ram`: sample VmRSS/VmHWM per span. Linux-only; zeros elsewhere.
+/// - `streaming`: emit per-span `[texray] <name>: <dur>  ── RAM Δ +X peak Y`
+///   lines on stderr as each span closes. The texray graph prints either
+///   way when the examined root span exits.
+#[unsafe(no_mangle)]
+extern "C" fn rs_texray_init(
+  name_prefixes: LeanString<LeanBorrowed<'_>>,
+  track_ram: bool,
+  streaming: bool,
+) -> LeanIOResult<LeanOwned> {
+  let prefixes: Vec<String> = name_prefixes
+    .to_string()
+    .split(',')
+    .map(|s| s.trim().to_string())
+    .filter(|s| !s.is_empty())
+    .collect();
+
+  let mut layer = tracing_texray::TeXRayLayer::new();
+  if track_ram {
+    layer = layer.track_ram();
+  }
+  if streaming {
+    layer = layer.streaming();
+  }
+
+  let filter = FilterFn::new(move |metadata| {
+    if !metadata.is_span() || prefixes.is_empty() {
+      return true;
+    }
+    let name = metadata.name();
+    prefixes.iter().any(|p| name.starts_with(p.as_str()))
+  });
+
+  let _ = Registry::default().with(layer.with_filter(filter)).try_init();
+  LeanIOResult::ok(LeanOwned::box_usize(0))
+}