pyo3/impl_/
concat.rs

1/// Calculates the total byte length of all byte pieces in the array.
2///
3/// This is a useful utility in order to determine the size needed for the constant
4/// `combine` function.
5pub const fn combined_len(pieces: &[&[u8]]) -> usize {
6    let mut len = 0;
7    let mut pieces_idx = 0;
8    while pieces_idx < pieces.len() {
9        len += pieces[pieces_idx].len();
10        pieces_idx += 1;
11    }
12    len
13}
14
15/// Combines all bytes pieces into a single byte array.
16///
17/// `out` should be a buffer at the correct size of `combined_len(pieces)`, else this will panic.
18#[cfg(mut_ref_in_const_fn)] // requires MSRV 1.83
19#[allow(clippy::incompatible_msrv)] // `split_at_mut` also requires MSRV 1.83
20const fn combine(pieces: &[&[u8]], mut out: &mut [u8]) {
21    let mut pieces_idx = 0;
22    while pieces_idx < pieces.len() {
23        let piece = pieces[pieces_idx];
24        slice_copy_from_slice(out, piece);
25        // using split_at_mut because range indexing not yet supported in const fn
26        out = out.split_at_mut(piece.len()).1;
27        pieces_idx += 1;
28    }
29    // should be no trailing buffer
30    assert!(out.is_empty(), "output buffer too large");
31}
32
33/// Wrapper around `combine` which has a const generic parameter, this is going to be more codegen
34/// at compile time (?)
35///
36/// Unfortunately the `&mut [u8]` buffer needs MSRV 1.83
37pub const fn combine_to_array<const LEN: usize>(pieces: &[&[u8]]) -> [u8; LEN] {
38    let mut out: [u8; LEN] = [0u8; LEN];
39    #[cfg(mut_ref_in_const_fn)]
40    combine(pieces, &mut out);
41    #[cfg(not(mut_ref_in_const_fn))] // inlined here for higher code
42    {
43        let mut out_idx = 0;
44        let mut pieces_idx = 0;
45        while pieces_idx < pieces.len() {
46            let piece = pieces[pieces_idx];
47            let mut piece_idx = 0;
48            while piece_idx < piece.len() {
49                out[out_idx] = piece[piece_idx];
50                out_idx += 1;
51                piece_idx += 1;
52            }
53            pieces_idx += 1;
54        }
55        assert!(out_idx == out.len(), "output buffer too large");
56    }
57    out
58}
59
60/// Replacement for `slice::copy_from_slice`, which is const from 1.87
61#[cfg(mut_ref_in_const_fn)] // requires MSRV 1.83
62const fn slice_copy_from_slice(out: &mut [u8], src: &[u8]) {
63    let mut i = 0;
64    while i < src.len() {
65        out[i] = src[i];
66        i += 1;
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn test_combined_len() {
76        let pieces: [&[u8]; 3] = [b"foo", b"bar", b"baz"];
77        assert_eq!(combined_len(&pieces), 9);
78        let empty: [&[u8]; 0] = [];
79        assert_eq!(combined_len(&empty), 0);
80    }
81
82    #[test]
83    fn test_combine_to_array() {
84        let pieces: [&[u8]; 2] = [b"foo", b"bar"];
85        let combined = combine_to_array::<6>(&pieces);
86        assert_eq!(&combined, b"foobar");
87    }
88
89    #[test]
90    #[should_panic(expected = "index out of bounds")]
91    fn test_combine_to_array_buffer_too_small() {
92        let pieces: [&[u8]; 2] = [b"foo", b"bar"];
93        // Intentionally wrong length
94        let _ = combine_to_array::<5>(&pieces);
95    }
96
97    #[test]
98    #[should_panic(expected = "output buffer too large")]
99    fn test_combine_to_array_buffer_too_big() {
100        let pieces: [&[u8]; 2] = [b"foo", b"bar"];
101        // Intentionally wrong length
102        let _ = combine_to_array::<10>(&pieces);
103    }
104}
pyo3/impl_/concat.rs

pyo3/impl_/
concat.rs
