qir_stdlib/
arrays.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use crate::{strings::convert, update_counts};
5use std::{mem::ManuallyDrop, rc::Rc};
6
7#[cfg(not(feature = "fail-support"))]
8#[allow(improper_ctypes)]
9extern "C" {
10    fn __quantum__rt__fail(str: *const std::ffi::CString);
11}
12
13#[cfg(feature = "fail-support")]
14use crate::__quantum__rt__fail;
15
16#[derive(Debug, Clone)]
17pub struct QirArray {
18    pub(crate) elem_size: usize,
19    pub(crate) data: Vec<u8>,
20}
21
22#[no_mangle]
23pub extern "C" fn __quantum__rt__array_create_1d(elem_size: u32, count: u64) -> *const QirArray {
24    let elem_size = elem_size
25        .try_into()
26        .expect("The `elem_size` argument should fit in the `usize` type for this platform.");
27    let count: usize = count
28        .try_into()
29        .expect("The `count` argument should fit in the `usize` type for this platform.");
30    let data = vec![0_u8; elem_size * count];
31    Rc::into_raw(Rc::new(QirArray { elem_size, data }))
32}
33
34#[no_mangle]
35pub unsafe extern "C" fn __quantum__rt__array_copy(
36    arr: *const QirArray,
37    force: bool,
38) -> *const QirArray {
39    // Wrap the array in a `ManuallyDrop` to effectively borrow it and ensure the array
40    // won't be dropped, refcount decremented, and cleaned up.
41    let rc = ManuallyDrop::new(Rc::from_raw(arr));
42    if force || Rc::weak_count(&rc) > 0 {
43        let copy = rc.as_ref().clone();
44        Rc::into_raw(Rc::new(copy))
45    } else {
46        let _ = Rc::into_raw(Rc::clone(&rc));
47        arr
48    }
49}
50
51#[no_mangle]
52pub unsafe extern "C" fn __quantum__rt__array_concatenate(
53    arr1: *const QirArray,
54    arr2: *const QirArray,
55) -> *const QirArray {
56    let array1 = &*arr1;
57    let array2 = &*arr2;
58    if array1.elem_size != array2.elem_size {
59        __quantum__rt__fail(convert(&format!(
60            "Cannot concatenate arrays with differing element sizes: {} vs {}",
61            array1.elem_size, array2.elem_size
62        )));
63    }
64
65    let mut new_array = QirArray {
66        elem_size: array1.elem_size,
67        data: Vec::new(),
68    };
69    new_array.data.resize(array1.data.len(), 0_u8);
70    new_array.data.copy_from_slice(array1.data.as_slice());
71
72    let mut copy = vec![0; array2.data.len()];
73    copy.copy_from_slice(array2.data.as_slice());
74
75    new_array.data.append(&mut copy);
76    Rc::into_raw(Rc::new(new_array))
77}
78
79#[no_mangle]
80pub unsafe extern "C" fn __quantum__rt__array_get_size_1d(arr: *const QirArray) -> u64 {
81    let array = &*arr;
82    let len = array.data.len() / array.elem_size;
83    len.try_into()
84        .expect("Length of array should always fit in a 64-bit integer.")
85}
86
87#[no_mangle]
88pub unsafe extern "C" fn __quantum__rt__array_get_element_ptr_1d(
89    arr: *const QirArray,
90    index: u64,
91) -> *mut i8 {
92    let array = &*arr;
93    let index: usize = index
94        .try_into()
95        .expect("Indices into an array should fit into the `usize` ");
96    array.data.as_ptr().add(array.elem_size * index) as *mut i8
97}
98
99#[no_mangle]
100pub unsafe extern "C" fn __quantum__rt__array_update_reference_count(
101    arr: *const QirArray,
102    update: i32,
103) {
104    update_counts(arr, update, false);
105}
106
107#[no_mangle]
108pub unsafe extern "C" fn __quantum__rt__array_update_alias_count(
109    arr: *const QirArray,
110    update: i32,
111) {
112    update_counts(arr, update, true);
113}
114
115#[cfg(test)]
116mod tests {
117    use std::{convert::TryInto, mem::size_of};
118
119    use super::*;
120
121    #[test]
122    fn test_array_creation_simple() {
123        let arr = __quantum__rt__array_create_1d(1, 3);
124        unsafe {
125            assert_eq!(__quantum__rt__array_get_size_1d(arr), 3);
126            // Note: array reference counts must be decremented to prevent leaking array from test
127            // and failing address sanitized runs.
128            __quantum__rt__array_update_reference_count(arr, -1);
129        }
130    }
131
132    #[test]
133    fn test_array_item_zero_init() {
134        let arr = __quantum__rt__array_create_1d(1, 3);
135        unsafe {
136            assert_eq!(__quantum__rt__array_get_size_1d(arr), 3);
137            // Array contents should always be zero initialized.
138            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 0), 0);
139            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 1), 0);
140            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 2), 0);
141            __quantum__rt__array_update_reference_count(arr, -1);
142        }
143    }
144
145    #[test]
146    fn test_array_item_updates() {
147        let arr = __quantum__rt__array_create_1d(1, 3);
148        unsafe {
149            assert_eq!(__quantum__rt__array_get_size_1d(arr), 3);
150            // Confirm that array contents can be updated.
151            let first = __quantum__rt__array_get_element_ptr_1d(arr, 0);
152            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 0), 0);
153            *first = 42;
154            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 0), 42);
155            let second = __quantum__rt__array_get_element_ptr_1d(arr, 1);
156            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 1), 0);
157            *second = 31;
158            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 1), 31);
159            let third = __quantum__rt__array_get_element_ptr_1d(arr, 2);
160            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 2), 0);
161            *third = 20;
162            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 2), 20);
163            __quantum__rt__array_update_reference_count(arr, -1);
164        }
165    }
166
167    #[test]
168    fn test_array_copy() {
169        let arr = __quantum__rt__array_create_1d(1, 3);
170        unsafe {
171            assert_eq!(__quantum__rt__array_get_size_1d(arr), 3);
172            let first = __quantum__rt__array_get_element_ptr_1d(arr, 0);
173            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 0), 0);
174            *first = 42;
175            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 0), 42);
176            let second = __quantum__rt__array_get_element_ptr_1d(arr, 1);
177            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 1), 0);
178            *second = 31;
179            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 1), 31);
180            // First array has contents [42, 31, 0], copy to second array and confirm contents.
181            let arr2 = __quantum__rt__array_copy(arr, true);
182            assert_eq!(__quantum__rt__array_get_size_1d(arr2), 3);
183            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr2, 0), 42);
184            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr2, 1), 31);
185            // Update first array to [42, 31, 20], confirm second array is independent copy that still
186            // has original value.
187            let third = __quantum__rt__array_get_element_ptr_1d(arr, 2);
188            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 2), 0);
189            *third = 20;
190            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 2), 20);
191            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr2, 2), 0);
192            // Decrement ref count on first array to trigger deletion, confirm the second array is
193            // independent copy with original values.
194            __quantum__rt__array_update_reference_count(arr, -1);
195            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr2, 0), 42);
196            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr2, 1), 31);
197            __quantum__rt__array_update_reference_count(arr2, -1);
198        }
199    }
200
201    #[test]
202    fn test_array_concat() {
203        let arr = __quantum__rt__array_create_1d(1, 3);
204        unsafe {
205            assert_eq!(__quantum__rt__array_get_size_1d(arr), 3);
206            let first = __quantum__rt__array_get_element_ptr_1d(arr, 0);
207            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 0), 0);
208            *first = 42;
209            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 0), 42);
210            let second = __quantum__rt__array_get_element_ptr_1d(arr, 1);
211            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 1), 0);
212            *second = 31;
213            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr, 1), 31);
214            // First array has contents [42, 31, 0], copy to second array and confirm contents.
215            let arr2 = __quantum__rt__array_copy(arr, true);
216            assert_eq!(__quantum__rt__array_get_size_1d(arr2), 3);
217            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr2, 0), 42);
218            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr2, 1), 31);
219            // Create third array with concatenated contents [42, 31, 0, 42, 31, 0].
220            let arr3 = __quantum__rt__array_concatenate(arr, arr2);
221            assert_eq!(__quantum__rt__array_get_size_1d(arr3), 6);
222            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 0), 42);
223            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 1), 31);
224            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 2), 0);
225            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 3), 42);
226            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 4), 31);
227            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 5), 0);
228            // Decrement counts on first and secon array to trigger deletion, confirm third array
229            // maintains contents.
230            __quantum__rt__array_update_reference_count(arr, -1);
231            __quantum__rt__array_update_reference_count(arr2, -1);
232            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 0), 42);
233            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 1), 31);
234            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 2), 0);
235            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 3), 42);
236            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 4), 31);
237            assert_eq!(*__quantum__rt__array_get_element_ptr_1d(arr3, 5), 0);
238            __quantum__rt__array_update_reference_count(arr3, -1);
239        }
240    }
241
242    #[allow(clippy::cast_ptr_alignment)]
243    #[test]
244    fn test_array_creation_large_size() {
245        // Use a struct that is known to have a size greater than 1.
246        struct Data {
247            first: i64,
248            second: i64,
249            third: i64,
250        }
251        let arr = __quantum__rt__array_create_1d(size_of::<Data>().try_into().unwrap(), 3);
252        unsafe {
253            assert_eq!(__quantum__rt__array_get_size_1d(arr), 3);
254            let first = __quantum__rt__array_get_element_ptr_1d(arr, 0).cast::<Data>();
255            *first = Data {
256                first: 1,
257                second: 2,
258                third: 3,
259            };
260            let second = __quantum__rt__array_get_element_ptr_1d(arr, 1).cast::<Data>();
261            *second = Data {
262                first: 10,
263                second: 20,
264                third: 30,
265            };
266            let third = __quantum__rt__array_get_element_ptr_1d(arr, 2).cast::<Data>();
267            *third = Data {
268                first: 100,
269                second: 200,
270                third: 300,
271            };
272            assert_eq!(
273                (*(__quantum__rt__array_get_element_ptr_1d(arr, 0).cast::<Data>())).first,
274                1
275            );
276            assert_eq!(
277                (*(__quantum__rt__array_get_element_ptr_1d(arr, 0).cast::<Data>())).second,
278                2
279            );
280            assert_eq!(
281                (*(__quantum__rt__array_get_element_ptr_1d(arr, 0).cast::<Data>())).third,
282                3
283            );
284            assert_eq!(
285                (*(__quantum__rt__array_get_element_ptr_1d(arr, 1).cast::<Data>())).first,
286                10
287            );
288            assert_eq!(
289                (*(__quantum__rt__array_get_element_ptr_1d(arr, 1).cast::<Data>())).second,
290                20
291            );
292            assert_eq!(
293                (*(__quantum__rt__array_get_element_ptr_1d(arr, 1).cast::<Data>())).third,
294                30
295            );
296            assert_eq!(
297                (*(__quantum__rt__array_get_element_ptr_1d(arr, 2).cast::<Data>())).first,
298                100
299            );
300            assert_eq!(
301                (*(__quantum__rt__array_get_element_ptr_1d(arr, 2).cast::<Data>())).second,
302                200
303            );
304            assert_eq!(
305                (*(__quantum__rt__array_get_element_ptr_1d(arr, 2).cast::<Data>())).third,
306                300
307            );
308            __quantum__rt__array_update_reference_count(arr, -1);
309        }
310    }
311}