Adding ability to access and iterate over function attributes.

pyqir/pyqir/_native.pyi

@@ -19,6 +19,11 @@ class ArrayType(Type):
 class Attribute:
     """An attribute."""
 
+    @property
+    def string_kind(self) -> str:
+        """The kind of this attribute as a string."""
+        ...
+
     @property
     def string_value(self) -> Optional[str]:
         """The value of this attribute as a string, or `None` if this is not a string attribute."""
@@ -44,6 +49,15 @@ class AttributeList:
         """The attributes for the function itself."""
         ...
 
+class AttributeIterator:
+    """An iterator of attributes for a specific part of a function."""
+
+    def __iter__(self) -> object:
+        ...
+
+    def __next__(self) -> Optional[Attribute]:
+        ...
+
 class AttributeSet:
     """A set of attributes for a specific part of a function."""
 
@@ -64,6 +78,9 @@ class AttributeSet:
         """
         ...
 
+    def __iter__(self) -> AttributeIterator:
+        ...
+
 class BasicBlock(Value):
     """A basic block."""
 

pyqir/src/values.rs

@@ -23,7 +23,7 @@ use pyo3::{
     types::{PyBytes, PyLong, PyString},
     PyRef,
 };
-use qirlib::values;
+use qirlib::values::{self, get_string_attribute_kind, get_string_attribute_value};
 use std::{
     borrow::Borrow,
     collections::hash_map::DefaultHasher,
@@ -33,6 +33,7 @@ use std::{
     ops::Deref,
     ptr::NonNull,
     slice, str,
+    vec::IntoIter,
 };
 
 /// A value.
@@ -522,21 +523,20 @@ pub(crate) struct Attribute(LLVMAttributeRef);
 
 #[pymethods]
 impl Attribute {
+    /// The id of this attribute as a string.
+    ///
+    /// :type: str
+    #[getter]
+    fn string_kind(&self) -> String {
+        unsafe { get_string_attribute_kind(self.0) }
+    }
+
     /// The value of this attribute as a string, or `None` if this is not a string attribute.
     ///
     /// :type: typing.Optional[str]
     #[getter]
-    fn string_value(&self) -> Option<&str> {
-        unsafe {
-            if LLVMIsStringAttribute(self.0) == 0 {
-                None
-            } else {
-                let mut len = 0;
-                let value = LLVMGetStringAttributeValue(self.0, &mut len).cast();
-                let value = slice::from_raw_parts(value, len.try_into().unwrap());
-                Some(str::from_utf8(value).unwrap())
-            }
-        }
+    fn string_value(&self) -> Option<String> {
+        unsafe { get_string_attribute_value(self.0) }
     }
 }
 
@@ -588,6 +588,24 @@ pub(crate) struct AttributeSet {
     index: LLVMAttributeIndex,
 }
 
+/// An iterator of attributes for a specific part of a function.
+#[pyclass]
+struct AttributeIterator {
+    iter: IntoIter<Py<Attribute>>,
+}
+
+#[pymethods]
+impl AttributeIterator {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+    // Returning `None` from `__next__` indicates that that there are no further items.
+    // and maps to StopIteration
+    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<Py<Attribute>> {
+        slf.iter.next()
+    }
+}
+
 #[pymethods]
 impl AttributeSet {
     /// Tests if an attribute is a member of the set.
@@ -622,6 +640,24 @@ impl AttributeSet {
             Ok(Attribute(attr))
         }
     }
+
+    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<Py<AttributeIterator>> {
+        let function = slf.function.borrow(slf.py()).into_super().into_super();
+
+        unsafe {
+            let attrs = qirlib::values::get_attributes(function.as_ptr(), slf.index);
+            let items: Vec<Py<Attribute>> = attrs
+                .into_iter()
+                .map(|a| Py::new(slf.py(), Attribute(a)).expect("msg"))
+                .collect();
+            Py::new(
+                slf.py(),
+                AttributeIterator {
+                    iter: items.clone().into_iter(),
+                },
+            )
+        }
+    }
 }
 
 #[derive(FromPyObject)]

qirlib/src/values.rs

@@ -9,7 +9,13 @@ use llvm_sys::{
     core::*, prelude::*, LLVMAttributeFunctionIndex, LLVMAttributeIndex, LLVMLinkage,
     LLVMOpaqueAttributeRef, LLVMOpcode, LLVMTypeKind, LLVMValueKind,
 };
-use std::{convert::TryFrom, ffi::CStr, ptr::NonNull, str};
+use std::{
+    convert::TryFrom,
+    ffi::CStr,
+    mem::{ManuallyDrop, MaybeUninit},
+    ptr::NonNull,
+    str,
+};
 
 pub unsafe fn qubit(context: LLVMContextRef, id: u64) -> LLVMValueRef {
     let i64 = LLVMInt64TypeInContext(context);
@@ -225,6 +231,52 @@ unsafe fn get_string_attribute(
     ))
 }
 
+pub unsafe fn get_attribute_count(function: LLVMValueRef, index: LLVMAttributeIndex) -> usize {
+    LLVMGetAttributeCountAtIndex(function, index)
+        .try_into()
+        .expect("Attribute count larger than usize.")
+}
+
+pub unsafe fn get_string_attribute_kind(attr: *mut LLVMOpaqueAttributeRef) -> String {
+    let mut len = 0;
+    let value = LLVMGetStringAttributeKind(attr, &mut len).cast();
+    let value = slice::from_raw_parts(value, len.try_into().unwrap());
+    str::from_utf8(value)
+        .expect("Attribute kind is not valid UTF-8.")
+        .to_string()
+}
+
+pub unsafe fn get_string_attribute_value(attr: *mut LLVMOpaqueAttributeRef) -> Option<String> {
+    if LLVMIsStringAttribute(attr) == 0 {
+        None
+    } else {
+        let mut len = 0;
+        let value = LLVMGetStringAttributeValue(attr, &mut len).cast();
+        let value = slice::from_raw_parts(value, len.try_into().unwrap());
+        Some(
+            str::from_utf8(value)
+                .expect("Attribute kind is not valid UTF-8.")
+                .to_string(),
+        )
+    }
+}
+
+pub unsafe fn get_attributes(
+    function: LLVMValueRef,
+    index: LLVMAttributeIndex,
+) -> Vec<*mut LLVMOpaqueAttributeRef> {
+    let count = get_attribute_count(function, index);
+    let attrs: Vec<MaybeUninit<*mut LLVMOpaqueAttributeRef>> = Vec::with_capacity(count);
+    let mut attrs = ManuallyDrop::new(attrs);
+    for _ in 0..count {
+        attrs.push(MaybeUninit::uninit());
+    }
+
+    LLVMGetAttributesAtIndex(function, index, attrs.as_mut_ptr() as *mut _);
+
+    Vec::from_raw_parts(attrs.as_mut_ptr() as *mut _, attrs.len(), attrs.capacity())
+}
+
 unsafe fn pointer_to_int(value: LLVMValueRef) -> Option<u64> {
     let ty = LLVMTypeOf(value);
     if LLVMGetTypeKind(ty) == LLVMTypeKind::LLVMPointerTypeKind && LLVMIsConstant(value) != 0 {