Skip to main content

pliron/
attribute.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) The pliron contributors
3
4//! Attributes are non-SSA data stored in [Operation](crate::operation::Operation)s.
5//!
6//! See [MLIR Attributes](https://mlir.llvm.org/docs/LangRef/#attributes).
7//! Unlike in MLIR, we do not unique attributes, and hence they are mutable.
8//! These are similar in concept to [Properties](https://discourse.llvm.org/t/rfc-introducing-mlir-operation-properties/67846).
9//! Attribute objects are boxed and not wrapped with [Ptr](crate::context::Ptr).
10//! They are heavy (i.e., not just a pointer, handle or reference),
11//! making clones potentially expensive.
12//!
13//! The [def_attribute](pliron::derive::def_attribute) proc macro from the
14//! pliron-derive create can be used to implement [Attribute] for a rust type.
15//!
16//! Common semantics, API and behaviour of [Attribute]s are
17//! abstracted into interfaces. Interfaces in pliron capture MLIR
18//! functionality of both [Traits](https://mlir.llvm.org/docs/Traits/)
19//! and [Interfaces](https://mlir.llvm.org/docs/Interfaces/).
20//! Interfaces must all implement an associated function named `verify` with
21//! the type [AttrInterfaceVerifier].
22//!
23//! Interfaces are rust Trait definitions annotated with the attribute macro
24//! [attr_interface](pliron::derive::attr_interface). The attribute ensures that any
25//! verifiers of super-interfaces are run prior to the verifier of this interface.
26//! Note: Super-interface verifiers *may* run multiple times for the same attribute.
27//!
28//! [Attribute]s that implement an interface must annotate the implementation with
29//! [attr_interface_impl](pliron::derive::attr_interface_impl) macro to ensure that
30//! the interface verifier is automatically called during verification
31//! and that a `&dyn Attribute` object can be [cast](attr_cast) into an interface object,
32//! (or that it can be checked if the interface is [implemented](attr_impls))
33//! with ease.
34//!
35//! Use [verify_attr] to verify an [Attribute] object.
36//! This function verifies all interfaces implemented by the attribute, and then the attribute itself.
37//! The attribute's verifier must explicitly invoke verifiers on any sub-objects it contains.
38//!
39//! [AttrObj]s can be downcasted to their concrete types using
40//! [downcast_rs](https://docs.rs/downcast-rs/latest/downcast_rs/#example-without-generics).
41
42use alloc::{boxed::Box, string::String, vec::Vec};
43use core::{
44    fmt::{Debug, Display},
45    ops::Deref,
46};
47use downcast_rs::{Downcast, impl_downcast};
48use dyn_clone::DynClone;
49
50use crate::{
51    builtin::attr_interfaces::OutlinedAttr,
52    combine::{Parser, parser, token},
53    common_traits::Verify,
54    context::{Context, collect_deduped_interface_verifiers},
55    dialect::{Dialect, DialectName},
56    identifier::Identifier,
57    impl_printable_for_display, input_err,
58    irfmt::{
59        parsers::{attr_parser, delimited_list_parser, spaced},
60        printers::iter_with_sep,
61    },
62    location::Located,
63    parsable::{Parsable, ParseResult, StateStream},
64    printable::{self, Printable},
65    result::Result,
66    std_deps::{hash::FxHashMap, sync::LazyLock},
67    utils::trait_cast::impls_trait_static,
68};
69
70/// Convenience type to easily print and parse key-value pairs in an [AttributeDict].
71#[derive(Clone)]
72struct AttributeDictKeyVal<'a> {
73    key: &'a Identifier,
74    val: &'a AttrObj,
75}
76
77impl<'a> Printable for AttributeDictKeyVal<'a> {
78    fn fmt(
79        &self,
80        ctx: &Context,
81        _state: &printable::State,
82        f: &mut core::fmt::Formatter<'_>,
83    ) -> core::fmt::Result {
84        write!(f, "{}: {}", self.key, self.val.disp(ctx))
85    }
86}
87
88impl<'b> Parsable for AttributeDictKeyVal<'b> {
89    type Arg = ();
90
91    type Parsed = (Identifier, AttrObj);
92
93    fn parse<'a>(
94        state_stream: &mut StateStream<'a>,
95        _arg: Self::Arg,
96    ) -> ParseResult<'a, Self::Parsed> {
97        (Identifier::parser(()), spaced(token(':')), attr_parser())
98            .map(|(key, _, val)| (key, val))
99            .parse_stream(state_stream)
100            .into_result()
101    }
102}
103
104impl Printable for AttributeDict {
105    fn fmt(
106        &self,
107        ctx: &Context,
108        _state: &printable::State,
109        f: &mut core::fmt::Formatter<'_>,
110    ) -> core::fmt::Result {
111        write!(
112            f,
113            "[{}]",
114            iter_with_sep(
115                self.0
116                    .iter()
117                    .map(|(key, val)| AttributeDictKeyVal { key, val }),
118                printable::ListSeparator::CharSpace(','),
119            )
120            .disp(ctx)
121        )
122    }
123}
124
125impl Parsable for AttributeDict {
126    type Arg = ();
127    type Parsed = Self;
128
129    fn parse<'a>(
130        state_stream: &mut StateStream<'a>,
131        _arg: Self::Arg,
132    ) -> ParseResult<'a, Self::Parsed> {
133        delimited_list_parser('[', ']', ',', AttributeDictKeyVal::parser(()))
134            .map(|key_vals| AttributeDict(key_vals.into_iter().collect()))
135            .parse_stream(state_stream)
136            .into_result()
137    }
138}
139
140/// A dictionary of attributes, mapping keys to attribute objects.
141#[derive(Default, Debug, Clone, PartialEq, Eq)]
142pub struct AttributeDict(pub FxHashMap<Identifier, AttrObj>);
143
144impl AttributeDict {
145    /// Get reference to attribute value that is mapped to key `k`.
146    pub fn get<T: Attribute>(&self, k: &Identifier) -> Option<&T> {
147        self.0.get(k).and_then(|ao| ao.downcast_ref::<T>())
148    }
149
150    /// Get mutable reference to attribute value that is mapped to key `k`.
151    pub fn get_mut<T: Attribute>(&mut self, k: &Identifier) -> Option<&mut T> {
152        self.0.get_mut(k).and_then(|ao| ao.downcast_mut::<T>())
153    }
154
155    /// Reference to the attribute value (that is mapped to key `k`) as an interface reference.
156    pub fn get_as<T: ?Sized + AttrInterfaceMarker + 'static>(&self, k: &Identifier) -> Option<&T> {
157        self.0.get(k).and_then(|ao| attr_cast::<T>(&**ao))
158    }
159
160    /// Set the attribute value for key `k`.
161    pub fn set<T: Attribute>(&mut self, k: Identifier, v: T) {
162        self.0.insert(k, Box::new(v));
163    }
164
165    /// Clone, but skip [Outlined](OutlinedAttr) attributes.
166    pub fn clone_skip_outlined(&self) -> Self {
167        self.0
168            .iter()
169            .filter_map(|(k, v)| {
170                if attr_impls::<dyn OutlinedAttr>(&**v) {
171                    None
172                } else {
173                    Some((k.clone(), dyn_clone::clone_box(&**v)))
174                }
175            })
176            .collect::<FxHashMap<Identifier, AttrObj>>()
177            .into()
178    }
179}
180
181impl From<FxHashMap<Identifier, AttrObj>> for AttributeDict {
182    fn from(value: FxHashMap<Identifier, AttrObj>) -> Self {
183        AttributeDict(value)
184    }
185}
186
187/// Basic functionality that every attribute in the IR must implement.
188///
189/// See [module](crate::attribute) documentation for more information.
190pub trait Attribute: Printable + Verify + Downcast + Sync + Send + DynClone + Debug {
191    /// Is self equal to an other Attribute?
192    fn eq_attr(&self, other: &dyn Attribute) -> bool;
193
194    /// Get an [Attribute]'s static name. This is *not* per instantnce.
195    /// It is mostly useful for printing and parsing the attribute.
196    fn get_attr_id(&self) -> AttrId;
197
198    /// Same as [get_attr_id](Self::get_attr_id), but without the self reference.
199    fn get_attr_id_static() -> AttrId
200    where
201        Self: Sized;
202
203    #[doc(hidden)]
204    /// Verify all interfaces implemented by this attribute.
205    fn verify_interfaces(&self, ctx: &Context) -> Result<()>;
206
207    /// Register this attribute's [AttrId] in the dialect it belongs to.
208    fn register<A: Attribute>(ctx: &mut Context)
209    where
210        Self: Sized + Parsable<Arg = (), Parsed = A>,
211    {
212        let attr_parser: AttrParserFn = |parsable_state, &()| {
213            Self::parse(parsable_state, ()).map(|(attr, r)| -> (AttrObj, _) { (Box::new(attr), r) })
214        };
215        let attrid = Self::get_attr_id_static();
216        Dialect::register(ctx, &attrid.dialect).add_attr(attrid.clone(), attr_parser);
217    }
218}
219impl_downcast!(Attribute);
220dyn_clone::clone_trait_object!(Attribute);
221
222/// [Attribute] objects are boxed and stored in the IR.
223pub type AttrObj = Box<dyn Attribute>;
224
225/// A storable function pointer to parse a specific [Attribute].
226/// The [Attribute]'s [Dialect] maps an [AttrId] to such a parser.
227pub(crate) type AttrParserFn = for<'a> fn(&mut StateStream<'a>, &'a ()) -> ParseResult<'a, AttrObj>;
228
229impl PartialEq for AttrObj {
230    fn eq(&self, other: &Self) -> bool {
231        (**self).eq_attr(&**other)
232    }
233}
234
235impl<T: Attribute> From<T> for AttrObj {
236    fn from(value: T) -> Self {
237        Box::new(value)
238    }
239}
240
241impl Eq for AttrObj {}
242
243impl Printable for AttrObj {
244    fn fmt(
245        &self,
246        ctx: &Context,
247        state: &printable::State,
248        f: &mut core::fmt::Formatter<'_>,
249    ) -> core::fmt::Result {
250        write!(f, "{} ", self.get_attr_id())?;
251        Printable::fmt(self.deref(), ctx, state, f)
252    }
253}
254
255impl Parsable for AttrObj {
256    type Arg = ();
257    type Parsed = AttrObj;
258
259    fn parse<'a>(
260        state_stream: &mut StateStream<'a>,
261        _arg: Self::Arg,
262    ) -> ParseResult<'a, Self::Parsed> {
263        let loc = state_stream.loc();
264        let attr_id_parser = spaced(AttrId::parser(()));
265
266        let mut attr_parser = attr_id_parser.then(move |attr_id: AttrId| {
267            let loc = loc.clone();
268            combine::parser(move |parsable_state: &mut StateStream<'a>| {
269                let state = &parsable_state.state;
270                let dialect = state
271                    .ctx
272                    .dialects
273                    .get(&attr_id.dialect)
274                    .expect("Dialect name parsed but dialect isn't registered");
275                let Some(attr_parser) = dialect.attributes.get(&attr_id) else {
276                    input_err!(
277                        loc.clone(),
278                        "Unregistered attribute {}",
279                        attr_id.disp(state.ctx)
280                    )?
281                };
282                attr_parser(parsable_state, &())
283            })
284        });
285
286        attr_parser.parse_stream(state_stream).into_result()
287    }
288}
289
290/// Verify an [Attribute] object.
291/// 1. Verify all interfaces implemented by this attribute.
292/// 2. Verify the attribute itself.
293pub fn verify_attr(attr: &dyn Attribute, ctx: &Context) -> Result<()> {
294    // Verify all interfaces implemented by this attribute.
295    attr.verify_interfaces(ctx)?;
296
297    // Verify the attribute itself.
298    Verify::verify(attr, ctx)
299}
300
301impl Verify for AttrObj {
302    fn verify(&self, ctx: &Context) -> Result<()> {
303        verify_attr(self.as_ref(), ctx)
304    }
305}
306
307/// Marker trait for attribute interface trait objects.
308///
309/// This is auto-implemented by the `#[attr_interface]` macro for `dyn Interface`
310/// objects and is used to restrict [attr_cast] and [attr_impls] to interface casts.
311#[diagnostic::on_unimplemented(
312    message = "`{Self}` not an attribute interface.",
313    label = "If `{Self}` is a trait, annotate it with #[attr_interface] to be able to cast to it from a `&dyn Attribute`",
314    note = "If you want to cast to a concrete `Attribute`, use `downcast_ref` instead."
315)]
316pub trait AttrInterfaceMarker {}
317
318/// Cast reference to an [Attribute] object to an interface reference.
319///
320/// Right usage: cast to an interface trait object.
321/// ```
322/// use pliron::attribute::{Attribute, attr_cast};
323/// use pliron::builtin::attr_interfaces::TypedAttrInterface;
324///
325/// fn right_cast(attr: &dyn Attribute) {
326///     let _ = attr_cast::<dyn TypedAttrInterface>(attr);
327/// }
328/// ```
329///
330/// Casting to concrete [Attribute] types are intentionally rejected.
331/// ```compile_fail
332/// use pliron::attribute::{Attribute, attr_cast};
333/// use pliron::builtin::attributes::IntegerAttr;
334///
335/// fn wrong_cast(attr: &dyn Attribute) {
336///     let _ = attr_cast::<IntegerAttr>(attr);
337/// }
338/// ```
339/// Use [downcast_rs](https://docs.rs/downcast-rs/latest/downcast_rs/#example-without-generics)
340/// to cast to concrete [Attribute] types.
341pub fn attr_cast<T: ?Sized + AttrInterfaceMarker + 'static>(attr: &dyn Attribute) -> Option<&T> {
342    crate::utils::trait_cast::any_to_trait::<T>(attr.as_any())
343}
344
345/// Does this [Attribute] object implement interface `T`?
346///
347/// Right usage: query using an interface trait object.
348/// ```
349/// use pliron::attribute::{Attribute, attr_impls};
350/// use pliron::builtin::attr_interfaces::TypedAttrInterface;
351///
352/// fn right_query(attr: &dyn Attribute) {
353///     let _ = attr_impls::<dyn TypedAttrInterface>(attr);
354/// }
355/// ```
356///
357/// Querying with a concrete [Attribute] type is intentionally rejected.
358/// ```compile_fail
359/// use pliron::attribute::{Attribute, attr_impls};
360/// use pliron::builtin::attributes::IntegerAttr;
361///
362/// fn wrong_query(attr: &dyn Attribute) {
363///     let _ = attr_impls::<IntegerAttr>(attr);
364/// }
365/// ```
366pub fn attr_impls<T: ?Sized + AttrInterfaceMarker + 'static>(attr: &dyn Attribute) -> bool {
367    attr_cast::<T>(attr).is_some()
368}
369
370/// Does [Attribute] `A` implement interface `I`?
371/// See also: [`attr_impls`].
372///
373/// Example:
374/// ```
375/// use pliron::attribute::{Attribute, attr_impls_static};
376/// use pliron::builtin::attr_interfaces::{FloatAttr, TypedAttrInterface};
377/// use pliron::builtin::attributes::IntegerAttr;
378/// assert!(attr_impls_static::<IntegerAttr, dyn TypedAttrInterface>());
379/// assert!(!attr_impls_static::<IntegerAttr, dyn FloatAttr>());
380/// ```
381pub fn attr_impls_static<A: Attribute, I: ?Sized + AttrInterfaceMarker + 'static>() -> bool {
382    impls_trait_static::<A, I>()
383}
384
385#[derive(Clone, Hash, PartialEq, Eq)]
386/// An [Attribute]'s name (not including it's dialect).
387pub struct AttrName(Identifier);
388
389impl AttrName {
390    /// Create a new AttrName.
391    pub fn try_new(name: &str) -> Result<AttrName> {
392        Identifier::try_from(name).map(AttrName)
393    }
394}
395
396impl_printable_for_display!(AttrName);
397
398impl Display for AttrName {
399    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
400        write!(f, "{}", self.0)
401    }
402}
403
404impl Parsable for AttrName {
405    type Arg = ();
406    type Parsed = AttrName;
407
408    fn parse<'a>(
409        state_stream: &mut crate::parsable::StateStream<'a>,
410        _arg: Self::Arg,
411    ) -> ParseResult<'a, Self::Parsed>
412    where
413        Self: Sized,
414    {
415        Identifier::parser(())
416            .map(AttrName)
417            .parse_stream(state_stream)
418            .into()
419    }
420}
421
422impl Deref for AttrName {
423    type Target = Identifier;
424
425    fn deref(&self) -> &Self::Target {
426        &self.0
427    }
428}
429
430impl From<Identifier> for AttrName {
431    fn from(value: Identifier) -> Self {
432        AttrName(value)
433    }
434}
435
436impl From<AttrName> for Identifier {
437    fn from(value: AttrName) -> Self {
438        value.0
439    }
440}
441
442impl TryFrom<&str> for AttrName {
443    type Error = crate::result::Error;
444
445    fn try_from(value: &str) -> Result<Self> {
446        Identifier::try_from(value).map(AttrName)
447    }
448}
449
450impl TryFrom<String> for AttrName {
451    type Error = crate::result::Error;
452
453    fn try_from(value: String) -> Result<Self> {
454        Identifier::try_from(value).map(AttrName)
455    }
456}
457
458/// A combination of a Attr's name and its dialect.
459#[derive(Clone, Hash, PartialEq, Eq)]
460pub struct AttrId {
461    pub dialect: DialectName,
462    pub name: AttrName,
463}
464
465impl_printable_for_display!(AttrId);
466
467impl Display for AttrId {
468    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
469        write!(f, "{}.{}", self.dialect, self.name)
470    }
471}
472
473impl Parsable for AttrId {
474    type Arg = ();
475    type Parsed = AttrId;
476
477    // Parses (but does not validate) a TypeId.
478    fn parse<'a>(
479        state_stream: &mut StateStream<'a>,
480        _arg: Self::Arg,
481    ) -> ParseResult<'a, Self::Parsed>
482    where
483        Self: Sized,
484    {
485        let mut parser = DialectName::parser(())
486            .skip(parser::char::char('.'))
487            .and(AttrName::parser(()))
488            .map(|(dialect, name)| AttrId { dialect, name });
489        parser.parse_stream(state_stream).into()
490    }
491}
492
493/// Every attribute interface must have a function named `verify` with this type.
494pub type AttrInterfaceVerifier = fn(&dyn Attribute, &Context) -> Result<()>;
495/// Function returns the list of super verifiers, followed by a self verifier, for an interface.
496pub type AttrInterfaceAllVerifiers = fn() -> Vec<AttrInterfaceVerifier>;
497
498#[doc(hidden)]
499/// An [Attribute] paired with an interface it implements
500/// (specifically the verifiers (including super verifiers) for that interface).
501type AttrInterfaceVerifierInfo = (core::any::TypeId, AttrInterfaceAllVerifiers);
502
503#[doc(hidden)]
504#[cfg(not(target_family = "wasm"))]
505pub mod statics {
506    use super::*;
507
508    #[::pliron::linkme::distributed_slice]
509    pub static ATTR_INTERFACE_VERIFIERS: [AttrInterfaceVerifierInfo] = [..];
510
511    pub(super) fn get_attr_interface_verifiers()
512    -> impl Iterator<Item = &'static AttrInterfaceVerifierInfo> {
513        ATTR_INTERFACE_VERIFIERS.iter()
514    }
515}
516#[doc(hidden)]
517#[cfg(not(target_family = "wasm"))]
518pub use statics::ATTR_INTERFACE_VERIFIERS;
519
520#[doc(hidden)]
521#[cfg(target_family = "wasm")]
522pub mod statics {
523    use super::*;
524    use crate::InventoryWrapper;
525
526    ::pliron::inventory::collect!(InventoryWrapper<AttrInterfaceVerifierInfo>);
527
528    pub(super) fn get_attr_interface_verifiers()
529    -> impl Iterator<Item = &'static AttrInterfaceVerifierInfo> {
530        ::pliron::inventory::iter::<InventoryWrapper<AttrInterfaceVerifierInfo>>().map(|llw| llw.0)
531    }
532}
533
534#[doc(hidden)]
535/// A map from every [Attribute] to its ordered (as per interface deps) list of interface verifiers.
536/// An interface's super-interfaces are to be verified before it itself is.
537pub static ATTR_INTERFACE_VERIFIERS_MAP: LazyLock<
538    FxHashMap<core::any::TypeId, Vec<AttrInterfaceVerifier>>,
539> = LazyLock::new(|| collect_deduped_interface_verifiers(statics::get_attr_interface_verifiers()));