hugr_llvm/emit/func/
mailbox.rs

1use std::{borrow::Cow, rc::Rc};
2
3use anyhow::{bail, Result};
4use delegate::delegate;
5use inkwell::{
6    builder::Builder,
7    types::{BasicType, BasicTypeEnum},
8    values::{BasicValue, BasicValueEnum, PointerValue},
9};
10use itertools::{zip_eq, Itertools as _};
11
12#[derive(Eq, PartialEq, Clone)]
13pub struct ValueMailBox<'c> {
14    typ: BasicTypeEnum<'c>,
15    ptr: PointerValue<'c>,
16    name: Cow<'static, str>,
17}
18
19fn join_names<'a>(names: impl IntoIterator<Item = &'a str>) -> String {
20    names
21        .into_iter()
22        .filter(|x| !x.is_empty())
23        .join("_")
24        .to_string()
25}
26
27impl<'c> ValueMailBox<'c> {
28    pub(super) fn new(
29        typ: impl BasicType<'c>,
30        ptr: PointerValue<'c>,
31        name: Option<String>,
32    ) -> Self {
33        Self {
34            typ: typ.as_basic_type_enum(),
35            ptr,
36            name: name.map_or(Cow::Borrowed(""), Cow::Owned),
37        }
38    }
39    pub fn get_type(&self) -> BasicTypeEnum<'c> {
40        self.typ
41    }
42
43    pub fn name(&self) -> &str {
44        self.name.as_ref()
45    }
46
47    pub fn promise(&self) -> ValuePromise<'c> {
48        ValuePromise(self.clone())
49    }
50
51    pub fn read<'a>(
52        &'a self,
53        builder: &Builder<'c>,
54        labels: impl IntoIterator<Item = &'a str>,
55    ) -> Result<BasicValueEnum<'c>> {
56        let r = builder.build_load(
57            self.ptr,
58            &join_names(
59                labels
60                    .into_iter()
61                    .chain(std::iter::once(self.name.as_ref())),
62            ),
63        )?;
64        debug_assert_eq!(r.get_type(), self.get_type());
65        Ok(r)
66    }
67
68    fn write(&self, builder: &Builder<'c>, v: impl BasicValue<'c>) -> Result<()> {
69        builder.build_store(self.ptr, v)?;
70        Ok(())
71    }
72}
73
74#[must_use]
75pub struct ValuePromise<'c>(ValueMailBox<'c>);
76
77impl<'c> ValuePromise<'c> {
78    pub fn finish(self, builder: &Builder<'c>, v: impl BasicValue<'c>) -> Result<()> {
79        self.0.write(builder, v)
80    }
81
82    delegate! {
83        to self.0 {
84            pub fn get_type(&self) -> BasicTypeEnum<'c>;
85        }
86    }
87}
88
89/// Holds a vector of [PointerValue]s pointing to `alloca`s in the first block
90/// of a function.
91#[derive(Eq, PartialEq, Clone)]
92#[allow(clippy::len_without_is_empty)]
93pub struct RowMailBox<'c>(Rc<Vec<ValueMailBox<'c>>>, Cow<'static, str>);
94
95impl<'c> RowMailBox<'c> {
96    pub fn new_empty() -> Self {
97        Self::new(std::iter::empty(), None)
98    }
99
100    pub(super) fn new(
101        mbs: impl IntoIterator<Item = ValueMailBox<'c>>,
102        name: Option<String>,
103    ) -> Self {
104        Self(
105            Rc::new(mbs.into_iter().collect_vec()),
106            name.map_or(Cow::Borrowed(""), Cow::Owned),
107        )
108    }
109
110    /// Returns a [RowPromise] that when [RowPromise::finish]ed will write to this `RowMailBox`.
111    pub fn promise(&self) -> RowPromise<'c> {
112        RowPromise(self.clone())
113    }
114
115    /// Get the LLVM types of this `RowMailBox`.
116    pub fn get_types(&'_ self) -> impl Iterator<Item = BasicTypeEnum<'c>> + '_ {
117        self.0.iter().map(ValueMailBox::get_type)
118    }
119
120    /// Returns the number of values in this `RowMailBox`.
121    pub fn len(&self) -> usize {
122        self.0.len()
123    }
124
125    /// Read from the inner pointers.
126    pub fn read_vec<'a>(
127        &'a self,
128        builder: &Builder<'c>,
129        labels: impl IntoIterator<Item = &'a str>,
130    ) -> Result<Vec<BasicValueEnum<'c>>> {
131        self.read(builder, labels)
132    }
133
134    /// Read from the inner pointers.
135    pub fn read<'a, R: FromIterator<BasicValueEnum<'c>>>(
136        &'a self,
137        builder: &Builder<'c>,
138        labels: impl IntoIterator<Item = &'a str>,
139    ) -> Result<R> {
140        let labels = labels.into_iter().collect_vec();
141        self.mailboxes()
142            .map(|mb| mb.read(builder, labels.clone()))
143            .collect::<Result<_>>()
144    }
145
146    pub(crate) fn write(
147        &self,
148        builder: &Builder<'c>,
149        vs: impl IntoIterator<Item = BasicValueEnum<'c>>,
150    ) -> Result<()> {
151        let vs = vs.into_iter().collect_vec();
152        #[cfg(debug_assertions)]
153        {
154            let actual_types = vs.clone().into_iter().map(|x| x.get_type()).collect_vec();
155            let expected_types = self.get_types().collect_vec();
156            if actual_types != expected_types {
157                bail!(
158                    "RowMailbox::write: Expected types {:?}, got {:?}",
159                    expected_types,
160                    actual_types
161                );
162            }
163        }
164        zip_eq(self.0.iter(), vs).try_for_each(|(mb, v)| mb.write(builder, v))
165    }
166
167    fn mailboxes(&'_ self) -> impl Iterator<Item = ValueMailBox<'c>> + '_ {
168        self.0.iter().cloned()
169    }
170}
171
172impl<'c> FromIterator<ValueMailBox<'c>> for RowMailBox<'c> {
173    fn from_iter<T: IntoIterator<Item = ValueMailBox<'c>>>(iter: T) -> Self {
174        Self::new(iter, None)
175    }
176}
177
178/// A promise to write values into a `RowMailBox`
179#[must_use]
180#[allow(clippy::len_without_is_empty)]
181pub struct RowPromise<'c>(RowMailBox<'c>);
182
183impl<'c> RowPromise<'c> {
184    /// Consumes the `RowPromise`, writing the values to the promised [RowMailBox].
185    pub fn finish(
186        self,
187        builder: &Builder<'c>,
188        vs: impl IntoIterator<Item = BasicValueEnum<'c>>,
189    ) -> Result<()> {
190        self.0.write(builder, vs)
191    }
192
193    delegate! {
194        to self.0 {
195            /// Get the LLVM types of this `RowMailBox`.
196            pub fn get_types(&'_ self) -> impl Iterator<Item=BasicTypeEnum<'c>> + '_;
197            /// Returns the number of values promised to be written.
198            pub fn len(&self) -> usize;
199        }
200    }
201}
hugr_llvm/emit/func/mailbox.rs

hugr_llvm/emit/func/
mailbox.rs
