mls_rs_codec/
iter.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// Copyright by contributors to this project.
3// SPDX-License-Identifier: (Apache-2.0 OR MIT)
4
5use crate::{MlsDecode, MlsEncode, MlsSize, VarInt};
6
7use alloc::vec::Vec;
8
9pub fn mls_encoded_len<T>(iter: impl Iterator<Item = T>) -> usize
10where
11    T: MlsSize,
12{
13    let len = iter.map(|x| x.mls_encoded_len()).sum::<usize>();
14
15    let header_length = VarInt::try_from(len).unwrap_or(VarInt(0)).mls_encoded_len();
16
17    header_length + len
18}
19
20#[cfg(feature = "preallocate")]
21pub fn mls_encode<I>(iter: I, writer: &mut Vec<u8>) -> Result<(), crate::Error>
22where
23    I: IntoIterator + Clone,
24    I::Item: MlsEncode,
25{
26    let len = iter
27        .clone()
28        .into_iter()
29        .map(|x| x.mls_encoded_len())
30        .sum::<usize>();
31
32    let header_length = VarInt::try_from(len)?;
33    header_length.mls_encode(writer)?;
34
35    writer.reserve(len);
36
37    iter.into_iter().try_for_each(|x| x.mls_encode(writer))?;
38
39    Ok(())
40}
41
42#[cfg(not(feature = "preallocate"))]
43pub fn mls_encode<I>(iter: I, writer: &mut Vec<u8>) -> Result<(), crate::Error>
44where
45    I: IntoIterator + Clone,
46    I::Item: MlsEncode,
47{
48    let mut buffer = Vec::new();
49
50    iter.into_iter()
51        .try_for_each(|x| x.mls_encode(&mut buffer))?;
52
53    let len = VarInt::try_from(buffer.len())?;
54
55    len.mls_encode(writer)?;
56    writer.extend(buffer);
57
58    Ok(())
59}
60
61pub fn mls_decode_collection<T, F>(reader: &mut &[u8], item_decode: F) -> Result<T, crate::Error>
62where
63    F: Fn(&mut &[u8]) -> Result<T, crate::Error>,
64{
65    let (mut data, rest) = mls_decode_split_on_collection(reader)?;
66
67    let items = item_decode(&mut data)?;
68
69    *reader = rest;
70
71    Ok(items)
72}
73
74pub fn mls_decode_split_on_collection<'b>(
75    reader: &mut &'b [u8],
76) -> Result<(&'b [u8], &'b [u8]), crate::Error> {
77    let len = VarInt::mls_decode(reader)?.0 as usize;
78
79    if len > reader.len() {
80        return Err(crate::Error::UnexpectedEOF);
81    }
82
83    Ok(reader.split_at(len))
84}