Skip to main content

tcontract_exec/
workspace.rs

1//! Owning workspace: page-aligned, grow-only, reused per worker.
2//!
3//! A workspace belongs to exactly one owner — a `tcontract_exec::Pool` or a caller
4//! that wants reuse — and is *lent* to the driver for one call at a time. There
5//! is deliberately no process-global arena: storage handed to a team is
6//! returned to the provider that issued it, and a second owner can never see
7//! it.
8//!
9//! Two payload shapes:
10//!
11//! * **worker buffers** — the packed A block, the accumulator tile and the
12//!   scratch vector of one thread. They are allocated by the thread that will
13//!   write them, which is what makes their first touch local to the node that
14//!   keeps using them; a thread-local handle names its slot in the owner, and a
15//!   handle outlives nothing: the storage is the owner's, the handle is only an
16//!   index, and a slot that has been trimmed is simply re-created.
17//! * **team sets** — the shared packed B panel, the team's scatter vectors and
18//!   its barriers. They are taken exclusively for the duration of a call and
19//!   returned on drop, so two concurrent executes can never share one.
20//!
21//! # Exclusivity
22//!
23//! A slot that is *in the owner's map* is touched only while the map lock is
24//! held. A slot that is *checked out* is not in the map at all: the thread that
25//! checked it out is its only owner until it is returned, so a re-entrant call
26//! on that thread finds nothing and allocates fresh call-local buffers instead
27//! of aliasing the ones its outer call is using (waiting for itself would
28//! deadlock). `trim` and the accounting snapshot therefore never observe a slot
29//! another thread is writing.
30//!
31//! `ArenaProvider::traced()` additionally records every growth's thread and
32//! size, which is how the tests prove that a worker's buffers were allocated by
33//! the worker. Tracing is off by default and only covers buffers the owner
34//! keeps; the fresh re-entrant path is not traced.
35
36use std::collections::HashMap;
37use std::sync::atomic::{AtomicU64, Ordering};
38use std::sync::{Barrier, Mutex};
39
40/// The 4096-byte alignment every buffer in the workspace has.
41const PAGE: usize = 4096;
42
43/// Bytes and counts one execute needs. Buffer sizes are bytes because the
44/// provider is element-type erased: the driver knows the element type and
45/// converts its real counts once.
46///
47/// A zero entry never allocates.
48#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
49pub struct WorkspaceReq {
50    /// This thread's packed A block.
51    pub a_bytes: usize,
52    /// This thread's accumulator tile, including any induced-method scratch.
53    pub tile_bytes: usize,
54    /// Elements reserved for this thread's scatter scratch.
55    pub worker_scatter: usize,
56    /// The team's shared packed B panel.
57    pub b_bytes: usize,
58    /// Elements reserved for the team's block-scatter vectors.
59    pub team_scatter: usize,
60    /// Barriers the team needs, or zero when nothing is shared.
61    pub barriers: usize,
62}
63
64impl WorkspaceReq {
65    /// The payload this request needs, in bytes: the buffers, the `i64` scatter
66    /// elements and the barriers. It is the value a [`Workspace`] is sized to and
67    /// the amount [`Workspace::have`] reports.
68    ///
69    /// # Examples
70    ///
71    /// ```
72    /// use tcontract_exec::WorkspaceReq;
73    /// let req = WorkspaceReq { a_bytes: 16, team_scatter: 2, ..Default::default() };
74    /// assert_eq!(req.bytes(), 16 + 2 * core::mem::size_of::<i64>());
75    /// ```
76    pub fn bytes(&self) -> usize {
77        self.a_bytes
78            + self.tile_bytes
79            + self.b_bytes
80            + (self.worker_scatter + self.team_scatter) * core::mem::size_of::<i64>()
81            + self.barriers * core::mem::size_of::<Barrier>()
82    }
83
84    /// Whether `have` covers every component of this request.
85    ///
86    /// Capacity is checked component-wise, never by the byte total: a large A
87    /// block does not stand in for a missing B panel, and the driver would grow
88    /// the arena for the shortage rather than refuse it. [`Workspace::reserve`]
89    /// keeps the same component-wise maxima.
90    pub fn covered_by(&self, have: &WorkspaceReq) -> bool {
91        self.a_bytes <= have.a_bytes
92            && self.tile_bytes <= have.tile_bytes
93            && self.worker_scatter <= have.worker_scatter
94            && self.b_bytes <= have.b_bytes
95            && self.team_scatter <= have.team_scatter
96            && self.barriers <= have.barriers
97    }
98}
99
100/// What an owner currently holds and currently lends out.
101///
102/// Exactness: both values are settled at every checkout, return, team take,
103/// lease drop and `trim`, so `ArenaProvider::stats` is exact whenever no
104/// execution is in flight. While one is, a checked-out slot is counted at the
105/// capacity it had when it was checked out, so `retained_bytes` is an unsettled
106/// estimate that can lag growth or shrinkage until the call returns.
107#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
108pub struct WorkspaceStats {
109    /// Capacity the owner holds: worker A/tile/scatter and team
110    /// panel/scatter/barriers, idle or in use.
111    pub retained_bytes: usize,
112    /// The part of `retained_bytes` lent to a call right now.
113    pub leased_bytes: usize,
114}
115
116/// Page-aligned, grow-only, never zeroed storage.
117///
118/// The caller writes before reading, which is the contract the packing and
119/// kernel code already has; zeroing would only double the traffic.
120#[derive(Default, Debug)]
121struct PageBuf {
122    ptr: *mut u8,
123    cap: usize,
124    /// Owner key this buffer reports growths under, or zero when untraced.
125    traced: u64,
126}
127
128// SAFETY: the buffer owns its allocation and hands out pointers only through
129// `&mut self` or an explicit, contracted call.
130unsafe impl Send for PageBuf {}
131
132impl PageBuf {
133    /// A buffer that reports its growths under `owner`, for tests that prove
134    /// where storage was first touched. Zero means untraced.
135    fn traced(owner: u64) -> Self {
136        Self {
137            traced: owner,
138            ..Self::default()
139        }
140    }
141
142    /// A pointer to at least `bytes` writable, page-aligned bytes, growing the
143    /// allocation if needed and reusing it otherwise. `ensure(0)` returns a
144    /// dangling but well-aligned pointer and allocates nothing.
145    fn ensure(&mut self, bytes: usize) -> *mut u8 {
146        if bytes == 0 {
147            return std::ptr::NonNull::<u8>::dangling().as_ptr();
148        }
149        if bytes > self.cap {
150            // Grow by doubling, so a sequence of slightly larger calls does not
151            // reallocate every time.
152            let want = bytes.next_power_of_two();
153            let layout = std::alloc::Layout::from_size_align(want, PAGE).expect("workspace layout");
154            // SAFETY: `want` is nonzero, and any old block was allocated with
155            // this same alignment and a power-of-two size.
156            let new = unsafe {
157                match self.ptr.is_null() {
158                    true => std::alloc::alloc(layout),
159                    false => {
160                        let old = std::alloc::Layout::from_size_align_unchecked(self.cap, PAGE);
161                        std::alloc::realloc(self.ptr, old, want)
162                    }
163                }
164            };
165            assert!(
166                !new.is_null(),
167                "workspace allocation of {want} bytes failed"
168            );
169            self.ptr = new;
170            self.cap = want;
171            if self.traced != 0 {
172                record(self.traced, want);
173            }
174        }
175        self.ptr
176    }
177
178    /// Pointer to the current allocation, or a dangling pointer when empty.
179    fn as_ptr(&self) -> *mut u8 {
180        if self.ptr.is_null() {
181            std::ptr::NonNull::<u8>::dangling().as_ptr()
182        } else {
183            self.ptr
184        }
185    }
186
187    /// Bytes currently allocated.
188    fn cap_bytes(&self) -> usize {
189        self.cap
190    }
191
192    /// Release the allocation. Callers must hold the buffer exclusively.
193    fn trim(&mut self) {
194        if !self.ptr.is_null() {
195            // SAFETY: allocated with this alignment and a power-of-two size.
196            unsafe {
197                std::alloc::dealloc(
198                    self.ptr,
199                    std::alloc::Layout::from_size_align_unchecked(self.cap, PAGE),
200                )
201            };
202            self.ptr = std::ptr::null_mut();
203            self.cap = 0;
204        }
205    }
206}
207
208impl Drop for PageBuf {
209    fn drop(&mut self) {
210        self.trim();
211    }
212}
213
214/// One thread's buffers, as the owner keeps them between calls.
215#[derive(Debug)]
216struct WorkerSlot {
217    a: PageBuf,
218    tile: PageBuf,
219    scratch: Vec<i64>,
220}
221
222impl WorkerSlot {
223    /// `owner` is the tracing key, or zero for an untraced owner.
224    fn new(owner: u64) -> Self {
225        Self {
226            a: PageBuf::traced(owner),
227            tile: PageBuf::traced(owner),
228            scratch: Vec::new(),
229        }
230    }
231
232    fn bytes(&self) -> usize {
233        self.a.cap_bytes()
234            + self.tile.cap_bytes()
235            + self.scratch.capacity() * core::mem::size_of::<i64>()
236    }
237
238    fn trim(&mut self) {
239        self.a.trim();
240        self.tile.trim();
241        self.scratch = Vec::new();
242    }
243}
244
245/// A team's shared storage: one packed B panel, the scatter vectors and the
246/// barriers that publish the panel between its threads.
247#[derive(Debug, Default)]
248struct TeamSet {
249    /// The packed B panel.
250    b: PageBuf,
251    /// Block-scatter vectors for this call, capacity-reused.
252    scatter: Vec<i64>,
253    /// One barrier per column group when `pm > 1`, else empty.
254    barriers: Vec<Barrier>,
255    pm: usize,
256    pn: usize,
257}
258
259impl TeamSet {
260    fn prepare(&mut self, req: &WorkspaceReq, pm: usize, pn: usize) {
261        self.scatter.clear();
262        self.scatter.reserve(req.team_scatter);
263        if (self.pm, self.pn) != (pm, pn) {
264            // A barrier cannot be reset, so a differently shaped team needs new
265            // ones; that is why the free list keeps sets and rebuilds here.
266            self.barriers = (0..req.barriers).map(|_| Barrier::new(pm)).collect();
267            self.pm = pm;
268            self.pn = pn;
269        }
270    }
271
272    fn bytes(&self) -> usize {
273        self.b.cap_bytes()
274            + self.scatter.capacity() * core::mem::size_of::<i64>()
275            + self.barriers.capacity() * core::mem::size_of::<Barrier>()
276    }
277}
278
279/// An exclusive loan of a team set, returned to its owner on drop.
280///
281/// The set is deliberately opaque: the panel grows only through
282/// [`TeamLease::panel`], so every byte the owner hands out is accounted.
283#[derive(Debug)]
284pub struct TeamLease<'a> {
285    owner: &'a ArenaProvider,
286    set: Option<TeamSet>,
287    /// Bytes this lease has added to the owner's leased total.
288    billed: usize,
289}
290
291impl<'a> TeamLease<'a> {
292    /// Size the shared panel and return its pointer, accounting the growth to
293    /// the owner. The driver sizes it because only the driver knows the element
294    /// type; the provider keeps the storage.
295    pub fn panel(&mut self, bytes: usize) -> *mut u8 {
296        let set = self.set.as_mut().expect("team lease holds its set");
297        let before = set.b.cap_bytes();
298        let ptr = set.b.ensure(bytes);
299        let after = set.b.cap_bytes();
300        if after != before {
301            let delta = after as isize - before as isize;
302            let mut st = lock(&self.owner.state);
303            st.owned = add_signed(st.owned, delta, "retained bytes");
304            st.leased = add_signed(st.leased, delta, "leased bytes");
305            st.panel = add_signed(st.panel, delta, "panel bytes");
306            self.billed = add_signed(self.billed, delta, "billed bytes");
307        }
308        ptr
309    }
310
311    /// The team's block-scatter buffer, sized by the driver each call.
312    pub fn scatter_mut(&mut self) -> &mut Vec<i64> {
313        &mut self.set.as_mut().expect("team lease holds its set").scatter
314    }
315
316    /// The team's block-scatter buffer.
317    pub fn scatter(&self) -> &[i64] {
318        &self.set.as_ref().expect("team lease holds its set").scatter
319    }
320
321    /// One barrier per column group, empty when nothing is shared.
322    pub fn barriers(&self) -> &[Barrier] {
323        &self
324            .set
325            .as_ref()
326            .expect("team lease holds its set")
327            .barriers
328    }
329}
330
331impl Drop for TeamLease<'_> {
332    fn drop(&mut self) {
333        if let Some(set) = self.set.take() {
334            let mut st = lock(&self.owner.state);
335            // The scatter vector may have been resized through `scatter_mut`
336            // without going through `panel`, so settle both directions.
337            let delta = set.bytes() as isize - self.billed as isize;
338            st.owned = add_signed(st.owned, delta, "retained bytes");
339            st.leased = sub(st.leased, self.billed, "leased bytes");
340            st.teams.push(set);
341        }
342    }
343}
344
345mod sealed {
346    /// Only the crate's own owner implements [`WorkspaceProvider`](super::WorkspaceProvider):
347    /// the driver writes through the raw pointers it hands out, so an
348    /// implementation must honour the arena's buffer and exclusivity contracts.
349    pub trait Sealed {}
350    impl Sealed for super::ArenaProvider {}
351}
352
353/// Storage an execution may borrow. Implemented by the owner; the driver only
354/// ever sees this trait, so nothing here depends on the driver.
355///
356/// Sealed: [`ArenaProvider`] is the supported way to lend reusable storage, and
357/// the trait is deliberately not implementable outside this crate.
358pub trait WorkspaceProvider: sealed::Sealed + Sync {
359    /// Run `f` with this thread's A block, tile and scatter scratch. A
360    /// re-entrant call on the same thread gets fresh buffers instead of the
361    /// ones its outer call is using.
362    fn with_worker(&self, req: &WorkspaceReq, f: &mut dyn FnMut(*mut u8, *mut u8, &mut Vec<i64>));
363
364    /// Take the team set for `pm x pn`, exclusively, until the lease drops.
365    fn take_team(&self, req: &WorkspaceReq, pm: usize, pn: usize) -> TeamLease<'_>;
366
367    /// Release idle storage. Live leases and checked-out worker slots are kept.
368    fn trim(&self);
369}
370
371/// What one owner holds and lends. One lock covers the storage maps and the
372/// counters, so a snapshot and an update share one synchronization boundary.
373#[derive(Debug, Default)]
374struct State {
375    workers: HashMap<u64, WorkerSlot>,
376    teams: Vec<TeamSet>,
377    /// Capacity of everything owned, idle or in use.
378    owned: usize,
379    /// The part of `owned` lent to a call right now.
380    leased: usize,
381    /// Capacity of the team panels alone, idle or leased. Independent of which
382    /// workers happened to take part, so a test can compare two runs of one
383    /// shape.
384    panel: usize,
385}
386
387/// A workspace owner: per-thread worker slots plus a free list of team sets.
388///
389/// One per owner — never process-global — so a lease always returns to the
390/// provider that issued it and no second owner can borrow it.
391#[derive(Debug)]
392pub struct ArenaProvider {
393    /// Process-unique and never recycled, so a thread-local handle can never be
394    /// confused with another provider that happens to share an address.
395    key: u64,
396    state: Mutex<State>,
397    next_worker: AtomicU64,
398    traced: u64,
399}
400
401impl Default for ArenaProvider {
402    fn default() -> Self {
403        Self::new()
404    }
405}
406
407thread_local! {
408    /// `(provider key, worker slot)` pairs this thread has used. Thread-local
409    /// by construction, so a slot is never handed to a different thread after
410    /// its owner exits: the handle dies with the thread.
411    static HANDLES: std::cell::RefCell<Vec<(u64, u64)>> =
412        const { std::cell::RefCell::new(Vec::new()) };
413}
414
415static NEXT_PROVIDER: AtomicU64 = AtomicU64::new(1);
416static TRACE: Mutex<Vec<(u64, std::thread::ThreadId, usize)>> = Mutex::new(Vec::new());
417
418fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
419    m.lock().unwrap_or_else(|e| e.into_inner())
420}
421
422/// `value + delta`, rejecting a decrement below zero. Every caller establishes
423/// the bound; this turns a mistake into a loud failure instead of a wrap.
424fn add_signed(value: usize, delta: isize, what: &str) -> usize {
425    value
426        .checked_add_signed(delta)
427        .unwrap_or_else(|| panic!("workspace {what} went out of range"))
428}
429
430/// `value - amount`, rejecting an underflow for the same reason.
431fn sub(value: usize, amount: usize, what: &str) -> usize {
432    value
433        .checked_sub(amount)
434        .unwrap_or_else(|| panic!("workspace {what} went out of range"))
435}
436
437fn record(owner: u64, bytes: usize) {
438    lock(&TRACE).push((owner, std::thread::current().id(), bytes));
439}
440
441impl ArenaProvider {
442    /// A fresh owner with no storage.
443    pub fn new() -> Self {
444        Self {
445            key: NEXT_PROVIDER.fetch_add(1, Ordering::Relaxed),
446            state: Mutex::new(State::default()),
447            next_worker: AtomicU64::new(0),
448            traced: 0,
449        }
450    }
451
452    /// A fresh owner that records where its storage is first touched, readable
453    /// through [`ArenaProvider::trace_take`].
454    pub fn traced() -> Self {
455        let mut arena = Self::new();
456        arena.traced = arena.key;
457        arena
458    }
459
460    /// Take this owner's recorded buffer growths, leaving its log empty.
461    ///
462    /// Empty unless the owner was built with [`ArenaProvider::traced`].
463    pub fn trace_take(&self) -> Vec<(std::thread::ThreadId, usize)> {
464        let mut mine = Vec::new();
465        lock(&TRACE).retain(|(owner, thread, bytes)| {
466            if *owner == self.key {
467                mine.push((*thread, *bytes));
468                false
469            } else {
470                true
471            }
472        });
473        mine
474    }
475
476    /// Bytes currently retained and leased, for accounting by the host.
477    pub fn stats(&self) -> WorkspaceStats {
478        let st = lock(&self.state);
479        WorkspaceStats {
480            retained_bytes: st.owned,
481            leased_bytes: st.leased,
482        }
483    }
484
485    /// Bytes currently retained, for accounting in tests and diagnostics.
486    pub fn retained_bytes(&self) -> usize {
487        self.stats().retained_bytes
488    }
489
490    /// Panel bytes retained, idle or leased. Unlike
491    /// [`ArenaProvider::retained_bytes`] this does not depend on which workers
492    /// happened to take part, so it is what a test can compare between two runs
493    /// of one shape.
494    pub fn retained_panel_bytes(&self) -> usize {
495        lock(&self.state).panel
496    }
497
498    /// Check this thread's slot out of the map, or `None` when the thread
499    /// already holds it (a re-entrant call).
500    fn checkout(&self) -> Option<(u64, WorkerSlot, usize)> {
501        let known = HANDLES.with(|h| {
502            h.borrow()
503                .iter()
504                .find(|(k, _)| *k == self.key)
505                .map(|(_, i)| *i)
506        });
507        match known {
508            // The handle exists; a missing slot means an outer call holds it.
509            Some(index) => {
510                let mut st = lock(&self.state);
511                let slot = st.workers.remove(&index)?;
512                let billed = slot.bytes();
513                st.leased += billed;
514                Some((index, slot, billed))
515            }
516            None => {
517                let index = self.next_worker.fetch_add(1, Ordering::Relaxed);
518                HANDLES.with(|h| h.borrow_mut().push((self.key, index)));
519                // A fresh slot owns nothing yet, so the counters do not move.
520                Some((index, WorkerSlot::new(self.traced), 0))
521            }
522        }
523    }
524
525    /// Return a checked-out slot, settling both accounting directions.
526    fn give_back(&self, index: u64, slot: WorkerSlot, billed: usize) {
527        let delta = slot.bytes() as isize - billed as isize;
528        let mut st = lock(&self.state);
529        // `owned` follows the slot's capacity; `leased` gives back exactly what
530        // checkout added, since growth was never lent to anyone else.
531        st.owned = add_signed(st.owned, delta, "retained bytes");
532        st.leased = sub(st.leased, billed, "leased bytes");
533        st.workers.insert(index, slot);
534    }
535}
536
537/// Returns a checked-out slot even when the callback unwinds.
538struct SlotGuard<'a> {
539    owner: &'a ArenaProvider,
540    index: u64,
541    slot: Option<WorkerSlot>,
542    billed: usize,
543}
544
545impl Drop for SlotGuard<'_> {
546    fn drop(&mut self) {
547        if let Some(slot) = self.slot.take() {
548            self.owner.give_back(self.index, slot, self.billed);
549        }
550    }
551}
552
553impl WorkspaceProvider for ArenaProvider {
554    fn with_worker(&self, req: &WorkspaceReq, f: &mut dyn FnMut(*mut u8, *mut u8, &mut Vec<i64>)) {
555        let Some((index, slot, billed)) = self.checkout() else {
556            return fresh_worker(req, f);
557        };
558        // Armed before any growth, so a preparation unwind returns the slot too.
559        let mut guard = SlotGuard {
560            owner: self,
561            index,
562            slot: Some(slot),
563            billed,
564        };
565        let slot = guard.slot.as_mut().expect("guard holds its slot");
566        slot.a.ensure(req.a_bytes);
567        slot.tile.ensure(req.tile_bytes);
568        slot.scratch.clear();
569        slot.scratch.reserve(req.worker_scatter);
570        let (a, tile) = (slot.a.as_ptr(), slot.tile.as_ptr());
571        f(a, tile, &mut slot.scratch);
572    }
573
574    fn take_team(&self, req: &WorkspaceReq, pm: usize, pn: usize) -> TeamLease<'_> {
575        let (mut set, before, before_panel) = {
576            let mut st = lock(&self.state);
577            let set = st.teams.pop().unwrap_or_default();
578            let (bytes, panel) = (set.bytes(), set.b.cap_bytes());
579            (set, bytes, panel)
580        };
581        // `prepare` allocates, so it can panic. A panic drops `set` and its
582        // pages, and the owner has to stop counting the capacity it had when it
583        // was popped; the guard arms before `prepare` for that reason.
584        let mut teardown = Teardown {
585            owner: self,
586            before,
587            before_panel,
588            armed: true,
589        };
590        set.prepare(req, pm, pn);
591        let after = set.bytes();
592        teardown.armed = false;
593        let mut st = lock(&self.state);
594        st.owned = add_signed(st.owned, after as isize - before as isize, "retained bytes");
595        // `prepare` touches the scatter vector and the barriers, never the panel.
596        st.leased += after;
597        drop(st);
598        TeamLease {
599            owner: self,
600            set: Some(set),
601            billed: after,
602        }
603    }
604
605    fn trim(&self) {
606        let mut st = lock(&self.state);
607        let idle_slots: usize = st.workers.values().map(|s| s.bytes()).sum();
608        let idle_teams: usize = st.teams.iter().map(|t| t.bytes()).sum();
609        let idle_panels: usize = st.teams.iter().map(|t| t.b.cap_bytes()).sum();
610        for slot in st.workers.values_mut() {
611            slot.trim();
612        }
613        st.teams.clear();
614        st.owned = sub(st.owned, idle_slots + idle_teams, "retained bytes");
615        st.panel = sub(st.panel, idle_panels, "panel bytes");
616    }
617}
618
619/// Returns a popped team set's accounting when `prepare` unwinds.
620struct Teardown<'a> {
621    owner: &'a ArenaProvider,
622    before: usize,
623    before_panel: usize,
624    armed: bool,
625}
626
627impl Drop for Teardown<'_> {
628    fn drop(&mut self) {
629        if !self.armed {
630            return;
631        }
632        let mut st = lock(&self.owner.state);
633        st.owned = sub(st.owned, self.before, "retained bytes");
634        st.panel = sub(st.panel, self.before_panel, "panel bytes");
635    }
636}
637
638/// Buffers for one call that either re-entered its owner or has none. They are
639/// freed when the call returns, which is the price of re-entrancy, and they are
640/// not traced or accounted: the owner never holds them.
641fn fresh_worker(req: &WorkspaceReq, f: &mut dyn FnMut(*mut u8, *mut u8, &mut Vec<i64>)) {
642    let mut a = PageBuf::default();
643    let mut tile = PageBuf::default();
644    let mut scratch = Vec::new();
645    a.ensure(req.a_bytes);
646    tile.ensure(req.tile_bytes);
647    scratch.reserve(req.worker_scatter);
648    f(a.as_ptr(), tile.as_ptr(), &mut scratch);
649}
650
651/// A caller-owned, bounded workspace: an [`ArenaProvider`] pre-sized to one
652/// [`WorkspaceReq`], so a steady-state execute whose requirement is within the
653/// bound allocates nothing.
654///
655/// A plan computes a conservative requirement (`Plan::workspace_req` in the
656/// owning crate) and `execute_raw_in` checks it component-wise against
657/// [`Workspace::req`] before any write. The bound is grow-only and component-wise:
658/// [`Workspace::reserve`] keeps the maximum of each field, so a large A block
659/// does not stand in for a missing B panel. A parallel executor must call
660/// [`Workspace::prepare_for`] so every worker slot is warmed before the execute.
661///
662/// # Examples
663///
664/// ```
665/// use tcontract_exec::{Workspace, WorkspaceReq};
666/// let mut ws = Workspace::with_req(&WorkspaceReq { a_bytes: 64, ..Default::default() });
667/// assert!(ws.req().a_bytes >= 64);
668/// ws.reserve(&WorkspaceReq { b_bytes: 128, ..Default::default() });
669/// // Component-wise: the A block is kept and the B panel is added.
670/// assert!(ws.req().a_bytes >= 64 && ws.req().b_bytes >= 128);
671/// let _provider = ws.as_provider();
672/// ```
673#[derive(Debug, Default)]
674pub struct Workspace {
675    arena: ArenaProvider,
676    req: WorkspaceReq,
677    /// Worker slots warmed for an executor; see [`Workspace::prepare_for`].
678    slots: usize,
679}
680
681impl Workspace {
682    /// A workspace sized and warmed for `req`: the worker slot of the calling
683    /// thread and one team set are grown to hold it, so a later execute within
684    /// the bound does not grow the arena.
685    pub fn with_req(req: &WorkspaceReq) -> Self {
686        let mut ws = Self {
687            arena: ArenaProvider::new(),
688            req: *req,
689            slots: 0,
690        };
691        ws.warm();
692        ws.slots = 1;
693        ws
694    }
695
696    /// Grow the bound to cover `req` as well, component-wise. The bound never
697    /// shrinks, and every field keeps its own maximum.
698    pub fn reserve(&mut self, req: &WorkspaceReq) {
699        let before = self.req;
700        self.req.a_bytes = self.req.a_bytes.max(req.a_bytes);
701        self.req.tile_bytes = self.req.tile_bytes.max(req.tile_bytes);
702        self.req.worker_scatter = self.req.worker_scatter.max(req.worker_scatter);
703        self.req.b_bytes = self.req.b_bytes.max(req.b_bytes);
704        self.req.team_scatter = self.req.team_scatter.max(req.team_scatter);
705        self.req.barriers = self.req.barriers.max(req.barriers);
706        if self.req != before {
707            self.warm();
708        }
709    }
710
711    /// The bound this workspace was sized to, in bytes.
712    pub fn have(&self) -> usize {
713        self.req.bytes()
714    }
715
716    /// The per-component bound this workspace was sized to.
717    pub fn req(&self) -> &WorkspaceReq {
718        &self.req
719    }
720
721    /// How many worker slots have been warmed. [`Workspace::prepare_for`] raises
722    /// it to the executor's budget.
723    pub fn slots(&self) -> usize {
724        self.slots
725    }
726
727    /// Warm one worker slot per worker the executor may run on, so a parallel
728    /// execute within the bound does not allocate a slot per lane. `broadcast`
729    /// runs the closure on `budget` distinct workers, one call to `with_worker`
730    /// each; a context that cannot co-schedule (a serial one or a same-pool
731    /// worker) keeps the calling thread's slot only.
732    pub fn prepare_for(&mut self, exec: &crate::Exec<'_>) {
733        let slots = exec.budget().max(1);
734        if slots <= self.slots {
735            return;
736        }
737        let (arena, req) = (&self.arena, &self.req);
738        let warmed = exec.broadcast(slots, &|_| arena.with_worker(req, &mut |_, _, _| {}));
739        if warmed.is_ok() {
740            self.slots = slots;
741        }
742    }
743
744    /// The provider a plan lends to the driver for one call.
745    pub fn as_provider(&self) -> &dyn WorkspaceProvider {
746        &self.arena
747    }
748
749    /// Bytes the owner currently retains, idle or in use. A test observes it to
750    /// prove an execute did not grow the arena.
751    pub fn retained_bytes(&self) -> usize {
752        self.arena.retained_bytes()
753    }
754
755    /// First-touch the worker buffers (this thread's slot) and one team set so
756    /// the first execute within the bound allocates only the arena's bookkeeping.
757    fn warm(&self) {
758        self.arena.with_worker(&self.req, &mut |_, _, _| {});
759        let mut lease = self.arena.take_team(&self.req, 1, 1);
760        if self.req.b_bytes > 0 {
761            lease.panel(self.req.b_bytes);
762        }
763        drop(lease);
764    }
765}