aither_core/
ids.rs

1//! Model-facing ids: the words a model reads and types back to name what it
2//! can address, such as a background job, a stored output or a finished
3//! subagent.
4//!
5//! aither mints none of them itself: each is leased from an [`IdAllocator`]
6//! the embedder supplies, and the lease lasts as long as what the word names
7//! can still be addressed. An embedder that also names its own things for the
8//! model leases their words from the same allocator, so no two live ids share
9//! a word. [`WordAllocator`] is the allocator used when none is supplied.
10
11use alloc::{
12    boxed::Box,
13    collections::BTreeMap,
14    string::String,
15    sync::{Arc, Weak},
16};
17use core::{
18    fmt,
19    ops::Deref,
20    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
21};
22
23/// Hands out the words that name what a model can address.
24///
25/// [`lease`](Self::lease) hands out a word no live lease holds, and
26/// [`release`](Self::release) takes it back once what it names can no longer
27/// be addressed. Callers lease through [`AnyIdAllocator::lease`], whose
28/// [`IdLease`] releases the word when its last clone drops, so every lease is
29/// released exactly once.
30///
31/// A word is one or more lowercase ASCII letters: it names files, is written
32/// into commands and is recognized in text the model wrote.
33pub trait IdAllocator: Send + Sync + 'static {
34    /// Leases a word that no live lease holds.
35    ///
36    /// # Errors
37    ///
38    /// Returns [`IdsExhausted`] when every word the allocator can hand out is
39    /// leased.
40    fn lease(&self) -> Result<String, IdsExhausted>;
41
42    /// Ends the lease of `word`, which [`lease`](Self::lease) handed out and
43    /// which has not been released since.
44    fn release(&self, word: &str);
45}
46
47/// Every word an [`IdAllocator`] can hand out is leased.
48#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
49#[error("every model-facing id is leased")]
50pub struct IdsExhausted;
51
52/// A shared [`IdAllocator`].
53///
54/// Clones lease from the same allocator. The default is a [`WordAllocator`]
55/// over its built-in word list.
56#[derive(Clone)]
57pub struct AnyIdAllocator(Arc<dyn IdAllocator>);
58
59impl AnyIdAllocator {
60    /// Shares `allocator`.
61    pub fn new(allocator: impl IdAllocator) -> Self {
62        Self(Arc::new(allocator))
63    }
64
65    /// Leases a word that no live lease holds; it is released when the last
66    /// clone of the returned lease drops.
67    ///
68    /// # Errors
69    ///
70    /// Returns [`IdsExhausted`] when every word is leased.
71    ///
72    /// # Panics
73    ///
74    /// Panics when the allocator hands out a word that is not one or more
75    /// lowercase ASCII letters.
76    pub fn lease(&self) -> Result<IdLease, IdsExhausted> {
77        let word = self.0.lease()?;
78        assert!(
79            is_word(&word),
80            "the id allocator handed out '{word}', which is not one or more lowercase ASCII letters"
81        );
82        Ok(IdLease(Arc::new(Held {
83            word,
84            allocator: self.clone(),
85        })))
86    }
87}
88
89/// Whether `word` has the shape every model-facing id has.
90fn is_word(word: &str) -> bool {
91    !word.is_empty() && word.bytes().all(|byte| byte.is_ascii_lowercase())
92}
93
94impl Default for AnyIdAllocator {
95    fn default() -> Self {
96        Self::new(WordAllocator::default())
97    }
98}
99
100impl fmt::Debug for AnyIdAllocator {
101    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
102        f.debug_struct("AnyIdAllocator").finish_non_exhaustive()
103    }
104}
105
106/// A word leased from an [`AnyIdAllocator`].
107///
108/// Clones share the lease, which ends when the last of them drops: whatever
109/// can still address the named thing holds a clone.
110#[derive(Clone)]
111pub struct IdLease(Arc<Held>);
112
113/// The word of a lease and the allocator it goes back to.
114struct Held {
115    word: String,
116    allocator: AnyIdAllocator,
117}
118
119impl Drop for Held {
120    fn drop(&mut self) {
121        self.allocator.0.release(&self.word);
122    }
123}
124
125impl IdLease {
126    /// The leased word.
127    #[must_use]
128    pub fn as_str(&self) -> &str {
129        &self.0.word
130    }
131
132    /// Whether `other` is a clone of this lease, not merely a lease of the
133    /// same word: a word released and leased again names something else.
134    #[must_use]
135    pub fn same_lease(&self, other: &Self) -> bool {
136        Arc::ptr_eq(&self.0, &other.0)
137    }
138
139    /// A reference to this lease that does not keep the word leased.
140    #[must_use]
141    pub fn downgrade(&self) -> WeakIdLease {
142        WeakIdLease(Arc::downgrade(&self.0))
143    }
144}
145
146/// A reference to an [`IdLease`] that does not keep its word leased, for
147/// what must still tell whether the word names the same thing later.
148#[derive(Clone)]
149pub struct WeakIdLease(Weak<Held>);
150
151impl WeakIdLease {
152    /// The lease, while a clone of it still holds the word; `None` once the
153    /// word went back to its allocator, after which it may name something
154    /// else.
155    #[must_use]
156    pub fn upgrade(&self) -> Option<IdLease> {
157        self.0.upgrade().map(IdLease)
158    }
159
160    /// Whether this refers to `lease` itself, not merely to a lease of the
161    /// same word.
162    #[must_use]
163    pub fn refers_to(&self, lease: &IdLease) -> bool {
164        // The weak reference keeps the lease's allocation, so no later lease
165        // can share its address.
166        core::ptr::eq(self.0.as_ptr(), Arc::as_ptr(&lease.0))
167    }
168}
169
170impl fmt::Debug for WeakIdLease {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        f.write_str("WeakIdLease")
173    }
174}
175
176impl Deref for IdLease {
177    type Target = str;
178
179    fn deref(&self) -> &str {
180        self.as_str()
181    }
182}
183
184impl AsRef<str> for IdLease {
185    fn as_ref(&self) -> &str {
186        self.as_str()
187    }
188}
189
190impl fmt::Display for IdLease {
191    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
192        f.write_str(self.as_str())
193    }
194}
195
196impl fmt::Debug for IdLease {
197    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
198        f.debug_tuple("IdLease").field(&self.as_str()).finish()
199    }
200}
201
202/// The words of the default [`WordAllocator`]: the BIP-39 English word list
203/// (CC0), 2048 common English words, SHA-256
204/// `2f5eed53a4727b4bf8880d8f3f199efc90e58503646d9ff8eff3a2ed3b24dbda`.
205const ENGLISH_WORDS: &str = include_str!("ids/english.txt");
206
207/// An [`IdAllocator`] over a fixed list of words.
208///
209/// It hands the words out in list order, moving on past each one it leases,
210/// so a word just released is the last to come back. [`Default`] uses a
211/// built-in list of 2048 common English words.
212pub struct WordAllocator {
213    words: Box<[Box<str>]>,
214    /// Each word's position in `words`.
215    positions: BTreeMap<Box<str>, usize>,
216    leased: Box<[AtomicBool]>,
217    /// Where the next lease starts looking.
218    cursor: AtomicUsize,
219}
220
221impl WordAllocator {
222    /// An allocator handing out `words`.
223    ///
224    /// # Panics
225    ///
226    /// Panics when `words` is empty, holds one word twice, or holds a word
227    /// that is not one or more lowercase ASCII letters.
228    pub fn new<W: Into<Box<str>>>(words: impl IntoIterator<Item = W>) -> Self {
229        let words: Box<[Box<str>]> = words.into_iter().map(Into::into).collect();
230        assert!(!words.is_empty(), "a word allocator needs words");
231        let mut positions = BTreeMap::new();
232        for (position, word) in words.iter().enumerate() {
233            assert!(
234                is_word(word),
235                "'{word}' is not one or more lowercase ASCII letters"
236            );
237            assert!(
238                positions.insert(word.clone(), position).is_none(),
239                "the word '{word}' is listed twice"
240            );
241        }
242        let leased = words.iter().map(|_| AtomicBool::new(false)).collect();
243        Self {
244            words,
245            positions,
246            leased,
247            cursor: AtomicUsize::new(0),
248        }
249    }
250}
251
252impl Default for WordAllocator {
253    fn default() -> Self {
254        Self::new(ENGLISH_WORDS.lines())
255    }
256}
257
258impl fmt::Debug for WordAllocator {
259    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
260        f.debug_struct("WordAllocator")
261            .field("words", &self.words.len())
262            .finish_non_exhaustive()
263    }
264}
265
266impl IdAllocator for WordAllocator {
267    fn lease(&self) -> Result<String, IdsExhausted> {
268        let count = self.words.len();
269        for _ in 0..count {
270            let position = self.cursor.fetch_add(1, Ordering::Relaxed) % count;
271            if !self.leased[position].swap(true, Ordering::AcqRel) {
272                return Ok(String::from(&*self.words[position]));
273            }
274        }
275        Err(IdsExhausted)
276    }
277
278    fn release(&self, word: &str) {
279        let position = *self
280            .positions
281            .get(word)
282            .unwrap_or_else(|| panic!("released '{word}', which this allocator never leased"));
283        assert!(
284            self.leased[position].swap(false, Ordering::AcqRel),
285            "released '{word}', which is not leased"
286        );
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use alloc::string::ToString;
293
294    use super::*;
295
296    #[test]
297    fn the_default_list_is_single_lowercase_words() {
298        assert_eq!(ENGLISH_WORDS.lines().count(), 2048);
299        // Building it checks every word's shape and that none is listed twice.
300        let _ = WordAllocator::default();
301    }
302
303    #[test]
304    #[should_panic(expected = "is not one or more lowercase ASCII letters")]
305    fn a_word_of_another_shape_is_refused() {
306        let _ = WordAllocator::new(["river", "task-1"]);
307    }
308
309    /// An allocator that hands out a word of the wrong shape.
310    struct Numbered;
311
312    impl IdAllocator for Numbered {
313        fn lease(&self) -> Result<String, IdsExhausted> {
314            Ok("7".to_string())
315        }
316
317        fn release(&self, _word: &str) {}
318    }
319
320    #[test]
321    #[should_panic(expected = "handed out '7'")]
322    fn an_embedder_word_of_another_shape_fails_the_lease() {
323        let _ = AnyIdAllocator::new(Numbered).lease();
324    }
325
326    #[test]
327    fn live_leases_never_share_a_word_and_a_released_word_comes_back() {
328        let ids = AnyIdAllocator::new(WordAllocator::new(["river", "stone"]));
329        let river = ids.lease().expect("a free word");
330        let stone = ids.lease().expect("a free word");
331        assert_eq!((river.as_str(), stone.as_str()), ("river", "stone"));
332        assert_eq!(ids.lease().unwrap_err(), IdsExhausted);
333
334        drop(river);
335        assert_eq!(
336            ids.lease().expect("river is free again").to_string(),
337            "river"
338        );
339    }
340
341    #[test]
342    fn a_lease_ends_with_its_last_clone() {
343        let ids = AnyIdAllocator::new(WordAllocator::new(["river"]));
344        let lease = ids.lease().expect("a free word");
345        let held = lease.clone();
346        drop(lease);
347        assert_eq!(
348            ids.lease().unwrap_err(),
349            IdsExhausted,
350            "a clone still holds it"
351        );
352        drop(held);
353        assert_eq!(
354            &*ids.lease().expect("released with the last clone"),
355            "river"
356        );
357    }
358
359    #[test]
360    fn only_clones_are_the_same_lease() {
361        let ids = AnyIdAllocator::new(WordAllocator::new(["river"]));
362        let lease = ids.lease().expect("a free word");
363        assert!(lease.same_lease(&lease.clone()));
364        drop(lease);
365        let again = ids.lease().expect("river is free again");
366        let other = AnyIdAllocator::new(WordAllocator::new(["river"]))
367            .lease()
368            .expect("a free word");
369        assert_eq!(again.as_str(), other.as_str());
370        assert!(!again.same_lease(&other), "a word alone is not the lease");
371
372        let weak = again.downgrade();
373        assert!(weak.refers_to(&again));
374        assert!(!weak.refers_to(&other));
375        drop(again);
376        let after = ids.lease().expect("river is free again");
377        assert!(
378            !weak.refers_to(&after),
379            "a released lease is not its successor"
380        );
381    }
382
383    #[test]
384    fn a_weak_lease_upgrades_only_while_the_word_is_held() {
385        let ids = AnyIdAllocator::new(WordAllocator::new(["river"]));
386        let lease = ids.lease().expect("a free word");
387        let weak = lease.downgrade();
388        assert_eq!(weak.upgrade().as_deref(), Some("river"));
389        drop(lease);
390        assert!(weak.upgrade().is_none(), "the word went back");
391        let again = ids.lease().expect("river is free again");
392        assert!(
393            weak.upgrade().is_none(),
394            "a word leased again names something else: {again}"
395        );
396    }
397
398    #[test]
399    #[should_panic(expected = "is listed twice")]
400    fn a_word_listed_twice_is_refused() {
401        let _ = WordAllocator::new(["river", "river"]);
402    }
403
404    #[test]
405    #[should_panic(expected = "which is not leased")]
406    fn releasing_a_free_word_is_a_bug() {
407        WordAllocator::new(["river"]).release("river");
408    }
409}
aither_core/ids.rs

aither_core/
ids.rs
