hugr_core/extension/resolution/
weak_registry.rs

1use std::collections::BTreeMap;
2use std::sync::{Arc, Weak};
3
4use itertools::Itertools;
5
6use derive_more::Display;
7
8use crate::extension::{ExtensionId, ExtensionRegistry};
9use crate::Extension;
10
11/// The equivalent to an [`ExtensionRegistry`] that only contains weak
12/// references.
13///
14/// This is used to resolve extensions pointers while the extensions themselves
15/// (and the [`Arc`] that contains them) are being initialized.
16#[derive(Debug, Display, Default, Clone)]
17#[display("WeakExtensionRegistry[{}]", exts.keys().join(", "))]
18pub struct WeakExtensionRegistry {
19    /// The extensions in the registry.
20    exts: BTreeMap<ExtensionId, Weak<Extension>>,
21}
22
23impl WeakExtensionRegistry {
24    /// Create a new weak registry from a list of extensions and their ids.
25    pub fn new(extensions: impl IntoIterator<Item = (ExtensionId, Weak<Extension>)>) -> Self {
26        let mut res = Self::default();
27        for (id, ext) in extensions.into_iter() {
28            res.register(id, ext);
29        }
30        res
31    }
32
33    /// Gets the Extension with the given name
34    pub fn get(&self, name: &str) -> Option<&Weak<Extension>> {
35        self.exts.get(name)
36    }
37
38    /// Returns `true` if the registry contains an extension with the given name.
39    pub fn contains(&self, name: &str) -> bool {
40        self.exts.contains_key(name)
41    }
42
43    /// Register a new extension in the registry.
44    ///
45    /// Returns `true` if the extension was added, `false` if it was already present.
46    pub fn register(&mut self, id: ExtensionId, ext: impl Into<Weak<Extension>>) -> bool {
47        self.exts.insert(id, ext.into()).is_none()
48    }
49
50    /// Returns an iterator over the weak references in the registry and their ids.
51    pub fn iter(&self) -> impl Iterator<Item = (&ExtensionId, &Weak<Extension>)> {
52        self.exts.iter()
53    }
54
55    /// Returns an iterator over the weak references in the registry.
56    pub fn extensions(&self) -> impl Iterator<Item = &Weak<Extension>> {
57        self.exts.values()
58    }
59
60    /// Returns an iterator over the extension ids in the registry.
61    pub fn ids(&self) -> impl Iterator<Item = &ExtensionId> {
62        self.exts.keys()
63    }
64}
65
66impl IntoIterator for WeakExtensionRegistry {
67    type Item = Weak<Extension>;
68    type IntoIter = std::collections::btree_map::IntoValues<ExtensionId, Weak<Extension>>;
69
70    fn into_iter(self) -> Self::IntoIter {
71        self.exts.into_values()
72    }
73}
74
75impl<'a> TryFrom<&'a WeakExtensionRegistry> for ExtensionRegistry {
76    type Error = ();
77
78    fn try_from(weak: &'a WeakExtensionRegistry) -> Result<Self, Self::Error> {
79        let exts: Vec<Arc<Extension>> = weak
80            .extensions()
81            .map(|w| w.upgrade().ok_or(()))
82            .try_collect()?;
83        Ok(ExtensionRegistry::new(exts))
84    }
85}
86
87impl TryFrom<WeakExtensionRegistry> for ExtensionRegistry {
88    type Error = ();
89
90    fn try_from(weak: WeakExtensionRegistry) -> Result<Self, Self::Error> {
91        let exts: Vec<Arc<Extension>> = weak
92            .into_iter()
93            .map(|w| w.upgrade().ok_or(()))
94            .try_collect()?;
95        Ok(ExtensionRegistry::new(exts))
96    }
97}
98
99impl<'a> From<&'a ExtensionRegistry> for WeakExtensionRegistry {
100    fn from(reg: &'a ExtensionRegistry) -> Self {
101        let exts = reg
102            .iter()
103            .map(|ext| (ext.name().clone(), Arc::downgrade(ext)))
104            .collect();
105        Self { exts }
106    }
107}
hugr_core/extension/resolution/weak_registry.rs

hugr_core/extension/resolution/
weak_registry.rs
