calyx_frontend/
attributes.rs

1use super::Attribute;
2use crate::{InlineAttributes, attribute::SetAttribute};
3use calyx_utils::{CalyxResult, GPosIdx, WithPos};
4use itertools::Itertools;
5use linked_hash_map::LinkedHashMap;
6use smallvec::SmallVec;
7use std::{
8    collections::{HashMap, HashSet},
9    convert::TryFrom,
10};
11
12#[derive(Debug, Clone, Default)]
13/// Attribute information stored on the Heap
14struct HeapAttrInfo {
15    attrs: LinkedHashMap<Attribute, u64>,
16    set_attrs: HashMap<SetAttribute, VecSet<u32, 4>>,
17    span: GPosIdx,
18}
19
20/// Attributes associated with a specific IR structure.
21#[derive(Default, Debug, Clone)]
22#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
23pub struct Attributes {
24    /// Inlined attributes
25    inl: InlineAttributes,
26    /// Attributes stored on the heap
27    hinfo: Box<HeapAttrInfo>,
28}
29
30pub enum ParseAttributeWrapper {
31    Attribute(Attribute, u64),
32    Set(SetAttribute, Vec<u32>),
33}
34
35impl From<(Attribute, u64)> for ParseAttributeWrapper {
36    fn from(value: (Attribute, u64)) -> Self {
37        Self::Attribute(value.0, value.1)
38    }
39}
40
41impl From<(SetAttribute, Vec<u32>)> for ParseAttributeWrapper {
42    fn from(value: (SetAttribute, Vec<u32>)) -> Self {
43        Self::Set(value.0, value.1)
44    }
45}
46
47impl TryFrom<Vec<ParseAttributeWrapper>> for Attributes {
48    type Error = calyx_utils::Error;
49
50    fn try_from(v: Vec<ParseAttributeWrapper>) -> CalyxResult<Self> {
51        let mut attrs = Attributes::default();
52
53        for item in v {
54            match item {
55                ParseAttributeWrapper::Attribute(k, v) => {
56                    if attrs.has(k) {
57                        return Err(Self::Error::malformed_structure(format!(
58                            "Multiple entries for attribute: {k}"
59                        )));
60                    }
61                    attrs.insert(k, v);
62                }
63                ParseAttributeWrapper::Set(set_attr, vec) => {
64                    if attrs.hinfo.set_attrs.contains_key(&set_attr) {
65                        return Err(Self::Error::malformed_structure(format!(
66                            "Multiple entries for attribute: {set_attr}"
67                        )));
68                    }
69
70                    attrs
71                        .hinfo
72                        .set_attrs
73                        .insert(set_attr, vec.into_iter().collect());
74                }
75            }
76        }
77        Ok(attrs)
78    }
79}
80
81impl WithPos for Attributes {
82    fn copy_span(&self) -> GPosIdx {
83        self.hinfo.span
84    }
85}
86
87/// Structs that can return an [`Attributes`] instance.
88pub trait GetAttributes {
89    /// Returns an [`Attributes`] instance
90    fn get_attributes(&self) -> &Attributes;
91
92    /// Returns a mutable [`Attributes`] instance
93    fn get_mut_attributes(&mut self) -> &mut Attributes;
94}
95
96impl Attributes {
97    /// Add a new attribute
98    pub fn insert<A>(&mut self, key: A, val: u64)
99    where
100        A: Into<Attribute>,
101    {
102        match key.into() {
103            Attribute::Bool(b) => {
104                assert!(
105                    val == 1,
106                    "{} is a boolean attribute and can only have a value of 1",
107                    b.as_ref(),
108                );
109                self.inl.insert(b);
110            }
111            attr => {
112                self.hinfo.attrs.insert(attr, val);
113            }
114        }
115    }
116    pub fn insert_set<S>(&mut self, key: S, val: u32)
117    where
118        S: Into<SetAttribute>,
119    {
120        self.hinfo
121            .set_attrs
122            .entry(key.into())
123            .or_default()
124            .insert(val);
125    }
126
127    /// Get the value associated with an attribute key
128    pub fn get<A>(&self, key: A) -> Option<u64>
129    where
130        A: Into<Attribute>,
131    {
132        match key.into() {
133            Attribute::Bool(b) => {
134                if self.inl.has(b) {
135                    Some(1)
136                } else {
137                    None
138                }
139            }
140            attr => self.hinfo.attrs.get(&attr).cloned(),
141        }
142    }
143
144    pub fn get_set<S>(&self, key: S) -> Option<&VecSet<u32>>
145    where
146        S: Into<SetAttribute>,
147    {
148        self.hinfo.set_attrs.get(&key.into())
149    }
150
151    /// Check if an attribute key has been set
152    pub fn has<A>(&self, key: A) -> bool
153    where
154        A: Into<Attribute>,
155    {
156        match key.into() {
157            Attribute::Bool(b) => self.inl.has(b),
158            attr => self.hinfo.attrs.contains_key(&attr),
159        }
160    }
161
162    /// Returns true if there are no attributes
163    pub fn is_empty(&self) -> bool {
164        self.inl.is_empty()
165            && self.hinfo.attrs.is_empty()
166            && self.hinfo.set_attrs.is_empty()
167    }
168
169    /// Remove attribute with the name `key`
170    pub fn remove<A>(&mut self, key: A)
171    where
172        A: Into<Attribute>,
173    {
174        match key.into() {
175            Attribute::Bool(b) => {
176                self.inl.remove(b);
177            }
178            attr => {
179                self.hinfo.attrs.remove(&attr);
180            }
181        }
182    }
183
184    /// `self` copys (i.e., assigns the same values) the attributes in `other`.
185    /// However, we only copy attributes in `keys` (i.e.. we don't copy
186    /// all attributes in `other`, only the ones that we specify).
187    /// If a `key` is not present in `other`, then we ignore that `key`.
188    /// Example: suppose
189    /// self: A->10, B->5
190    /// other: A->15, C->5
191    /// keys: A, D
192    /// Then self gets: A->15 B->5. (D is ignored since it's not present in other
193    /// and C is ignored since it's not keys.)
194    pub fn copy_from<A>(&mut self, other: Self, keys: Vec<A>)
195    where
196        A: Into<Attribute> + Clone,
197    {
198        for key in keys {
199            match other.get(key.clone()) {
200                None => (),
201                Some(val) => self.insert(key, val),
202            }
203        }
204    }
205
206    /// Copies the values of the given set attributes,`keys`, from `other`
207    /// into `self`. Note that this does not overwrite set values in `self` that
208    /// are already present.
209    pub fn copy_from_set<A>(&mut self, other: &Self, keys: Vec<A>)
210    where
211        A: Into<SetAttribute> + Clone,
212    {
213        for key in keys {
214            if let Some(vals) = other.get_set(key.clone()) {
215                self.hinfo
216                    .set_attrs
217                    .entry(key.clone().into())
218                    .or_default()
219                    .extend(vals.iter().cloned());
220            }
221        }
222    }
223
224    /// Set the span information
225    pub fn add_span(mut self, span: GPosIdx) -> Self {
226        self.hinfo.span = span;
227        self
228    }
229
230    pub fn to_string_with<F, S>(
231        &self,
232        sep: &'static str,
233        fmt: F,
234        set_fmt: S,
235    ) -> String
236    where
237        F: Fn(String, u64) -> String,
238        S: Fn(String, &[u32]) -> String,
239    {
240        if self.is_empty() {
241            return String::default();
242        }
243
244        self.hinfo
245            .attrs
246            .iter()
247            .map(|(k, v)| fmt(k.to_string(), *v))
248            .chain(self.inl.iter().map(|k| fmt(k.as_ref().to_string(), 1)))
249            .chain(
250                self.hinfo
251                    .set_attrs
252                    .iter()
253                    .sorted_by_key(|(k, _)| *k)
254                    .filter_map(|(k, v)| {
255                        if v.is_empty() {
256                            None
257                        } else {
258                            let formatted =
259                                set_fmt(k.to_string(), v.as_slice());
260                            if formatted.is_empty() {
261                                None
262                            } else {
263                                Some(formatted)
264                            }
265                        }
266                    }),
267            )
268            .collect::<Vec<_>>()
269            .join(sep)
270    }
271}
272
273impl PartialEq for Attributes {
274    fn eq(&self, other: &Self) -> bool {
275        self.inl == other.inl
276            && self.hinfo.attrs.len() == other.hinfo.attrs.len()
277            && self
278                .hinfo
279                .attrs
280                .iter()
281                .all(|(k, v)| other.hinfo.attrs.get(k) == Some(v))
282            && self
283                .hinfo
284                .set_attrs
285                .iter()
286                .all(|(k, v)| other.hinfo.set_attrs.get(k) == Some(v))
287    }
288}
289
290impl Eq for Attributes {}
291
292#[cfg(feature = "serialize")]
293impl serde::Serialize for HeapAttrInfo {
294    fn serialize<S>(&self, ser: S) -> Result<S::Ok, S::Error>
295    where
296        S: serde::Serializer,
297    {
298        ser.collect_map(self.to_owned().attrs.iter())
299    }
300}
301
302#[derive(Debug, Clone, Default, PartialEq, Eq)]
303pub struct VecSet<D, const ALLOC: usize = 4>
304where
305    D: Eq + std::hash::Hash + Clone,
306{
307    inner: SmallVec<[D; ALLOC]>,
308}
309
310impl<D, const ALLOC: usize> Extend<D> for VecSet<D, ALLOC>
311where
312    D: Eq + std::hash::Hash + Clone,
313{
314    fn extend<T: IntoIterator<Item = D>>(&mut self, iter: T) {
315        let mut set: HashSet<_> = self.iter().cloned().collect();
316        for i in iter {
317            if set.insert(i.clone()) {
318                self.inner.push(i);
319            }
320        }
321    }
322}
323
324impl<D, const ALLOC: usize> VecSet<D, ALLOC>
325where
326    D: Eq + std::hash::Hash + Clone,
327{
328    pub fn new() -> Self {
329        Self {
330            inner: SmallVec::new(),
331        }
332    }
333
334    pub fn insert(&mut self, d: D) {
335        if !self.inner.contains(&d) {
336            self.inner.push(d);
337        }
338    }
339
340    pub fn contains(&self, d: &D) -> bool {
341        self.inner.contains(d)
342    }
343
344    pub fn iter(&self) -> impl Iterator<Item = &D> {
345        self.inner.iter()
346    }
347
348    pub fn as_slice(&self) -> &[D] {
349        self.inner.as_slice()
350    }
351
352    pub fn is_empty(&self) -> bool {
353        self.inner.is_empty()
354    }
355}
356
357impl<D, const ALLOC: usize> FromIterator<D> for VecSet<D, ALLOC>
358where
359    D: Eq + std::hash::Hash + Clone,
360{
361    fn from_iter<T: IntoIterator<Item = D>>(iter: T) -> Self {
362        Self {
363            inner: iter.into_iter().unique().collect(),
364        }
365    }
366}
calyx_frontend/attributes.rs

calyx_frontend/
attributes.rs
