melior/ir/attribute/
dense_elements.rs

1use super::{Attribute, AttributeLike};
2use crate::{
3    ir::{Type, TypeLike},
4    Error,
5};
6use mlir_sys::{
7    mlirDenseElementsAttrGet, mlirDenseElementsAttrGetInt32Value,
8    mlirDenseElementsAttrGetInt64Value, mlirElementsAttrGetNumElements, MlirAttribute,
9};
10
11/// A dense elements attribute.
12#[derive(Clone, Copy)]
13pub struct DenseElementsAttribute<'c> {
14    attribute: Attribute<'c>,
15}
16
17impl<'c> DenseElementsAttribute<'c> {
18    /// Creates a dense elements attribute.
19    pub fn new(r#type: Type<'c>, values: &[Attribute<'c>]) -> Result<Self, Error> {
20        if r#type.is_shaped() {
21            Ok(unsafe {
22                Self::from_raw(mlirDenseElementsAttrGet(
23                    r#type.to_raw(),
24                    values.len() as isize,
25                    values.as_ptr() as *const _ as *const _,
26                ))
27            })
28        } else {
29            Err(Error::TypeExpected("shaped", r#type.to_string()))
30        }
31    }
32
33    /// Returns a length.
34    pub fn len(&self) -> usize {
35        (unsafe { mlirElementsAttrGetNumElements(self.attribute.to_raw()) }) as usize
36    }
37
38    /// Checks if an array is empty.
39    pub fn is_empty(&self) -> bool {
40        self.len() == 0
41    }
42
43    /// Returns an i32 element.
44    // TODO Prevent calling these type specific methods on other types.
45    pub fn i32_element(&self, index: usize) -> Result<i32, Error> {
46        if !self.is_dense_int_elements() {
47            Err(Error::ElementExpected {
48                r#type: "integer",
49                value: self.to_string(),
50            })
51        } else if index < self.len() {
52            Ok(unsafe {
53                mlirDenseElementsAttrGetInt32Value(self.attribute.to_raw(), index as isize)
54            })
55        } else {
56            Err(Error::PositionOutOfBounds {
57                name: "dense element",
58                value: self.to_string(),
59                index,
60            })
61        }
62    }
63
64    /// Returns an i64 element.
65    pub fn i64_element(&self, index: usize) -> Result<i64, Error> {
66        if !self.is_dense_int_elements() {
67            Err(Error::ElementExpected {
68                r#type: "integer",
69                value: self.to_string(),
70            })
71        } else if index < self.len() {
72            Ok(unsafe {
73                mlirDenseElementsAttrGetInt64Value(self.attribute.to_raw(), index as isize)
74            })
75        } else {
76            Err(Error::PositionOutOfBounds {
77                name: "dense element",
78                value: self.to_string(),
79                index,
80            })
81        }
82    }
83}
84
85attribute_traits!(DenseElementsAttribute, is_dense_elements, "dense elements");
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use crate::{
91        ir::{
92            attribute::IntegerAttribute,
93            r#type::{IntegerType, MemRefType},
94        },
95        test::create_test_context,
96    };
97
98    #[test]
99    fn i32_element() {
100        let context = create_test_context();
101        let integer_type = IntegerType::new(&context, 32).into();
102        let attribute = DenseElementsAttribute::new(
103            MemRefType::new(integer_type, &[3], None, None).into(),
104            &[IntegerAttribute::new(integer_type, 42).into()],
105        )
106        .unwrap();
107
108        assert_eq!(attribute.i32_element(0), Ok(42));
109        assert_eq!(attribute.i32_element(1), Ok(42));
110        assert_eq!(attribute.i32_element(2), Ok(42));
111        assert_eq!(
112            attribute.i32_element(3),
113            Err(Error::PositionOutOfBounds {
114                name: "dense element",
115                value: attribute.to_string(),
116                index: 3,
117            })
118        );
119    }
120
121    #[test]
122    fn i64_element() {
123        let context = create_test_context();
124        let integer_type = IntegerType::new(&context, 64).into();
125        let attribute = DenseElementsAttribute::new(
126            MemRefType::new(integer_type, &[3], None, None).into(),
127            &[IntegerAttribute::new(integer_type, 42).into()],
128        )
129        .unwrap();
130
131        assert_eq!(attribute.i64_element(0), Ok(42));
132        assert_eq!(attribute.i64_element(1), Ok(42));
133        assert_eq!(attribute.i64_element(2), Ok(42));
134        assert_eq!(
135            attribute.i64_element(3),
136            Err(Error::PositionOutOfBounds {
137                name: "dense element",
138                value: attribute.to_string(),
139                index: 3,
140            })
141        );
142    }
143
144    #[test]
145    fn len() {
146        let context = create_test_context();
147        let integer_type = IntegerType::new(&context, 64).into();
148        let attribute = DenseElementsAttribute::new(
149            MemRefType::new(integer_type, &[3], None, None).into(),
150            &[IntegerAttribute::new(integer_type, 0).into()],
151        )
152        .unwrap();
153
154        assert_eq!(attribute.len(), 3);
155    }
156}
melior/ir/attribute/dense_elements.rs

melior/ir/attribute/
dense_elements.rs
