array_init/lib.rs
1#![no_std]
2
3//! The `array-init` crate allows you to initialize arrays
4//! with an initializer closure that will be called
5//! once for each element until the array is filled.
6//!
7//! This way you do not need to default-fill an array
8//! before running initializers. Rust currently only
9//! lets you either specify all initializers at once,
10//! individually (`[a(), b(), c(), ...]`), or specify
11//! one initializer for a `Copy` type (`[a(); N]`),
12//! which will be called once with the result copied over.
13//!
14//! Care is taken not to leak memory shall the initialization
15//! fail.
16//!
17//! # Examples:
18//! ```rust
19//! # #![allow(unused)]
20//! # extern crate array_init;
21//! #
22//! // Initialize an array of length 50 containing
23//! // successive squares
24//!
25//! let arr: [u32; 50] = array_init::array_init(|i: usize| (i * i) as u32);
26//!
27//! // Initialize an array from an iterator
28//! // producing an array of [1,2,3,4] repeated
29//!
30//! let four = [1,2,3,4];
31//! let mut iter = four.iter().copied().cycle();
32//! let arr: [u32; 50] = array_init::from_iter(iter).unwrap();
33//!
34//! // Closures can also mutate state. We guarantee that they will be called
35//! // in order from lower to higher indices.
36//!
37//! let mut last = 1u64;
38//! let mut secondlast = 0;
39//! let fibonacci: [u64; 50] = array_init::array_init(|_| {
40//! let this = last + secondlast;
41//! secondlast = last;
42//! last = this;
43//! this
44//! });
45//! ```
46
47use ::core::{
48 mem::{self, MaybeUninit},
49 ptr, slice,
50};
51
52#[inline]
53/// Initialize an array given an initializer expression.
54///
55/// The initializer is given the index of the element. It is allowed
56/// to mutate external state; we will always initialize the elements in order.
57///
58/// # Examples
59///
60/// ```rust
61/// # #![allow(unused)]
62/// # extern crate array_init;
63/// #
64/// // Initialize an array of length 50 containing
65/// // successive squares
66/// let arr: [usize; 50] = array_init::array_init(|i| i * i);
67///
68/// assert!(arr.iter().enumerate().all(|(i, &x)| x == i * i));
69/// ```
70pub fn array_init<F, T, const N: usize>(mut initializer: F) -> [T; N]
71where
72 F: FnMut(usize) -> T,
73{
74 enum Unreachable {}
75
76 try_array_init(
77 // monomorphise into an infallible version
78 move |i| -> Result<T, Unreachable> { Ok(initializer(i)) },
79 )
80 .unwrap_or_else(
81 // zero-cost unwrap
82 |unreachable| match unreachable { /* ! */ },
83 )
84}
85
86#[inline]
87/// Initialize an array given an iterator
88///
89/// We will iterate until the array is full or the iterator is exhausted. Returns
90/// `None` if the iterator is exhausted before we can fill the array.
91///
92/// - Once the array is full, extra elements from the iterator (if any)
93/// won't be consumed.
94///
95/// # Examples
96///
97/// ```rust
98/// # #![allow(unused)]
99/// # extern crate array_init;
100/// #
101/// // Initialize an array from an iterator
102/// // producing an array of [1,2,3,4] repeated
103///
104/// let four = [1,2,3,4];
105/// let mut iter = four.iter().copied().cycle();
106/// let arr: [u32; 10] = array_init::from_iter(iter).unwrap();
107/// assert_eq!(arr, [1, 2, 3, 4, 1, 2, 3, 4, 1, 2]);
108/// ```
109pub fn from_iter<Iterable, T, const N: usize>(iterable: Iterable) -> Option<[T; N]>
110where
111 Iterable: IntoIterator<Item = T>,
112{
113 try_array_init_impl::<_, _, T, N, 1>({
114 let mut iterator = iterable.into_iter();
115 move |_| iterator.next().ok_or(())
116 })
117 .ok()
118}
119
120#[inline]
121/// Initialize an array in reverse given an iterator
122///
123/// We will iterate until the array is full or the iterator is exhausted. Returns
124/// `None` if the iterator is exhausted before we can fill the array.
125///
126/// - Once the array is full, extra elements from the iterator (if any)
127/// won't be consumed.
128///
129/// # Examples
130///
131/// ```rust
132/// # #![allow(unused)]
133/// # extern crate array_init;
134/// #
135/// // Initialize an array from an iterator
136/// // producing an array of [4,3,2,1] repeated, finishing with 1.
137///
138/// let four = [1,2,3,4];
139/// let mut iter = four.iter().copied().cycle();
140/// let arr: [u32; 10] = array_init::from_iter_reversed(iter).unwrap();
141/// assert_eq!(arr, [2, 1, 4, 3, 2, 1, 4, 3, 2, 1]);
142/// ```
143pub fn from_iter_reversed<Iterable, T, const N: usize>(iterable: Iterable) -> Option<[T; N]>
144where
145 Iterable: IntoIterator<Item = T>,
146{
147 try_array_init_impl::<_, _, T, N, -1>({
148 let mut iterator = iterable.into_iter();
149 move |_| iterator.next().ok_or(())
150 })
151 .ok()
152}
153
154#[inline]
155/// Initialize an array given an initializer expression that may fail.
156///
157/// The initializer is given the index (between 0 and `N - 1` included) of the element, and returns a `Result<T, Err>,`. It is allowed
158/// to mutate external state; we will always initialize from lower to higher indices.
159///
160/// # Examples
161///
162/// ```rust
163/// # #![allow(unused)]
164/// # extern crate array_init;
165/// #
166/// #[derive(PartialEq,Eq,Debug)]
167/// struct DivideByZero;
168///
169/// fn inv(i : usize) -> Result<f64,DivideByZero> {
170/// if i == 0 {
171/// Err(DivideByZero)
172/// } else {
173/// Ok(1./(i as f64))
174/// }
175/// }
176///
177/// // If the initializer does not fail, we get an initialized array
178/// let arr: [f64; 3] = array_init::try_array_init(|i| inv(3-i)).unwrap();
179/// assert_eq!(arr,[1./3., 1./2., 1./1.]);
180///
181/// // The initializer fails
182/// let res : Result<[f64;4], DivideByZero> = array_init::try_array_init(|i| inv(3-i));
183/// assert_eq!(res,Err(DivideByZero));
184/// ```
185pub fn try_array_init<Err, F, T, const N: usize>(initializer: F) -> Result<[T; N], Err>
186where
187 F: FnMut(usize) -> Result<T, Err>,
188{
189 try_array_init_impl::<Err, F, T, N, 1>(initializer)
190}
191
192#[inline]
193/// Initialize an array given a source array and a mapping expression. The size of the source array
194/// is the same as the size of the returned array.
195///
196/// The mapper is given an element from the source array and maps it to an element in the
197/// destination.
198///
199/// # Examples
200///
201/// ```rust
202/// # #![allow(unused)]
203/// # extern crate array_init;
204/// #
205/// // Initialize an array of length 50 containing successive squares
206/// let arr: [usize; 50] = array_init::array_init(|i| i * i);
207///
208/// // Map each usize element to a u64 element.
209/// let u64_arr: [u64; 50] = array_init::map_array_init(&arr, |element| *element as u64);
210///
211/// assert!(u64_arr.iter().enumerate().all(|(i, &x)| x == (i * i) as u64));
212/// ```
213pub fn map_array_init<M, T, U, const N: usize>(source: &[U; N], mut mapper: M) -> [T; N]
214where
215 M: FnMut(&U) -> T,
216{
217 // # Safety
218 // - The array size is known at compile time so we are certain that both the source and
219 // desitination have the same size. If the two arrays are of the same size we know that a
220 // valid index for one would be a valid index for the other.
221 array_init(|index| unsafe { mapper(source.get_unchecked(index)) })
222}
223
224#[inline]
225fn try_array_init_impl<Err, F, T, const N: usize, const D: i8>(
226 mut initializer: F,
227) -> Result<[T; N], Err>
228where
229 F: FnMut(usize) -> Result<T, Err>,
230{
231 // The implementation differentiates two cases:
232 // A) `T` does not need to be dropped. Even if the initializer panics
233 // or returns `Err` we will not leak memory.
234 // B) `T` needs to be dropped. We must keep track of which elements have
235 // been initialized so far, and drop them if we encounter a panic or `Err` midway.
236 if !mem::needs_drop::<T>() {
237 let mut array: MaybeUninit<[T; N]> = MaybeUninit::uninit();
238 // pointer to array = *mut [T; N] <-> *mut T = pointer to first element
239 let mut ptr_i = array.as_mut_ptr() as *mut T;
240
241 // # Safety
242 //
243 // - for D > 0, we are within the array since we start from the
244 // beginning of the array, and we have `0 <= i < N`.
245 // - for D < 0, we start at the end of the array and go back one
246 // place before writing, going back N times in total, finishing
247 // at the start of the array.
248 unsafe {
249 if D < 0 {
250 ptr_i = ptr_i.add(N);
251 }
252 for i in 0..N {
253 let value_i = initializer(i)?;
254 // We overwrite *ptr_i previously undefined value without reading or dropping it.
255 if D < 0 {
256 ptr_i = ptr_i.sub(1);
257 }
258 ptr_i.write(value_i);
259 if D > 0 {
260 ptr_i = ptr_i.add(1);
261 }
262 }
263 Ok(array.assume_init())
264 }
265 } else {
266 // else: `mem::needs_drop::<T>()`
267
268 /// # Safety
269 ///
270 /// - `base_ptr[.. initialized_count]` must be a slice of init elements...
271 ///
272 /// - ... that must be sound to `ptr::drop_in_place` if/when
273 /// `UnsafeDropSliceGuard` is dropped: "symbolic ownership"
274 struct UnsafeDropSliceGuard<Item> {
275 base_ptr: *mut Item,
276 initialized_count: usize,
277 }
278
279 impl<Item> Drop for UnsafeDropSliceGuard<Item> {
280 fn drop(self: &'_ mut Self) {
281 unsafe {
282 // # Safety
283 //
284 // - the contract of the struct guarantees that this is sound
285 ptr::drop_in_place(slice::from_raw_parts_mut(
286 self.base_ptr,
287 self.initialized_count,
288 ));
289 }
290 }
291 }
292
293 // If the `initializer(i)` call panics, `panic_guard` is dropped,
294 // dropping `array[.. initialized_count]` => no memory leak!
295 //
296 // # Safety
297 //
298 // 1. - For D > 0, by construction, array[.. initiliazed_count] only
299 // contains init elements, thus there is no risk of dropping
300 // uninit data;
301 // - For D < 0, by construction, array[N - initialized_count..] only
302 // contains init elements.
303 //
304 // 2. - for D > 0, we are within the array since we start from the
305 // beginning of the array, and we have `0 <= i < N`.
306 // - for D < 0, we start at the end of the array and go back one
307 // place before writing, going back N times in total, finishing
308 // at the start of the array.
309 //
310 unsafe {
311 let mut array: MaybeUninit<[T; N]> = MaybeUninit::uninit();
312 // pointer to array = *mut [T; N] <-> *mut T = pointer to first element
313 let mut ptr_i = array.as_mut_ptr() as *mut T;
314 if D < 0 {
315 ptr_i = ptr_i.add(N);
316 }
317 let mut panic_guard = UnsafeDropSliceGuard {
318 base_ptr: ptr_i,
319 initialized_count: 0,
320 };
321
322 for i in 0..N {
323 // Invariant: `i` elements have already been initialized
324 panic_guard.initialized_count = i;
325 // If this panics or fails, `panic_guard` is dropped, thus
326 // dropping the elements in `base_ptr[.. i]` for D > 0 or
327 // `base_ptr[N - i..]` for D < 0.
328 let value_i = initializer(i)?;
329 // this cannot panic
330 // the previously uninit value is overwritten without being read or dropped
331 if D < 0 {
332 ptr_i = ptr_i.sub(1);
333 panic_guard.base_ptr = ptr_i;
334 }
335 ptr_i.write(value_i);
336 if D > 0 {
337 ptr_i = ptr_i.add(1);
338 }
339 }
340 // From now on, the code can no longer `panic!`, let's take the
341 // symbolic ownership back
342 mem::forget(panic_guard);
343
344 Ok(array.assume_init())
345 }
346 }
347}
348
349#[cfg(test)]
350mod tests {
351 use super::*;
352
353 #[test]
354 fn seq() {
355 let seq: [usize; 5] = array_init(|i| i);
356 assert_eq!(&[0, 1, 2, 3, 4], &seq);
357 }
358
359 #[test]
360 fn array_from_iter() {
361 let array = [0, 1, 2, 3, 4];
362 let seq: [usize; 5] = from_iter(array.iter().copied()).unwrap();
363 assert_eq!(array, seq,);
364 }
365
366 #[test]
367 fn array_init_no_drop() {
368 DropChecker::with(|drop_checker| {
369 let result: Result<[_; 5], ()> = try_array_init(|i| {
370 if i < 3 {
371 Ok(drop_checker.new_element())
372 } else {
373 Err(())
374 }
375 });
376 assert!(result.is_err());
377 });
378 }
379
380 #[test]
381 fn from_iter_no_drop() {
382 DropChecker::with(|drop_checker| {
383 let iterator = (0..3).map(|_| drop_checker.new_element());
384 let result: Option<[_; 5]> = from_iter(iterator);
385 assert!(result.is_none());
386 });
387 }
388
389 #[test]
390 fn from_iter_reversed_no_drop() {
391 DropChecker::with(|drop_checker| {
392 let iterator = (0..3).map(|_| drop_checker.new_element());
393 let result: Option<[_; 5]> = from_iter_reversed(iterator);
394 assert!(result.is_none());
395 });
396 }
397
398 #[test]
399 fn test_513_seq() {
400 let seq: [usize; 513] = array_init(|i| i);
401 assert_eq!(
402 [
403 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
404 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43,
405 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64,
406 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85,
407 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104,
408 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120,
409 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136,
410 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152,
411 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168,
412 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184,
413 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200,
414 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216,
415 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232,
416 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248,
417 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264,
418 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280,
419 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296,
420 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312,
421 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328,
422 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344,
423 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360,
424 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376,
425 377, 378, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 392,
426 393, 394, 395, 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, 408,
427 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, 420, 421, 422, 423, 424,
428 425, 426, 427, 428, 429, 430, 431, 432, 433, 434, 435, 436, 437, 438, 439, 440,
429 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, 456,
430 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, 468, 469, 470, 471, 472,
431 473, 474, 475, 476, 477, 478, 479, 480, 481, 482, 483, 484, 485, 486, 487, 488,
432 489, 490, 491, 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, 504,
433 505, 506, 507, 508, 509, 510, 511, 512
434 ],
435 seq
436 );
437 }
438
439 use self::drop_checker::DropChecker;
440 mod drop_checker {
441 use ::core::cell::Cell;
442
443 pub(super) struct DropChecker {
444 slots: [Cell<bool>; 512],
445 next_uninit_slot: Cell<usize>,
446 }
447
448 pub(super) struct Element<'drop_checker> {
449 slot: &'drop_checker Cell<bool>,
450 }
451
452 impl Drop for Element<'_> {
453 fn drop(self: &'_ mut Self) {
454 assert!(self.slot.replace(false), "Double free!");
455 }
456 }
457
458 impl DropChecker {
459 pub(super) fn with(f: impl FnOnce(&Self)) {
460 let drop_checker = Self::new();
461 f(&drop_checker);
462 drop_checker.assert_no_leaks();
463 }
464
465 pub(super) fn new_element(self: &'_ Self) -> Element<'_> {
466 let i = self.next_uninit_slot.get();
467 self.next_uninit_slot.set(i + 1);
468 self.slots[i].set(true);
469 Element {
470 slot: &self.slots[i],
471 }
472 }
473
474 fn new() -> Self {
475 Self {
476 slots: crate::array_init(|_| Cell::new(false)),
477 next_uninit_slot: Cell::new(0),
478 }
479 }
480
481 fn assert_no_leaks(self: Self) {
482 let leak_count: usize = self.slots[..self.next_uninit_slot.get()]
483 .iter()
484 .map(|slot| usize::from(slot.get() as u8))
485 .sum();
486 assert_eq!(leak_count, 0, "{} elements leaked", leak_count);
487 }
488 }
489 }
490}