multiparty/
utils.rs

1use bytes::{BufMut, Bytes, BytesMut};
2
3/// Search for `needle` in `haystack`
4pub fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
5    memchr::memmem::find(haystack, needle)
6}
7
8/// Search for a `needle` that sits in `haystack1` and may continue in `haystack2`
9pub fn find_bytes_split(mut haystack1: &[u8], haystack2: &[u8], needle: &[u8]) -> Option<usize> {
10    let mut i = 0;
11
12    while !haystack1.is_empty() && haystack1.len() + haystack2.len() >= needle.len() {
13        if starts_with_between(haystack1, haystack2, needle) {
14            return Some(i);
15        }
16
17        haystack1 = &haystack1[1..];
18        i += 1;
19    }
20
21    None
22}
23
24/// Determine if `(haystack1 + haystack2).starts_with(needle)`
25pub fn starts_with_between(haystack1: &[u8], haystack2: &[u8], needle: &[u8]) -> bool {
26    let skip1 = haystack1.len().min(needle.len());
27
28    let (needle1, needle2) = needle.split_at(skip1);
29    &haystack1[..skip1] == needle1 && haystack2.starts_with(needle2)
30}
31
32/// Join `bytes1` and `bytes2` into a single allocation
33pub fn join_bytes(bytes1: Bytes, bytes2: Bytes) -> Bytes {
34    if bytes1.is_empty() {
35        bytes2
36    } else if bytes2.is_empty() {
37        bytes1
38    } else {
39        let mut buf = BytesMut::with_capacity(bytes1.len() + bytes2.len());
40        buf.put(bytes1);
41        buf.put(bytes2);
42        buf.freeze()
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn search_bytes() {
52        assert_eq!(find_bytes(b"abcdefgh", b"abc"), Some(0));
53        assert_eq!(find_bytes(b"abc", b"abc"), Some(0));
54        assert_eq!(find_bytes(b"abcdefgh", b"bcde"), Some(1));
55        assert_eq!(find_bytes(b"abcdefgh", b"bc"), Some(1));
56    }
57
58    #[test]
59    fn search_bytes_split() {
60        assert_eq!(find_bytes_split(b"abcd", b"efgh", b"abc"), Some(0));
61        assert_eq!(find_bytes_split(b"abc", b"", b"abc"), Some(0));
62        assert_eq!(find_bytes_split(b"abcd", b"efgh", b"bcde"), Some(1));
63        assert_eq!(find_bytes_split(b"abcd", b"efgh", b"bc"), Some(1));
64        assert_eq!(find_bytes_split(b"abcd", b"efgh", b"fh"), None);
65    }
66
67    #[test]
68    fn join() {
69        assert_eq!(
70            join_bytes(Bytes::copy_from_slice(b"abcd"), Bytes::new()),
71            Bytes::copy_from_slice(b"abcd")
72        );
73        assert_eq!(
74            join_bytes(Bytes::new(), Bytes::copy_from_slice(b"efgh")),
75            Bytes::copy_from_slice(b"efgh")
76        );
77        assert_eq!(
78            join_bytes(
79                Bytes::copy_from_slice(b"abcd"),
80                Bytes::copy_from_slice(b"efgh")
81            ),
82            Bytes::copy_from_slice(b"abcdefgh")
83        );
84    }
85}