proof_stream/
sink.rs

1//! `ProofEventSink` trait + `ProofSink` wrapper.
2//!
3//! `ProofSink` is the handle that stwo-ml stores in a thread-local and calls
4//! into during proving. It wraps an optional boxed sink so that the inactive
5//! (no visualization) path is a single `Option::is_some()` branch with no
6//! heap allocation.
7
8use crate::events::ProofEvent;
9
10// ─── Trait ───────────────────────────────────────────────────────────────────
11
12/// A consumer of proof events.
13///
14/// Implementations must be `Send + Sync + 'static` so they can be stored in a
15/// thread-local and potentially forwarded across threads.
16pub trait ProofEventSink: Send + Sync + 'static {
17    /// Handle one event. Must not block the calling thread.
18    fn emit(&self, event: ProofEvent);
19
20    /// Called at the end of a proof session. Default: no-op.
21    fn flush(&self) {}
22}
23
24// ─── NullSink ────────────────────────────────────────────────────────────────
25
26/// A sink that discards every event. Used as the default when visualization is
27/// disabled. The compiler should inline and eliminate all calls through this.
28pub struct NullSink;
29
30impl ProofEventSink for NullSink {
31    #[inline(always)]
32    fn emit(&self, _: ProofEvent) {}
33}
34
35// ─── ProofSink ───────────────────────────────────────────────────────────────
36
37/// Wrapper around an optional `ProofEventSink`.
38///
39/// When `inner` is `None` (the common, non-visualization path), all hot-path
40/// calls go through a single branch and generate zero allocations.
41pub struct ProofSink {
42    inner: Option<Box<dyn ProofEventSink>>,
43}
44
45impl ProofSink {
46    /// Create an inactive sink (no visualization).
47    pub fn none() -> Self {
48        Self { inner: None }
49    }
50
51    /// Create an active sink wrapping `s`.
52    pub fn new(s: impl ProofEventSink) -> Self {
53        Self {
54            inner: Some(Box::new(s)),
55        }
56    }
57
58    /// Returns `true` if a sink is installed (visualization is active).
59    #[inline]
60    pub fn is_active(&self) -> bool {
61        self.inner.is_some()
62    }
63
64    /// Emit an event. No-op when inactive.
65    #[inline]
66    pub fn emit(&self, event: ProofEvent) {
67        if let Some(s) = &self.inner {
68            s.emit(event);
69        }
70    }
71
72    /// Emit the result of `f()` only when active.
73    ///
74    /// The closure is **never called** when no sink is installed, so constructing
75    /// an event in a hot loop (e.g. sumcheck rounds) costs nothing when disabled.
76    #[inline]
77    pub fn emit_if(&self, f: impl FnOnce() -> ProofEvent) {
78        if self.inner.is_some() {
79            self.emit(f());
80        }
81    }
82
83    /// Flush the inner sink. Called at the end of a proof session.
84    pub fn flush(&self) {
85        if let Some(s) = &self.inner {
86            s.flush();
87        }
88    }
89}
90
91// ─── ChannelSink ─────────────────────────────────────────────────────────────
92
93/// A sink backed by a `crossbeam` bounded channel.
94///
95/// Events are `try_send`'d — if the channel is full the event is silently
96/// dropped rather than blocking the prover.
97pub struct ChannelSink {
98    tx: crossbeam::channel::Sender<ProofEvent>,
99}
100
101impl ChannelSink {
102    /// Create a new `ChannelSink`/receiver pair with `capacity` slots.
103    pub fn new(capacity: usize) -> (Self, crossbeam::channel::Receiver<ProofEvent>) {
104        let (tx, rx) = crossbeam::channel::bounded(capacity);
105        (Self { tx }, rx)
106    }
107}
108
109impl ProofEventSink for ChannelSink {
110    #[inline]
111    fn emit(&self, event: ProofEvent) {
112        // Silently drop events if the consumer is falling behind.
113        let _ = self.tx.try_send(event);
114    }
115}
116
117// ─── CollectingSink (useful in tests) ────────────────────────────────────────
118
119use std::sync::{Arc, Mutex};
120
121/// A sink that collects all emitted events into a `Vec`. Useful in integration
122/// tests that need to assert on the sequence of events.
123#[derive(Clone)]
124pub struct CollectingSink {
125    events: Arc<Mutex<Vec<ProofEvent>>>,
126}
127
128impl CollectingSink {
129    pub fn new() -> Self {
130        Self {
131            events: Arc::new(Mutex::new(Vec::new())),
132        }
133    }
134
135    /// Drain and return all collected events.
136    pub fn drain(&self) -> Vec<ProofEvent> {
137        self.events.lock().unwrap().drain(..).collect()
138    }
139
140    /// Return a snapshot of collected events (clone).
141    pub fn snapshot(&self) -> Vec<ProofEvent> {
142        self.events.lock().unwrap().clone()
143    }
144}
145
146impl Default for CollectingSink {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152impl ProofEventSink for CollectingSink {
153    fn emit(&self, event: ProofEvent) {
154        self.events.lock().unwrap().push(event);
155    }
156}
proof_stream/sink.rs

proof_stream/
sink.rs
