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.
18const fn combine(pieces: &[&[u8]], mut out: &mut [u8]) {
19    let mut pieces_idx = 0;
20    while pieces_idx < pieces.len() {
21        let piece = pieces[pieces_idx];
22        slice_copy_from_slice(out, piece);
23        // using split_at_mut because range indexing not yet supported in const fn
24        out = out.split_at_mut(piece.len()).1;
25        pieces_idx += 1;
26    }
27    // should be no trailing buffer
28    assert!(out.is_empty(), "output buffer too large");
29}
30
31/// Wrapper around `combine` which has a const generic parameter, this is going to be more codegen
32/// at compile time (?)
33pub const fn combine_to_array<const LEN: usize>(pieces: &[&[u8]]) -> [u8; LEN] {
34    let mut out: [u8; LEN] = [0u8; LEN];
35    combine(pieces, &mut out);
36    out
37}
38
39/// Replacement for `slice::copy_from_slice`, which is const from 1.87
40pub(crate) const fn slice_copy_from_slice(out: &mut [u8], src: &[u8]) {
41    let mut i = 0;
42    while i < src.len() {
43        out[i] = src[i];
44        i += 1;
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn test_combined_len() {
54        let pieces: [&[u8]; 3] = [b"foo", b"bar", b"baz"];
55        assert_eq!(combined_len(&pieces), 9);
56        let empty: [&[u8]; 0] = [];
57        assert_eq!(combined_len(&empty), 0);
58    }
59
60    #[test]
61    fn test_combine_to_array() {
62        let pieces: [&[u8]; 2] = [b"foo", b"bar"];
63        let combined = combine_to_array::<6>(&pieces);
64        assert_eq!(&combined, b"foobar");
65    }
66
67    #[test]
68    #[should_panic(expected = "index out of bounds")]
69    fn test_combine_to_array_buffer_too_small() {
70        let pieces: [&[u8]; 2] = [b"foo", b"bar"];
71        // Intentionally wrong length
72        let _ = combine_to_array::<5>(&pieces);
73    }
74
75    #[test]
76    #[should_panic(expected = "output buffer too large")]
77    fn test_combine_to_array_buffer_too_big() {
78        let pieces: [&[u8]; 2] = [b"foo", b"bar"];
79        // Intentionally wrong length
80        let _ = combine_to_array::<10>(&pieces);
81    }
82}
pyo3/impl_/concat.rs

pyo3/impl_/
concat.rs
