rsrs_core/fragmentation/
fragment.rs1use recode::bytes::BytesMut;
2
3use crate::frame::Payload;
4use crate::Frame;
5use super::payload::PayloadFrame;
6
7#[derive(Debug)]
10pub struct Fragmenter {
11 mtu: usize,
12}
13
14impl Fragmenter {
15 #[inline(always)]
20 pub const fn new(mtu: usize) -> Self {
21 Self { mtu }
22 }
23
24 #[inline]
27 pub fn fragment(&self, mut frame: Frame) -> impl Iterator<Item = Frame> {
28 let payload = self.trim_frame(&mut frame);
29
30 std::iter::once(frame).chain(PayloadIterator {
31 payload,
32 mtu: self.mtu,
33 })
34 }
35
36 #[inline]
37 fn trim_frame(&self, frame: &mut Frame) -> Option<(BytesMut, BytesMut)> {
38 match frame {
39 | Frame::RequestResponse(f) => f.trim_to(self.mtu),
40 | Frame::RequestFNF(f) => f.trim_to(self.mtu),
41 | Frame::RequestStream(f) => f.trim_to(self.mtu),
42 | Frame::RequestChannel(f) => f.trim_to(self.mtu),
43 | Frame::Payload(f) => f.trim_to(self.mtu),
44 | _ => Default::default(),
45 }
46 }
47}
48
49#[derive(Debug)]
50struct PayloadIterator {
51 payload: Option<(BytesMut, BytesMut)>,
52 mtu: usize,
53}
54
55impl Iterator for PayloadIterator {
56 type Item = Frame;
57
58 fn next(&mut self) -> Option<Self::Item> {
59 let Some((metadata, data)) = self.payload.take() else {
60 return None;
61 };
62
63 let mut payload = Payload::builder()
64 .metadata(metadata)
65 .data(data)
66 .build()
67 .unwrap();
68
69 if let Some((metadata, data)) = payload.trim_to(self.mtu) {
70 self.payload = Some((metadata, data));
71 payload.set_follow(true);
72 } else {
73 payload.set_follow(false);
74 }
75
76 Some(payload.into())
77 }
78}