weaveffi/abi/
iter.rs

1//! Lazily pulled iterators: the producer's [`Iter`] and the opaque handle a
2//! launcher returns for it.
3
4use std::sync::{Mutex, PoisonError};
5
6use crate::abi::error::{FfiError, MARSHAL_ERROR_CODE};
7
8/// An owned, type-erased iterator returned by a producer function whose IDL
9/// return type is `iter<T>`.
10///
11/// Build one from any iterator with [`Iter::new`]. The `#[weaveffi::module]`
12/// expansion boxes it behind an opaque [`IterHandle`], pulls one element per
13/// `_next` call, and drops it in `_destroy`. Pulling elements lazily (rather
14/// than materializing a `Vec`) is what distinguishes an `iter<T>` return from
15/// a `[T]` (list) return.
16pub struct Iter<T> {
17    inner: Box<dyn Iterator<Item = T> + Send>,
18}
19
20impl<T> Iter<T> {
21    /// Wrap any `Send + 'static` iterator as a WeaveFFI iterator.
22    ///
23    /// Accepts anything `IntoIterator`, so `Iter::new(vec)`, `Iter::new(0..n)`,
24    /// and `Iter::new(map.into_values())` all work.
25    pub fn new<I>(iter: I) -> Self
26    where
27        I: IntoIterator<Item = T>,
28        I::IntoIter: Send + 'static,
29    {
30        Self {
31            inner: Box::new(iter.into_iter()),
32        }
33    }
34}
35
36impl<T> Iterator for Iter<T> {
37    type Item = T;
38
39    fn next(&mut self) -> Option<T> {
40        self.inner.next()
41    }
42}
43
44impl<T> std::fmt::Debug for Iter<T> {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        f.debug_struct("Iter").finish_non_exhaustive()
47    }
48}
49
50/// The opaque iterator handle a launcher returns (`{prefix}_..._Iterator*`
51/// in C).
52///
53/// The handle never holds its lock while the producer's `next` runs: `_next`
54/// takes the iterator out, advances it, and puts it back. A `_next` that
55/// arrives while another is advancing the same handle (from another thread,
56/// or re-entrantly from inside the producer's `next`) fails with
57/// [`MARSHAL_ERROR_CODE`] instead of blocking, so no misuse can deadlock.
58#[derive(Debug)]
59pub struct IterHandle<T> {
60    inner: Mutex<Option<Iter<T>>>,
61}
62
63/// Box `iter` behind a new handle for the consumer, who releases it with
64/// `_destroy` (see [`iter_destroy`]).
65#[must_use]
66pub fn iter_into_raw<T>(iter: Iter<T>) -> *mut IterHandle<T> {
67    crate::abi::leak::track(crate::abi::leak::ITERATORS, 1);
68    Box::into_raw(Box::new(IterHandle {
69        inner: Mutex::new(Some(iter)),
70    }))
71}
72
73/// Puts the iterator back into its handle when dropped, including when the
74/// producer's `next` unwinds, so a panic doesn't leave the handle busy.
75struct Advancing<'a, T> {
76    handle: &'a IterHandle<T>,
77    iter: Option<Iter<T>>,
78}
79
80impl<T> Drop for Advancing<'_, T> {
81    fn drop(&mut self) {
82        let mut slot = self
83            .handle
84            .inner
85            .lock()
86            .unwrap_or_else(PoisonError::into_inner);
87        *slot = self.iter.take();
88    }
89}
90
91/// Pull the next element from a handle: `Ok(None)` when the iterator is
92/// exhausted.
93///
94/// A producer panic from an earlier `next` doesn't wedge the handle: the
95/// iterator is put back in whatever state the panic left it.
96///
97/// # Errors
98///
99/// Returns a [`MARSHAL_ERROR_CODE`] error when `handle` is null or another
100/// `_next` is advancing it right now.
101///
102/// # Safety
103///
104/// `handle` must be null or a live pointer from [`iter_into_raw`] that isn't
105/// destroyed during the call.
106pub unsafe fn iter_next<T>(handle: *const IterHandle<T>) -> Result<Option<T>, FfiError> {
107    // SAFETY: the caller guarantees a live handle or null.
108    let Some(handle) = (unsafe { handle.as_ref() }) else {
109        return Err(FfiError::new(MARSHAL_ERROR_CODE, "iterator is null"));
110    };
111    let taken = handle
112        .inner
113        .lock()
114        .unwrap_or_else(PoisonError::into_inner)
115        .take();
116    let Some(iter) = taken else {
117        return Err(FfiError::new(
118            MARSHAL_ERROR_CODE,
119            "iterator is already being advanced (a concurrent or re-entrant _next)",
120        ));
121    };
122    let mut advancing = Advancing {
123        handle,
124        iter: Some(iter),
125    };
126    let item = advancing.iter.as_mut().and_then(Iterator::next);
127    drop(advancing);
128    Ok(item)
129}
130
131/// The body of every iterator's `_destroy` symbol: drop the iterator and its
132/// handle. A null handle is a no-op, and a panicking `Drop` is swallowed
133/// because a destructor has no `out_err` slot and must never unwind into C.
134///
135/// # Safety
136///
137/// `handle` must be null or a live pointer from [`iter_into_raw`], released
138/// exactly once.
139pub unsafe fn iter_destroy<T>(handle: *mut IterHandle<T>) {
140    if handle.is_null() {
141        return;
142    }
143    crate::abi::leak::track(crate::abi::leak::ITERATORS, -1);
144    let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
145        // SAFETY: the caller hands back a pointer from `Box::into_raw`
146        // exactly once.
147        drop(unsafe { Box::from_raw(handle) });
148    }));
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    fn next<T>(h: *const IterHandle<T>) -> Option<T> {
156        unsafe { iter_next(h) }.unwrap()
157    }
158
159    #[test]
160    fn pulls_until_exhausted() {
161        let h = iter_into_raw(Iter::new(vec![1, 2]));
162        assert_eq!(next(h), Some(1));
163        assert_eq!(next(h), Some(2));
164        assert_eq!(next(h), None);
165        unsafe {
166            assert!(iter_next::<i32>(std::ptr::null()).is_err());
167            iter_destroy(h);
168            iter_destroy::<i32>(std::ptr::null_mut());
169        }
170    }
171
172    #[test]
173    fn concurrent_next_yields_each_element_once_or_reports_busy() {
174        let h = iter_into_raw(Iter::new(0..10_000u32));
175        let addr = h as usize;
176        let threads: Vec<_> = (0..8)
177            .map(|_| {
178                std::thread::spawn(move || {
179                    let h = addr as *const IterHandle<u32>;
180                    let mut got = Vec::new();
181                    loop {
182                        match unsafe { iter_next(h) } {
183                            Ok(Some(x)) => got.push(x),
184                            Ok(None) => break,
185                            Err(e) => {
186                                assert_eq!(e.code, MARSHAL_ERROR_CODE);
187                                std::thread::yield_now();
188                            }
189                        }
190                    }
191                    got
192                })
193            })
194            .collect();
195        let mut all: Vec<u32> = threads
196            .into_iter()
197            .flat_map(|t| t.join().unwrap())
198            .collect();
199        all.sort_unstable();
200        assert_eq!(all, (0..10_000).collect::<Vec<_>>());
201        unsafe { iter_destroy(h) };
202    }
203
204    #[test]
205    fn a_panicking_next_puts_the_iterator_back() {
206        let mut n = 0;
207        let h = iter_into_raw(Iter::new(std::iter::from_fn(move || {
208            n += 1;
209            assert!(n != 2, "second element");
210            Some(n)
211        })));
212        assert_eq!(next(h), Some(1));
213        let addr = h as usize;
214        assert!(std::panic::catch_unwind(move || {
215            let _ = unsafe { iter_next(addr as *const IterHandle<i32>) };
216        })
217        .is_err());
218        assert_eq!(next(h), Some(3));
219        unsafe { iter_destroy(h) };
220    }
221}
weaveffi/abi/iter.rs

weaveffi/abi/
iter.rs
